faiss 0.2.3 → 0.2.5
Sign up to get free protection for your applications and to get access to all the features.
- checksums.yaml +4 -4
- data/CHANGELOG.md +9 -0
- data/LICENSE.txt +1 -1
- data/README.md +23 -21
- data/ext/faiss/extconf.rb +11 -0
- data/ext/faiss/index.cpp +4 -4
- data/ext/faiss/index_binary.cpp +6 -6
- data/ext/faiss/product_quantizer.cpp +4 -4
- data/lib/faiss/version.rb +1 -1
- data/vendor/faiss/faiss/AutoTune.cpp +13 -0
- data/vendor/faiss/faiss/Clustering.cpp +32 -0
- data/vendor/faiss/faiss/Clustering.h +14 -0
- data/vendor/faiss/faiss/IVFlib.cpp +101 -2
- data/vendor/faiss/faiss/IVFlib.h +26 -2
- data/vendor/faiss/faiss/Index.cpp +36 -3
- data/vendor/faiss/faiss/Index.h +43 -6
- data/vendor/faiss/faiss/Index2Layer.cpp +24 -93
- data/vendor/faiss/faiss/Index2Layer.h +8 -17
- data/vendor/faiss/faiss/IndexAdditiveQuantizer.cpp +610 -0
- data/vendor/faiss/faiss/IndexAdditiveQuantizer.h +253 -0
- data/vendor/faiss/faiss/IndexAdditiveQuantizerFastScan.cpp +299 -0
- data/vendor/faiss/faiss/IndexAdditiveQuantizerFastScan.h +199 -0
- data/vendor/faiss/faiss/IndexBinary.cpp +20 -4
- data/vendor/faiss/faiss/IndexBinary.h +18 -3
- data/vendor/faiss/faiss/IndexBinaryFlat.cpp +9 -2
- data/vendor/faiss/faiss/IndexBinaryFlat.h +4 -2
- data/vendor/faiss/faiss/IndexBinaryFromFloat.cpp +4 -1
- data/vendor/faiss/faiss/IndexBinaryFromFloat.h +2 -1
- data/vendor/faiss/faiss/IndexBinaryHNSW.cpp +5 -1
- data/vendor/faiss/faiss/IndexBinaryHNSW.h +2 -1
- data/vendor/faiss/faiss/IndexBinaryHash.cpp +17 -4
- data/vendor/faiss/faiss/IndexBinaryHash.h +8 -4
- data/vendor/faiss/faiss/IndexBinaryIVF.cpp +28 -13
- data/vendor/faiss/faiss/IndexBinaryIVF.h +10 -7
- data/vendor/faiss/faiss/IndexFastScan.cpp +626 -0
- data/vendor/faiss/faiss/IndexFastScan.h +145 -0
- data/vendor/faiss/faiss/IndexFlat.cpp +52 -69
- data/vendor/faiss/faiss/IndexFlat.h +16 -19
- data/vendor/faiss/faiss/IndexFlatCodes.cpp +101 -0
- data/vendor/faiss/faiss/IndexFlatCodes.h +59 -0
- data/vendor/faiss/faiss/IndexHNSW.cpp +66 -138
- data/vendor/faiss/faiss/IndexHNSW.h +4 -2
- data/vendor/faiss/faiss/IndexIDMap.cpp +247 -0
- data/vendor/faiss/faiss/IndexIDMap.h +107 -0
- data/vendor/faiss/faiss/IndexIVF.cpp +200 -40
- data/vendor/faiss/faiss/IndexIVF.h +59 -22
- data/vendor/faiss/faiss/IndexIVFAdditiveQuantizer.cpp +393 -0
- data/vendor/faiss/faiss/IndexIVFAdditiveQuantizer.h +183 -0
- data/vendor/faiss/faiss/IndexIVFAdditiveQuantizerFastScan.cpp +590 -0
- data/vendor/faiss/faiss/IndexIVFAdditiveQuantizerFastScan.h +171 -0
- data/vendor/faiss/faiss/IndexIVFFastScan.cpp +1290 -0
- data/vendor/faiss/faiss/IndexIVFFastScan.h +213 -0
- data/vendor/faiss/faiss/IndexIVFFlat.cpp +43 -26
- data/vendor/faiss/faiss/IndexIVFFlat.h +4 -2
- data/vendor/faiss/faiss/IndexIVFPQ.cpp +238 -53
- data/vendor/faiss/faiss/IndexIVFPQ.h +6 -2
- data/vendor/faiss/faiss/IndexIVFPQFastScan.cpp +23 -852
- data/vendor/faiss/faiss/IndexIVFPQFastScan.h +7 -112
- data/vendor/faiss/faiss/IndexIVFPQR.cpp +3 -3
- data/vendor/faiss/faiss/IndexIVFPQR.h +1 -1
- data/vendor/faiss/faiss/IndexIVFSpectralHash.cpp +63 -40
- data/vendor/faiss/faiss/IndexIVFSpectralHash.h +23 -7
- data/vendor/faiss/faiss/IndexLSH.cpp +8 -32
- data/vendor/faiss/faiss/IndexLSH.h +4 -16
- data/vendor/faiss/faiss/IndexLattice.cpp +7 -1
- data/vendor/faiss/faiss/IndexLattice.h +3 -1
- data/vendor/faiss/faiss/IndexNNDescent.cpp +4 -5
- data/vendor/faiss/faiss/IndexNNDescent.h +2 -1
- data/vendor/faiss/faiss/IndexNSG.cpp +37 -5
- data/vendor/faiss/faiss/IndexNSG.h +25 -1
- data/vendor/faiss/faiss/IndexPQ.cpp +108 -120
- data/vendor/faiss/faiss/IndexPQ.h +21 -22
- data/vendor/faiss/faiss/IndexPQFastScan.cpp +15 -450
- data/vendor/faiss/faiss/IndexPQFastScan.h +15 -78
- data/vendor/faiss/faiss/IndexPreTransform.cpp +47 -8
- data/vendor/faiss/faiss/IndexPreTransform.h +15 -3
- data/vendor/faiss/faiss/IndexRefine.cpp +36 -4
- data/vendor/faiss/faiss/IndexRefine.h +14 -2
- data/vendor/faiss/faiss/IndexReplicas.cpp +4 -2
- data/vendor/faiss/faiss/IndexReplicas.h +2 -1
- data/vendor/faiss/faiss/IndexRowwiseMinMax.cpp +438 -0
- data/vendor/faiss/faiss/IndexRowwiseMinMax.h +92 -0
- data/vendor/faiss/faiss/IndexScalarQuantizer.cpp +28 -43
- data/vendor/faiss/faiss/IndexScalarQuantizer.h +8 -23
- data/vendor/faiss/faiss/IndexShards.cpp +4 -1
- data/vendor/faiss/faiss/IndexShards.h +2 -1
- data/vendor/faiss/faiss/MetaIndexes.cpp +5 -178
- data/vendor/faiss/faiss/MetaIndexes.h +3 -81
- data/vendor/faiss/faiss/VectorTransform.cpp +45 -1
- data/vendor/faiss/faiss/VectorTransform.h +25 -4
- data/vendor/faiss/faiss/clone_index.cpp +26 -3
- data/vendor/faiss/faiss/clone_index.h +3 -0
- data/vendor/faiss/faiss/cppcontrib/SaDecodeKernels.h +300 -0
- data/vendor/faiss/faiss/cppcontrib/detail/CoarseBitType.h +24 -0
- data/vendor/faiss/faiss/cppcontrib/detail/UintReader.h +195 -0
- data/vendor/faiss/faiss/cppcontrib/sa_decode/Level2-avx2-inl.h +2058 -0
- data/vendor/faiss/faiss/cppcontrib/sa_decode/Level2-inl.h +408 -0
- data/vendor/faiss/faiss/cppcontrib/sa_decode/Level2-neon-inl.h +2147 -0
- data/vendor/faiss/faiss/cppcontrib/sa_decode/MinMax-inl.h +460 -0
- data/vendor/faiss/faiss/cppcontrib/sa_decode/MinMaxFP16-inl.h +465 -0
- data/vendor/faiss/faiss/cppcontrib/sa_decode/PQ-avx2-inl.h +1618 -0
- data/vendor/faiss/faiss/cppcontrib/sa_decode/PQ-inl.h +251 -0
- data/vendor/faiss/faiss/cppcontrib/sa_decode/PQ-neon-inl.h +1452 -0
- data/vendor/faiss/faiss/gpu/GpuAutoTune.cpp +1 -0
- data/vendor/faiss/faiss/gpu/GpuCloner.cpp +2 -6
- data/vendor/faiss/faiss/gpu/GpuIcmEncoder.h +60 -0
- data/vendor/faiss/faiss/gpu/GpuIndex.h +28 -4
- data/vendor/faiss/faiss/gpu/GpuIndexBinaryFlat.h +2 -1
- data/vendor/faiss/faiss/gpu/GpuIndexFlat.h +10 -8
- data/vendor/faiss/faiss/gpu/GpuIndexIVF.h +75 -14
- data/vendor/faiss/faiss/gpu/GpuIndexIVFFlat.h +19 -32
- data/vendor/faiss/faiss/gpu/GpuIndexIVFPQ.h +22 -31
- data/vendor/faiss/faiss/gpu/GpuIndexIVFScalarQuantizer.h +22 -28
- data/vendor/faiss/faiss/gpu/GpuResources.cpp +14 -0
- data/vendor/faiss/faiss/gpu/GpuResources.h +16 -3
- data/vendor/faiss/faiss/gpu/StandardGpuResources.cpp +3 -3
- data/vendor/faiss/faiss/gpu/impl/IndexUtils.h +32 -0
- data/vendor/faiss/faiss/gpu/test/TestGpuIndexBinaryFlat.cpp +1 -0
- data/vendor/faiss/faiss/gpu/test/TestGpuIndexFlat.cpp +311 -75
- data/vendor/faiss/faiss/gpu/test/TestUtils.cpp +10 -0
- data/vendor/faiss/faiss/gpu/test/TestUtils.h +3 -0
- data/vendor/faiss/faiss/gpu/test/demo_ivfpq_indexing_gpu.cpp +2 -2
- data/vendor/faiss/faiss/gpu/utils/DeviceUtils.h +5 -4
- data/vendor/faiss/faiss/impl/AdditiveQuantizer.cpp +331 -29
- data/vendor/faiss/faiss/impl/AdditiveQuantizer.h +110 -19
- data/vendor/faiss/faiss/impl/AuxIndexStructures.cpp +0 -54
- data/vendor/faiss/faiss/impl/AuxIndexStructures.h +0 -76
- data/vendor/faiss/faiss/impl/DistanceComputer.h +64 -0
- data/vendor/faiss/faiss/impl/HNSW.cpp +133 -32
- data/vendor/faiss/faiss/impl/HNSW.h +19 -16
- data/vendor/faiss/faiss/impl/IDSelector.cpp +125 -0
- data/vendor/faiss/faiss/impl/IDSelector.h +135 -0
- data/vendor/faiss/faiss/impl/LocalSearchQuantizer.cpp +378 -217
- data/vendor/faiss/faiss/impl/LocalSearchQuantizer.h +106 -29
- data/vendor/faiss/faiss/impl/LookupTableScaler.h +77 -0
- data/vendor/faiss/faiss/impl/NNDescent.cpp +1 -0
- data/vendor/faiss/faiss/impl/NSG.cpp +1 -4
- data/vendor/faiss/faiss/impl/NSG.h +1 -1
- data/vendor/faiss/faiss/impl/ProductAdditiveQuantizer.cpp +383 -0
- data/vendor/faiss/faiss/impl/ProductAdditiveQuantizer.h +154 -0
- data/vendor/faiss/faiss/impl/ProductQuantizer.cpp +225 -145
- data/vendor/faiss/faiss/impl/ProductQuantizer.h +29 -10
- data/vendor/faiss/faiss/impl/Quantizer.h +43 -0
- data/vendor/faiss/faiss/impl/ResidualQuantizer.cpp +521 -55
- data/vendor/faiss/faiss/impl/ResidualQuantizer.h +94 -16
- data/vendor/faiss/faiss/impl/ResultHandler.h +96 -0
- data/vendor/faiss/faiss/impl/ScalarQuantizer.cpp +108 -191
- data/vendor/faiss/faiss/impl/ScalarQuantizer.h +18 -18
- data/vendor/faiss/faiss/impl/index_read.cpp +338 -24
- data/vendor/faiss/faiss/impl/index_write.cpp +300 -18
- data/vendor/faiss/faiss/impl/io.cpp +1 -1
- data/vendor/faiss/faiss/impl/io_macros.h +20 -0
- data/vendor/faiss/faiss/impl/kmeans1d.cpp +303 -0
- data/vendor/faiss/faiss/impl/kmeans1d.h +48 -0
- data/vendor/faiss/faiss/impl/pq4_fast_scan.cpp +56 -16
- data/vendor/faiss/faiss/impl/pq4_fast_scan.h +25 -8
- data/vendor/faiss/faiss/impl/pq4_fast_scan_search_1.cpp +66 -25
- data/vendor/faiss/faiss/impl/pq4_fast_scan_search_qbs.cpp +75 -27
- data/vendor/faiss/faiss/index_factory.cpp +772 -412
- data/vendor/faiss/faiss/index_factory.h +3 -0
- data/vendor/faiss/faiss/index_io.h +5 -0
- data/vendor/faiss/faiss/invlists/DirectMap.cpp +1 -0
- data/vendor/faiss/faiss/invlists/InvertedLists.cpp +4 -1
- data/vendor/faiss/faiss/invlists/OnDiskInvertedLists.cpp +2 -1
- data/vendor/faiss/faiss/python/python_callbacks.cpp +27 -0
- data/vendor/faiss/faiss/python/python_callbacks.h +15 -0
- data/vendor/faiss/faiss/utils/Heap.h +31 -15
- data/vendor/faiss/faiss/utils/distances.cpp +384 -58
- data/vendor/faiss/faiss/utils/distances.h +149 -18
- data/vendor/faiss/faiss/utils/distances_simd.cpp +776 -6
- data/vendor/faiss/faiss/utils/extra_distances.cpp +12 -7
- data/vendor/faiss/faiss/utils/extra_distances.h +3 -1
- data/vendor/faiss/faiss/utils/fp16-fp16c.h +21 -0
- data/vendor/faiss/faiss/utils/fp16-inl.h +101 -0
- data/vendor/faiss/faiss/utils/fp16.h +11 -0
- data/vendor/faiss/faiss/utils/hamming-inl.h +54 -0
- data/vendor/faiss/faiss/utils/hamming.cpp +0 -48
- data/vendor/faiss/faiss/utils/ordered_key_value.h +10 -0
- data/vendor/faiss/faiss/utils/quantize_lut.cpp +62 -0
- data/vendor/faiss/faiss/utils/quantize_lut.h +20 -0
- data/vendor/faiss/faiss/utils/random.cpp +53 -0
- data/vendor/faiss/faiss/utils/random.h +5 -0
- data/vendor/faiss/faiss/utils/simdlib_avx2.h +4 -0
- data/vendor/faiss/faiss/utils/simdlib_emulated.h +6 -1
- data/vendor/faiss/faiss/utils/simdlib_neon.h +7 -2
- data/vendor/faiss/faiss/utils/utils.h +1 -1
- metadata +46 -5
- data/vendor/faiss/faiss/IndexResidual.cpp +0 -291
- data/vendor/faiss/faiss/IndexResidual.h +0 -152
data/vendor/faiss/faiss/Index.h
CHANGED
@@ -18,7 +18,7 @@
|
|
18
18
|
|
19
19
|
#define FAISS_VERSION_MAJOR 1
|
20
20
|
#define FAISS_VERSION_MINOR 7
|
21
|
-
#define FAISS_VERSION_PATCH
|
21
|
+
#define FAISS_VERSION_PATCH 3
|
22
22
|
|
23
23
|
/**
|
24
24
|
* @namespace faiss
|
@@ -38,11 +38,24 @@
|
|
38
38
|
|
39
39
|
namespace faiss {
|
40
40
|
|
41
|
-
/// Forward declarations see AuxIndexStructures.h
|
41
|
+
/// Forward declarations see impl/AuxIndexStructures.h, impl/IDSelector.h and
|
42
|
+
/// impl/DistanceComputer.h
|
42
43
|
struct IDSelector;
|
43
44
|
struct RangeSearchResult;
|
44
45
|
struct DistanceComputer;
|
45
46
|
|
47
|
+
/** Parent class for the optional search paramenters.
|
48
|
+
*
|
49
|
+
* Sub-classes with additional search parameters should inherit this class.
|
50
|
+
* Ownership of the object fields is always to the caller.
|
51
|
+
*/
|
52
|
+
struct SearchParameters {
|
53
|
+
/// if non-null, only these IDs will be considered during search.
|
54
|
+
IDSelector* sel = nullptr;
|
55
|
+
/// make sure we can dynamic_cast this
|
56
|
+
virtual ~SearchParameters() {}
|
57
|
+
};
|
58
|
+
|
46
59
|
/** Abstract structure for an index, supports adding vectors and searching them.
|
47
60
|
*
|
48
61
|
* All vectors provided at add or search time are 32-bit float arrays,
|
@@ -114,7 +127,8 @@ struct Index {
|
|
114
127
|
const float* x,
|
115
128
|
idx_t k,
|
116
129
|
float* distances,
|
117
|
-
idx_t* labels
|
130
|
+
idx_t* labels,
|
131
|
+
const SearchParameters* params = nullptr) const = 0;
|
118
132
|
|
119
133
|
/** query n vectors of dimension d to the index.
|
120
134
|
*
|
@@ -130,7 +144,8 @@ struct Index {
|
|
130
144
|
idx_t n,
|
131
145
|
const float* x,
|
132
146
|
float radius,
|
133
|
-
RangeSearchResult* result
|
147
|
+
RangeSearchResult* result,
|
148
|
+
const SearchParameters* params = nullptr) const;
|
134
149
|
|
135
150
|
/** return the indexes of the k vectors closest to the query x.
|
136
151
|
*
|
@@ -157,6 +172,16 @@ struct Index {
|
|
157
172
|
*/
|
158
173
|
virtual void reconstruct(idx_t key, float* recons) const;
|
159
174
|
|
175
|
+
/** Reconstruct several stored vectors (or an approximation if lossy coding)
|
176
|
+
*
|
177
|
+
* this function may not be defined for some indexes
|
178
|
+
* @param n number of vectors to reconstruct
|
179
|
+
* @param keys ids of the vectors to reconstruct (size n)
|
180
|
+
* @param recons reconstucted vector (size n * d)
|
181
|
+
*/
|
182
|
+
virtual void reconstruct_batch(idx_t n, const idx_t* keys, float* recons)
|
183
|
+
const;
|
184
|
+
|
160
185
|
/** Reconstruct vectors i0 to i0 + ni - 1
|
161
186
|
*
|
162
187
|
* this function may not be defined for some indexes
|
@@ -178,7 +203,8 @@ struct Index {
|
|
178
203
|
idx_t k,
|
179
204
|
float* distances,
|
180
205
|
idx_t* labels,
|
181
|
-
float* recons
|
206
|
+
float* recons,
|
207
|
+
const SearchParameters* params = nullptr) const;
|
182
208
|
|
183
209
|
/** Computes a residual vector after indexing encoding.
|
184
210
|
*
|
@@ -234,13 +260,24 @@ struct Index {
|
|
234
260
|
*/
|
235
261
|
virtual void sa_encode(idx_t n, const float* x, uint8_t* bytes) const;
|
236
262
|
|
237
|
-
/**
|
263
|
+
/** decode a set of vectors
|
238
264
|
*
|
239
265
|
* @param n number of vectors
|
240
266
|
* @param bytes input encoded vectors, size n * sa_code_size()
|
241
267
|
* @param x output vectors, size n * d
|
242
268
|
*/
|
243
269
|
virtual void sa_decode(idx_t n, const uint8_t* bytes, float* x) const;
|
270
|
+
|
271
|
+
/** moves the entries from another dataset to self.
|
272
|
+
* On output, other is empty.
|
273
|
+
* add_id is added to all moved ids
|
274
|
+
* (for sequential ids, this would be this->ntotal) */
|
275
|
+
virtual void merge_from(Index& otherIndex, idx_t add_id = 0);
|
276
|
+
|
277
|
+
/** check that the two indexes are compatible (ie, they are
|
278
|
+
* trained in the same way and have the same
|
279
|
+
* parameters). Otherwise throw. */
|
280
|
+
virtual void check_compatible_for_merge(const Index& otherIndex) const;
|
244
281
|
};
|
245
282
|
|
246
283
|
} // namespace faiss
|
@@ -30,16 +30,6 @@
|
|
30
30
|
#include <faiss/utils/distances.h>
|
31
31
|
#include <faiss/utils/utils.h>
|
32
32
|
|
33
|
-
/*
|
34
|
-
#include <faiss/utils/Heap.h>
|
35
|
-
|
36
|
-
#include <faiss/Clustering.h>
|
37
|
-
|
38
|
-
#include <faiss/utils/hamming.h>
|
39
|
-
|
40
|
-
|
41
|
-
*/
|
42
|
-
|
43
33
|
namespace faiss {
|
44
34
|
|
45
35
|
/*************************************
|
@@ -52,7 +42,7 @@ Index2Layer::Index2Layer(
|
|
52
42
|
int M,
|
53
43
|
int nbit,
|
54
44
|
MetricType metric)
|
55
|
-
:
|
45
|
+
: IndexFlatCodes(0, quantizer->d, metric),
|
56
46
|
q1(quantizer, nlist),
|
57
47
|
pq(quantizer->d, M, nbit) {
|
58
48
|
is_trained = false;
|
@@ -116,83 +106,16 @@ void Index2Layer::train(idx_t n, const float* x) {
|
|
116
106
|
is_trained = true;
|
117
107
|
}
|
118
108
|
|
119
|
-
void Index2Layer::add(idx_t n, const float* x) {
|
120
|
-
idx_t bs = 32768;
|
121
|
-
if (n > bs) {
|
122
|
-
for (idx_t i0 = 0; i0 < n; i0 += bs) {
|
123
|
-
idx_t i1 = std::min(i0 + bs, n);
|
124
|
-
if (verbose) {
|
125
|
-
printf("Index2Layer::add: adding %" PRId64 ":%" PRId64
|
126
|
-
" / %" PRId64 "\n",
|
127
|
-
i0,
|
128
|
-
i1,
|
129
|
-
n);
|
130
|
-
}
|
131
|
-
add(i1 - i0, x + i0 * d);
|
132
|
-
}
|
133
|
-
return;
|
134
|
-
}
|
135
|
-
|
136
|
-
std::vector<idx_t> codes1(n);
|
137
|
-
q1.quantizer->assign(n, x, codes1.data());
|
138
|
-
std::vector<float> residuals(n * d);
|
139
|
-
for (idx_t i = 0; i < n; i++) {
|
140
|
-
q1.quantizer->compute_residual(
|
141
|
-
x + i * d, residuals.data() + i * d, codes1[i]);
|
142
|
-
}
|
143
|
-
std::vector<uint8_t> codes2(n * code_size_2);
|
144
|
-
|
145
|
-
pq.compute_codes(residuals.data(), codes2.data(), n);
|
146
|
-
|
147
|
-
codes.resize((ntotal + n) * code_size);
|
148
|
-
uint8_t* wp = &codes[ntotal * code_size];
|
149
|
-
|
150
|
-
{
|
151
|
-
int i = 0x11223344;
|
152
|
-
const char* ip = (char*)&i;
|
153
|
-
FAISS_THROW_IF_NOT_MSG(
|
154
|
-
ip[0] == 0x44, "works only on a little-endian CPU");
|
155
|
-
}
|
156
|
-
|
157
|
-
// copy to output table
|
158
|
-
for (idx_t i = 0; i < n; i++) {
|
159
|
-
memcpy(wp, &codes1[i], code_size_1);
|
160
|
-
wp += code_size_1;
|
161
|
-
memcpy(wp, &codes2[i * code_size_2], code_size_2);
|
162
|
-
wp += code_size_2;
|
163
|
-
}
|
164
|
-
|
165
|
-
ntotal += n;
|
166
|
-
}
|
167
|
-
|
168
109
|
void Index2Layer::search(
|
169
110
|
idx_t /*n*/,
|
170
111
|
const float* /*x*/,
|
171
112
|
idx_t /*k*/,
|
172
113
|
float* /*distances*/,
|
173
|
-
idx_t* /*labels
|
114
|
+
idx_t* /*labels*/,
|
115
|
+
const SearchParameters* /* params */) const {
|
174
116
|
FAISS_THROW_MSG("not implemented");
|
175
117
|
}
|
176
118
|
|
177
|
-
void Index2Layer::reconstruct_n(idx_t i0, idx_t ni, float* recons) const {
|
178
|
-
std::vector<float> recons1(d);
|
179
|
-
FAISS_THROW_IF_NOT(i0 >= 0 && i0 + ni <= ntotal);
|
180
|
-
const uint8_t* rp = &codes[i0 * code_size];
|
181
|
-
|
182
|
-
for (idx_t i = 0; i < ni; i++) {
|
183
|
-
idx_t key = 0;
|
184
|
-
memcpy(&key, rp, code_size_1);
|
185
|
-
q1.quantizer->reconstruct(key, recons1.data());
|
186
|
-
rp += code_size_1;
|
187
|
-
pq.decode(rp, recons);
|
188
|
-
for (idx_t j = 0; j < d; j++) {
|
189
|
-
recons[j] += recons1[j];
|
190
|
-
}
|
191
|
-
rp += code_size_2;
|
192
|
-
recons += d;
|
193
|
-
}
|
194
|
-
}
|
195
|
-
|
196
119
|
void Index2Layer::transfer_to_IVFPQ(IndexIVFPQ& other) const {
|
197
120
|
FAISS_THROW_IF_NOT(other.nlist == q1.nlist);
|
198
121
|
FAISS_THROW_IF_NOT(other.code_size == code_size_2);
|
@@ -211,15 +134,6 @@ void Index2Layer::transfer_to_IVFPQ(IndexIVFPQ& other) const {
|
|
211
134
|
other.ntotal = ntotal;
|
212
135
|
}
|
213
136
|
|
214
|
-
void Index2Layer::reconstruct(idx_t key, float* recons) const {
|
215
|
-
reconstruct_n(key, 1, recons);
|
216
|
-
}
|
217
|
-
|
218
|
-
void Index2Layer::reset() {
|
219
|
-
ntotal = 0;
|
220
|
-
codes.clear();
|
221
|
-
}
|
222
|
-
|
223
137
|
namespace {
|
224
138
|
|
225
139
|
struct Distance2Level : DistanceComputer {
|
@@ -259,7 +173,7 @@ struct DistanceXPQ4 : Distance2Level {
|
|
259
173
|
|
260
174
|
FAISS_ASSERT(quantizer);
|
261
175
|
M = storage.pq.M;
|
262
|
-
pq_l1_tab = quantizer->
|
176
|
+
pq_l1_tab = quantizer->get_xb();
|
263
177
|
}
|
264
178
|
|
265
179
|
float operator()(idx_t i) override {
|
@@ -368,12 +282,29 @@ DistanceComputer* Index2Layer::get_distance_computer() const {
|
|
368
282
|
}
|
369
283
|
|
370
284
|
/* The standalone codec interface */
|
371
|
-
|
372
|
-
|
373
|
-
|
285
|
+
|
286
|
+
// block size used in Index2Layer::sa_encode
|
287
|
+
int index2layer_sa_encode_bs = 32768;
|
374
288
|
|
375
289
|
void Index2Layer::sa_encode(idx_t n, const float* x, uint8_t* bytes) const {
|
376
290
|
FAISS_THROW_IF_NOT(is_trained);
|
291
|
+
|
292
|
+
idx_t bs = index2layer_sa_encode_bs;
|
293
|
+
if (n > bs) {
|
294
|
+
for (idx_t i0 = 0; i0 < n; i0 += bs) {
|
295
|
+
idx_t i1 = std::min(i0 + bs, n);
|
296
|
+
if (verbose) {
|
297
|
+
printf("Index2Layer::add: adding %" PRId64 ":%" PRId64
|
298
|
+
" / %" PRId64 "\n",
|
299
|
+
i0,
|
300
|
+
i1,
|
301
|
+
n);
|
302
|
+
}
|
303
|
+
sa_encode(i1 - i0, x + i0 * d, bytes + i0 * code_size);
|
304
|
+
}
|
305
|
+
return;
|
306
|
+
}
|
307
|
+
|
377
308
|
std::unique_ptr<int64_t[]> list_nos(new int64_t[n]);
|
378
309
|
q1.quantizer->assign(n, x, list_nos.get());
|
379
310
|
std::vector<float> residuals(n * d);
|
@@ -11,8 +11,10 @@
|
|
11
11
|
|
12
12
|
#include <vector>
|
13
13
|
|
14
|
+
#include <faiss/IndexFlatCodes.h>
|
14
15
|
#include <faiss/IndexIVF.h>
|
15
16
|
#include <faiss/IndexPQ.h>
|
17
|
+
#include <faiss/impl/platform_macros.h>
|
16
18
|
|
17
19
|
namespace faiss {
|
18
20
|
|
@@ -24,25 +26,19 @@ struct IndexIVFPQ;
|
|
24
26
|
* The class is mainly inteded to store encoded vectors that can be
|
25
27
|
* accessed randomly, the search function is not implemented.
|
26
28
|
*/
|
27
|
-
struct Index2Layer :
|
29
|
+
struct Index2Layer : IndexFlatCodes {
|
28
30
|
/// first level quantizer
|
29
31
|
Level1Quantizer q1;
|
30
32
|
|
31
33
|
/// second level quantizer is always a PQ
|
32
34
|
ProductQuantizer pq;
|
33
35
|
|
34
|
-
/// Codes. Size ntotal * code_size.
|
35
|
-
std::vector<uint8_t> codes;
|
36
|
-
|
37
36
|
/// size of the code for the first level (ceil(log8(q1.nlist)))
|
38
37
|
size_t code_size_1;
|
39
38
|
|
40
39
|
/// size of the code for the second level
|
41
40
|
size_t code_size_2;
|
42
41
|
|
43
|
-
/// code_size_1 + code_size_2
|
44
|
-
size_t code_size;
|
45
|
-
|
46
42
|
Index2Layer(
|
47
43
|
Index* quantizer,
|
48
44
|
size_t nlist,
|
@@ -55,21 +51,14 @@ struct Index2Layer : Index {
|
|
55
51
|
|
56
52
|
void train(idx_t n, const float* x) override;
|
57
53
|
|
58
|
-
void add(idx_t n, const float* x) override;
|
59
|
-
|
60
54
|
/// not implemented
|
61
55
|
void search(
|
62
56
|
idx_t n,
|
63
57
|
const float* x,
|
64
58
|
idx_t k,
|
65
59
|
float* distances,
|
66
|
-
idx_t* labels
|
67
|
-
|
68
|
-
void reconstruct_n(idx_t i0, idx_t ni, float* recons) const override;
|
69
|
-
|
70
|
-
void reconstruct(idx_t key, float* recons) const override;
|
71
|
-
|
72
|
-
void reset() override;
|
60
|
+
idx_t* labels,
|
61
|
+
const SearchParameters* params = nullptr) const override;
|
73
62
|
|
74
63
|
DistanceComputer* get_distance_computer() const override;
|
75
64
|
|
@@ -77,9 +66,11 @@ struct Index2Layer : Index {
|
|
77
66
|
void transfer_to_IVFPQ(IndexIVFPQ& other) const;
|
78
67
|
|
79
68
|
/* The standalone codec interface */
|
80
|
-
size_t sa_code_size() const override;
|
81
69
|
void sa_encode(idx_t n, const float* x, uint8_t* bytes) const override;
|
82
70
|
void sa_decode(idx_t n, const uint8_t* bytes, float* x) const override;
|
83
71
|
};
|
84
72
|
|
73
|
+
// block size used in Index2Layer::sa_encode
|
74
|
+
FAISS_API extern int index2layer_sa_encode_bs;
|
75
|
+
|
85
76
|
} // namespace faiss
|