From a0b50e669f5e362839adfc368f5687ba92626abf Mon Sep 17 00:00:00 2001 From: Matthijs Douze Date: Tue, 7 Dec 2021 04:34:51 -0800 Subject: [PATCH] Re-factor factory string parsing (#2134) Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/2134 The old parsing was very complex and grew out of hand. this diff just uses regex parsing. Reviewed By: wickedfoo Differential Revision: D32759110 fbshipit-source-id: 243029bba8a7fe70c71323f5edc7e2ce4e669757 --- .clang-format | 4 +- faiss/index_factory.cpp | 1090 ++++++++++++++++++++++----------------- faiss/index_factory.h | 3 + tests/test_factory.py | 73 ++- 4 files changed, 680 insertions(+), 490 deletions(-) diff --git a/.clang-format b/.clang-format index 6b808d997..1fe6508fa 100644 --- a/.clang-format +++ b/.clang-format @@ -15,8 +15,8 @@ AllowShortLoopsOnASingleLine: false AlwaysBreakAfterReturnType: None AlwaysBreakBeforeMultilineStrings: true AlwaysBreakTemplateDeclarations: true -BinPackArguments: false -BinPackParameters: false +BinPackArguments: false # at some point, set this to true +BinPackParameters: false # at some point, set this to true BraceWrapping: AfterClass: false AfterControlStatement: false diff --git a/faiss/index_factory.cpp b/faiss/index_factory.cpp index 2031b4fb4..5aec597ec 100644 --- a/faiss/index_factory.cpp +++ b/faiss/index_factory.cpp @@ -5,19 +5,19 @@ * LICENSE file in the root directory of this source tree. */ -// -*- c++ -*- - /* - * implementation of Hyper-parameter auto-tuning + * implementation of the index_factory function. Lots of regex parsing code. */ #include - -#include +#include "faiss/MetricType.h" +#include "faiss/impl/FaissAssert.h" #include #include +#include + #include #include @@ -34,6 +34,7 @@ #include #include #include +#include #include #include #include @@ -49,6 +50,7 @@ #include #include #include +#include namespace faiss { @@ -56,16 +58,49 @@ namespace faiss { * index_factory ***************************************************************/ +int index_factory_verbose = 0; + namespace { -struct VTChain { - std::vector chain; - ~VTChain() { - for (int i = 0; i < chain.size(); i++) { - delete chain[i]; +/*************************************************************** + * Small functions + */ + +bool re_match(const std::string& s, const std::string& pat, std::smatch& sm) { + return std::regex_match(s, sm, std::regex(pat)); +} + +// find first pair of matching parentheses +void find_matching_parentheses( + const std::string& s, + int& i0, + int& i1, + int begin = 0) { + int st = 0; + i0 = i1 = 0; + for (int i = begin; i < s.length(); i++) { + if (s[i] == '(') { + if (st == 0) { + i0 = i; + } + st++; + } + + if (s[i] == ')') { + st--; + if (st == 0) { + i1 = i; + return; + } + if (st < 0) { + FAISS_THROW_FMT( + "factory string %s: unbalanced parentheses", s.c_str()); + } } } -}; + FAISS_THROW_FMT( + "factory string %s: unbalanced parentheses st=%d", s.c_str(), st); +} /// what kind of training does this coarse quantizer require? char get_trains_alone(const Index* coarse_quantizer) { @@ -92,494 +127,575 @@ IndexIVF* fix_ivf_fields(IndexIVF* index_ivf) { return index_ivf; } -bool str_ends_with(const std::string& s, const std::string& suffix) { - return s.rfind(suffix) == std::abs(int(s.size() - suffix.size())); -} - -// check if ends with suffix followed by digits -bool str_ends_with_digits(const std::string& s, const std::string& suffix) { - int i; - for (i = s.length() - 1; i >= 0; i--) { - if (!isdigit(s[i])) - break; +int mres_to_int(const std::ssub_match& mr, int deflt = -1, int begin = 0) { + if (mr.length() == 0) { + return deflt; } - return str_ends_with(s.substr(0, i + 1), suffix); + return std::stoi(mr.str().substr(begin)); } -void find_matching_parentheses(const std::string& s, int& i0, int& i1) { - int st = 0; - i0 = i1 = 0; - for (int i = 0; i < s.length(); i++) { - if (s[i] == '(') { - if (st == 0) { - i0 = i; - } - st++; - } +std::map sq_types = { + {"SQ8", ScalarQuantizer::QT_8bit}, + {"SQ4", ScalarQuantizer::QT_4bit}, + {"SQ6", ScalarQuantizer::QT_6bit}, + {"SQfp16", ScalarQuantizer::QT_fp16}, +}; +const std::string sq_pattern = "(SQ4|SQ8|SQ6|SQfp16)"; - if (s[i] == ')') { - st--; - if (st == 0) { - i1 = i; - return; - } - if (st < 0) { - FAISS_THROW_FMT( - "factory string %s: unbalanced parentheses", s.c_str()); - } - } - } - FAISS_THROW_FMT( - "factory string %s: unbalanced parentheses st=%d", s.c_str(), st); -} +std::map aq_search_type = { + {"_Nfloat", AdditiveQuantizer::ST_norm_float}, + {"_Nnone", AdditiveQuantizer::ST_LUT_nonorm}, + {"_Nqint8", AdditiveQuantizer::ST_norm_qint8}, + {"_Nqint4", AdditiveQuantizer::ST_norm_qint4}, + {"_Ncqint8", AdditiveQuantizer::ST_norm_cqint8}, + {"_Ncqint4", AdditiveQuantizer::ST_norm_cqint4}, +}; -AdditiveQuantizer::Search_type_t parse_AQ_search_type( +const std::string aq_def_pattern = "[0-9]+x[0-9]+(_[0-9]+x[0-9]+)*"; +const std::string aq_norm_pattern = + "(|_Nnone|_Nfloat|_Nqint8|_Nqint4|_Ncqint8|_Ncqint4)"; + +AdditiveQuantizer::Search_type_t aq_parse_search_type( std::string stok, MetricType metric) { - if (str_ends_with(stok, "_Nfloat")) { - return AdditiveQuantizer::ST_norm_float; - } else if (str_ends_with(stok, "_Nnone")) { - return AdditiveQuantizer::ST_LUT_nonorm; - } else if (str_ends_with(stok, "_Nqint8")) { - return AdditiveQuantizer::ST_norm_qint8; - } else if (str_ends_with(stok, "_Nqint4")) { - return AdditiveQuantizer::ST_norm_qint4; - } else if (str_ends_with(stok, "_Ncqint8")) { - return AdditiveQuantizer::ST_norm_cqint8; - } else if (str_ends_with(stok, "_Ncqint4")) { - return AdditiveQuantizer::ST_norm_cqint4; - } else if (metric == METRIC_L2) { - return AdditiveQuantizer::ST_decompress; - } else { - return AdditiveQuantizer::ST_LUT_nonorm; + if (stok == "") { + return metric == METRIC_L2 ? AdditiveQuantizer::ST_decompress + : AdditiveQuantizer::ST_LUT_nonorm; } + int pos = stok.rfind("_"); + return aq_search_type[stok.substr(pos)]; +} + +std::vector aq_parse_nbits(std::string stok) { + std::vector nbits; + std::smatch sm; + while (std::regex_search(stok, sm, std::regex("([0-9]+)x([0-9]+)"))) { + int M = std::stoi(sm[1].str()); + int nbit = std::stoi(sm[2].str()); + nbits.resize(nbits.size() + M, nbit); + stok = sm.suffix(); + } + return nbits; +} + +/*************************************************************** + * Parse VectorTransform + */ + +VectorTransform* parse_VectorTransform(const std::string& description, int d) { + std::smatch sm; + auto match = [&sm, description](std::string pattern) { + return re_match(description, pattern, sm); + }; + if (match("PCA(W?)(R?)([0-9]+)")) { + bool white = sm[1].length() > 0; + bool rot = sm[2].length() > 0; + return new PCAMatrix(d, std::stoi(sm[3].str()), white ? -0.5 : 0, rot); + } + if (match("L2[nN]orm")) { + return new NormalizationTransform(d, 2.0); + } + if (match("RR([0-9]+)?")) { + return new RandomRotationMatrix(d, mres_to_int(sm[1], d)); + } + if (match("ITQ([0-9]+)?")) { + return new ITQTransform(d, mres_to_int(sm[1], d), sm[1].length() > 0); + } + if (match("OPQ([0-9]+)(_[0-9]+)?")) { + int M = std::stoi(sm[1].str()); + int d_out = mres_to_int(sm[2], d); + return new OPQMatrix(d, M, d_out); + } + if (match("Pad([0-9]+)")) { + int d_out = std::stoi(sm[1].str()); + return new RemapDimensionsTransform(d, std::max(d_out, d), false); + } + return nullptr; +}; + +/*************************************************************** + * Parse IndexIVF + */ + +// parsing guard + function +Index* parse_coarse_quantizer( + const std::string& description, + int d, + MetricType mt, + std::vector>& parenthesis_indexes, + size_t& nlist, + bool& use_2layer) { + std::smatch sm; + auto match = [&sm, description](std::string pattern) { + return re_match(description, pattern, sm); + }; + use_2layer = false; + + if (match("IVF([0-9]+)")) { + nlist = std::stoi(sm[1].str()); + return new IndexFlat(d, mt); + } + if (match("IMI2x([0-9]+)")) { + int nbit = std::stoi(sm[1].str()); + FAISS_THROW_IF_NOT_MSG( + mt == METRIC_L2, + "MultiIndex not implemented for inner prod search"); + nlist = (size_t)1 << (2 * nbit); + return new MultiIndexQuantizer(d, 2, nbit); + } + if (match("IVF([0-9]+)_HNSW([0-9]*)")) { + nlist = std::stoi(sm[1].str()); + int hnsw_M = sm[2].length() > 0 ? std::stoi(sm[2]) : 32; + return new IndexHNSWFlat(d, hnsw_M, mt); + } + if (match("IVF([0-9]+)_NSG([0-9]+)")) { + nlist = std::stoi(sm[1].str()); + int R = std::stoi(sm[2]); + return new IndexNSGFlat(d, R, mt); + } + if (match("IVF([0-9]+)\\(Index([0-9])\\)")) { + nlist = std::stoi(sm[1].str()); + int no = std::stoi(sm[2].str()); + FAISS_ASSERT(no >= 0 && no < parenthesis_indexes.size()); + return parenthesis_indexes[no].release(); + } + + // these two generate Index2Layer's not IndexIVF's + if (match("Residual([0-9]+)x([0-9]+)")) { + FAISS_THROW_IF_NOT_MSG( + mt == METRIC_L2, + "MultiIndex not implemented for inner prod search"); + int M = mres_to_int(sm[1]), nbit = mres_to_int(sm[2]); + nlist = (size_t)1 << (M * nbit); + use_2layer = true; + return new MultiIndexQuantizer(d, M, nbit); + } + if (match("Residual([0-9]+)")) { + FAISS_THROW_IF_NOT_MSG( + mt == METRIC_L2, + "Residual not implemented for inner prod search"); + use_2layer = true; + nlist = mres_to_int(sm[1]); + return new IndexFlatL2(d); + } + return nullptr; +} + +// parse the code description part of the IVF description + +IndexIVF* parse_IndexIVF( + const std::string& code_string, + std::unique_ptr& quantizer, + size_t nlist, + MetricType mt) { + std::smatch sm; + auto match = [&sm, &code_string](const std::string pattern) { + return re_match(code_string, pattern, sm); + }; + auto get_q = [&quantizer] { return quantizer.release(); }; + int d = quantizer->d; + + if (match("Flat")) { + return new IndexIVFFlat(get_q(), d, nlist, mt); + } + if (match("FlatDedup")) { + return new IndexIVFFlatDedup(get_q(), d, nlist, mt); + } + if (match(sq_pattern)) { + return new IndexIVFScalarQuantizer( + get_q(), d, nlist, sq_types[sm[1].str()], mt); + } + if (match("PQ([0-9]+)(x[0-9]+)?(np)?")) { + int M = mres_to_int(sm[1]), nbit = mres_to_int(sm[2], 8, 1); + IndexIVFPQ* index_ivf = new IndexIVFPQ(get_q(), d, nlist, M, nbit, mt); + index_ivf->do_polysemous_training = sm[3].str() != "np"; + return index_ivf; + } + if (match("PQ([0-9]+)\\+([0-9]+)")) { + FAISS_THROW_IF_NOT_MSG( + mt == METRIC_L2, + "IVFPQR not implemented for inner product search"); + int M1 = mres_to_int(sm[1]), M2 = mres_to_int(sm[2]); + return new IndexIVFPQR(get_q(), d, nlist, M1, 8, M2, 8); + } + if (match("PQ([0-9]+)x4fs(r?)(_[0-9]+)?")) { + int M = mres_to_int(sm[1]); + int bbs = mres_to_int(sm[3], 32, 1); + IndexIVFPQFastScan* index_ivf = + new IndexIVFPQFastScan(get_q(), d, nlist, M, 4, mt, bbs); + index_ivf->by_residual = sm[2].str() == "r"; + return index_ivf; + } + if (match("(RQ|LSQ)" + aq_def_pattern + aq_norm_pattern)) { + std::vector nbits = aq_parse_nbits(sm.str()); + AdditiveQuantizer::Search_type_t st = + aq_parse_search_type(sm[sm.size() - 1].str(), mt); + IndexIVF* index_ivf; + if (sm[1].str() == "RQ") { + index_ivf = new IndexIVFResidualQuantizer( + get_q(), d, nlist, nbits, mt, st); + } else { + FAISS_THROW_IF_NOT(nbits.size() > 0); + index_ivf = new IndexIVFLocalSearchQuantizer( + get_q(), d, nlist, nbits.size(), nbits[0], mt, st); + } + return index_ivf; + } + if (match("(ITQ|PCA|PCAR)([0-9]+)?,SH([-0-9.e]+)?([gcm])?")) { + int outdim = mres_to_int(sm[2], d); // is also the number of bits + std::unique_ptr vt; + if (sm[1] == "ITQ") { + vt.reset(new ITQTransform(d, outdim, d != outdim)); + } else if (sm[1] == "PCA") { + vt.reset(new PCAMatrix(d, outdim)); + } else if (sm[1] == "PCAR") { + vt.reset(new PCAMatrix(d, outdim, 0, true)); + } + // the rationale for -1e10 is that this corresponds to simple + // thresholding + float period = sm[3].length() > 0 ? std::stof(sm[3]) : -1e10; + IndexIVFSpectralHash* index_ivf = + new IndexIVFSpectralHash(get_q(), d, nlist, outdim, period); + index_ivf->replace_vt(vt.release(), true); + if (sm[4].length()) { + std::string s = sm[4].str(); + index_ivf->threshold_type = s == "g" + ? IndexIVFSpectralHash::Thresh_global + : s == "c" + ? IndexIVFSpectralHash::Thresh_centroid + : + /* s == "m" ? */ IndexIVFSpectralHash::Thresh_median; + } + return index_ivf; + } + return nullptr; +} + +/*************************************************************** + * Parse IndexHNSW + */ + +IndexHNSW* parse_IndexHNSW( + const std::string code_string, + int d, + MetricType mt, + int hnsw_M) { + std::smatch sm; + auto match = [&sm, &code_string](const std::string& pattern) { + return re_match(code_string, pattern, sm); + }; + + if (match("Flat|")) { + return new IndexHNSWFlat(d, hnsw_M, mt); + } + if (match("PQ([0-9]+)(np)?")) { + int M = std::stoi(sm[1].str()); + IndexHNSWPQ* ipq = new IndexHNSWPQ(d, M, hnsw_M); + dynamic_cast(ipq->storage)->do_polysemous_training = + sm[2].str() != "np"; + return ipq; + } + if (match(sq_pattern)) { + return new IndexHNSWSQ(d, sq_types[sm[1].str()], hnsw_M, mt); + } + if (match("([0-9]+)\\+PQ([0-9]+)?")) { + int ncent = mres_to_int(sm[1]); + int pq_m = mres_to_int(sm[2]); + IndexHNSW2Level* hidx2l = + new IndexHNSW2Level(new IndexFlatL2(d), ncent, pq_m, hnsw_M); + dynamic_cast(hidx2l->storage)->q1.own_fields = true; + return hidx2l; + } + if (match("2x([0-9]+)\\+PQ([0-9]+)?")) { + int nbit = mres_to_int(sm[1]); + int pq_m = mres_to_int(sm[2]); + Index* quant = new MultiIndexQuantizer(d, 2, nbit); + IndexHNSW2Level* hidx2l = new IndexHNSW2Level( + quant, (size_t)1 << (2 * nbit), pq_m, hnsw_M); + Index2Layer* idx2l = dynamic_cast(hidx2l->storage); + idx2l->q1.own_fields = true; + idx2l->q1.quantizer_trains_alone = 1; + return hidx2l; + } + + return nullptr; +} + +/*************************************************************** + * Parse basic indexes + */ + +Index* parse_other_indexes( + const std::string& description, + int d, + MetricType metric) { + std::smatch sm; + auto match = [&sm, description](const std::string& pattern) { + return re_match(description, pattern, sm); + }; + + // IndexFlat + if (description == "Flat") { + return new IndexFlat(d, metric); + } + + // IndexLSH + if (match("LSH(r?)(t?)")) { + bool rotate_data = sm[1].length() > 0; + bool train_thresholds = sm[2].length() > 0; + FAISS_THROW_IF_NOT(metric == METRIC_L2); + return new IndexLSH(d, d, rotate_data, train_thresholds); + } + + // IndexLattice + if (match("ZnLattice([0-9]+)x([0-9]+)_([0-9]+)")) { + int M = std::stoi(sm[1].str()), r2 = std::stoi(sm[2].str()); + int nbit = std::stoi(sm[3].str()); + return new IndexLattice(d, M, nbit, r2); + } + + // IndexNSGFlat + if (match("NSG([0-9]+)(,Flat)?")) { + return new IndexNSGFlat(d, std::stoi(sm[1].str()), metric); + } + + // IndexScalarQuantizer + if (match(sq_pattern)) { + return new IndexScalarQuantizer(d, sq_types[description], metric); + } + + // IndexPQ + if (match("PQ([0-9]+)(x[0-9]+)?(np)?")) { + int M = std::stoi(sm[1].str()); + int nbit = mres_to_int(sm[2], 8, 1); + IndexPQ* index_pq = new IndexPQ(d, M, nbit, metric); + index_pq->do_polysemous_training = sm[3].str() != "np"; + return index_pq; + } + + // IndexPQFastScan + if (match("PQ([0-9]+)x4fs(_[0-9]+)?")) { + int M = std::stoi(sm[1].str()); + int bbs = mres_to_int(sm[2], 32, 1); + return new IndexPQFastScan(d, M, 4, metric, bbs); + } + + // IndexResidualCoarseQuantizer and IndexResidualQuantizer + std::string pattern = "(RQ|RCQ)" + aq_def_pattern + aq_norm_pattern; + if (match(pattern)) { + std::vector nbits = aq_parse_nbits(description); + if (sm[1].str() == "RCQ") { + return new ResidualCoarseQuantizer(d, nbits, metric); + } + AdditiveQuantizer::Search_type_t st = + aq_parse_search_type(sm[sm.size() - 1].str(), metric); + return new IndexResidualQuantizer(d, nbits, metric, st); + } + + // LocalSearchCoarseQuantizer and IndexLocalSearchQuantizer + if (match("(LSQ|LSCQ)([0-9]+)x([0-9]+)" + aq_norm_pattern)) { + std::vector nbits = aq_parse_nbits(description); + int M = mres_to_int(sm[2]); + int nbit = mres_to_int(sm[3]); + if (sm[1].str() == "LSCQ") { + return new LocalSearchCoarseQuantizer(d, M, nbit, metric); + } + AdditiveQuantizer::Search_type_t st = + aq_parse_search_type(sm[sm.size() - 1].str(), metric); + return new IndexLocalSearchQuantizer(d, M, nbit, metric, st); + } + + return nullptr; +} + +/*************************************************************** + * Driver function + */ +std::unique_ptr index_factory_sub( + int d, + std::string description, + MetricType metric) { + // handle composite indexes + + bool verbose = index_factory_verbose; + + if (verbose) { + printf("begin parse VectorTransforms: %s \n", description.c_str()); + } + + // for the current match + std::smatch sm; + + // handle refines + if (re_match(description, "(.+),RFlat", sm) || + re_match(description, "(.+),Refine\\((.+)\\)", sm)) { + std::unique_ptr filter_index = + index_factory_sub(d, sm[1].str(), metric); + std::unique_ptr refine_index; + + if (sm.size() == 3) { // Refine + refine_index = index_factory_sub(d, sm[2].str(), metric); + } else { // RFlat + refine_index.reset(new IndexFlat(d, metric)); + } + IndexRefine* index_rf = + new IndexRefine(filter_index.get(), refine_index.get()); + index_rf->own_fields = true; + filter_index.release(); + refine_index.release(); + index_rf->own_refine_index = true; + return std::unique_ptr(index_rf); + } + + // IndexPreTransform + // should handle this first (even before parentheses) because it changes d + std::vector> vts; + VectorTransform* vt = nullptr; + while (re_match(description, "([^,]+),(.*)", sm) && + (vt = parse_VectorTransform(sm[1], d))) { + // reset loop + description = sm[sm.size() - 1]; + vts.emplace_back(vt); + d = vts.back()->d_out; + + } + + if (vts.size() > 0) { + std::unique_ptr sub_index = + index_factory_sub(d, description, metric); + IndexPreTransform* index_pt = new IndexPreTransform(sub_index.get()); + std::unique_ptr ret(index_pt); + index_pt->own_fields = true; + sub_index.release(); + while (vts.size() > 0) { + if (verbose) { + printf("prepend trans %d -> %d\n", + vts.back()->d_in, + vts.back()->d_out); + } + index_pt->prepend_transform(vts.back().release()); + vts.pop_back(); + } + return ret; + } + + // what we got from the parentheses + std::vector> parenthesis_indexes; + + int begin = 0; + while (description.find('(', begin) != std::string::npos) { + // replace indexes in () with Index0, Index1, etc. + int i0, i1; + find_matching_parentheses(description, i0, i1, begin); + std::string sub_description = description.substr(i0 + 1, i1 - i0 - 1); + int no = parenthesis_indexes.size(); + parenthesis_indexes.push_back( + index_factory_sub(d, sub_description, metric)); + description = description.substr(0, i0 + 1) + "Index" + + std::to_string(no) + description.substr(i1); + begin = i1 + 1; + } + + if (verbose) { + printf("after () normalization: %s %ld parenthesis indexes d=%d\n", + description.c_str(), + parenthesis_indexes.size(), + d); + } + + // IndexIDMap -- it turns out is was used both as a prefix and a suffix, so + // support both + if (re_match(description, "(.+),IDMap", sm) || + re_match(description, "IDMap,(.+)", sm)) { + IndexIDMap* idmap = new IndexIDMap( + index_factory_sub(d, sm[1].str(), metric).release()); + idmap->own_fields = true; + return std::unique_ptr(idmap); + } + + { // handle basic index types + Index* index = parse_other_indexes(description, d, metric); + if (index) { + return std::unique_ptr(index); + } + } + + // HNSW variants (it was unclear in the old version that the separator was a + // "," so we support both "_" and ",") + if (re_match(description, "HNSW([0-9]*)([,_].*)?", sm)) { + int hnsw_M = mres_to_int(sm[1], 32); + // We also accept empty code string (synonym of Flat) + std::string code_string = + sm[2].length() > 0 ? sm[2].str().substr(1) : ""; + if (verbose) { + printf("parsing HNSW string %s code_string=%s hnsw_M=%d\n", + description.c_str(), + code_string.c_str(), + hnsw_M); + } + + IndexHNSW* index = parse_IndexHNSW(code_string, d, metric, hnsw_M); + FAISS_THROW_IF_NOT_FMT( + index, + "could not parse HNSW code description %s in %s", + code_string.c_str(), + description.c_str()); + return std::unique_ptr(index); + } + + // IndexIVF + { + size_t nlist; + bool use_2layer; + size_t comma = description.find(","); + std::string coarse_string = description.substr(0, comma); + // Match coarse quantizer part first + std::unique_ptr quantizer(parse_coarse_quantizer( + description.substr(0, comma), + d, + metric, + parenthesis_indexes, + nlist, + use_2layer)); + + if (comma != std::string::npos && quantizer.get()) { + std::string code_description = description.substr(comma + 1); + if (use_2layer) { + bool ok = + re_match(code_description, "PQ([0-9]+)(x[0-9]+)?", sm); + FAISS_THROW_IF_NOT_FMT( + ok, + "could not parse 2 layer code description %s in %s", + code_description.c_str(), + description.c_str()); + int M = std::stoi(sm[1].str()), nbit = mres_to_int(sm[2], 8, 1); + Index2Layer* index_2l = + new Index2Layer(quantizer.release(), nlist, M, nbit); + index_2l->q1.own_fields = true; + index_2l->q1.quantizer_trains_alone = + get_trains_alone(index_2l->q1.quantizer); + return std::unique_ptr(index_2l); + } + + IndexIVF* index_ivf = + parse_IndexIVF(code_description, quantizer, nlist, metric); + + FAISS_THROW_IF_NOT_FMT( + index_ivf, + "could not parse code description %s in %s", + code_description.c_str(), + description.c_str()); + return std::unique_ptr(fix_ivf_fields(index_ivf)); + } + } + FAISS_THROW_FMT("could not parse index string %s", description.c_str()); + return nullptr; } } // anonymous namespace -Index* index_factory(int d, const char* description_in, MetricType metric) { - FAISS_THROW_IF_NOT(metric == METRIC_L2 || metric == METRIC_INNER_PRODUCT); - VTChain vts; - Index* coarse_quantizer = nullptr; - std::string parenthesis_ivf, parenthesis_refine; - Index* index = nullptr; - bool add_idmap = false; - int d_in = d; - - ScopeDeleter1 del_coarse_quantizer, del_index; - - std::string description(description_in); - char* ptr; - - // handle indexes in parentheses - while (description.find('(') != std::string::npos) { - // then we make a sub-index and remove the () from the description - int i0, i1; - find_matching_parentheses(description, i0, i1); - - std::string sub_description = description.substr(i0 + 1, i1 - i0 - 1); - - if (str_ends_with_digits(description.substr(0, i0), "IVF")) { - parenthesis_ivf = sub_description; - } else if (str_ends_with(description.substr(0, i0), "Refine")) { - parenthesis_refine = sub_description; - } else { - FAISS_THROW_MSG("don't know what to do with parenthesis index"); - } - description = description.erase(i0, i1 - i0 + 1); - } - - int64_t ncentroids = -1; - bool use_2layer = false; - int hnsw_M = -1; - int nsg_R = -1; - - for (char* tok = strtok_r(&description[0], " ,", &ptr); tok; - tok = strtok_r(nullptr, " ,", &ptr)) { - int d_out, opq_M, nbit, M, M2, pq_m, ncent, r2, R; - std::string stok(tok); - nbit = 8; - int bbs = -1; - char c; - - // to avoid mem leaks with exceptions: - // do all tests before any instanciation - - VectorTransform* vt_1 = nullptr; - Index* coarse_quantizer_1 = nullptr; - Index* index_1 = nullptr; - - // VectorTransforms - if (sscanf(tok, "PCA%d", &d_out) == 1) { - vt_1 = new PCAMatrix(d, d_out); - d = d_out; - } else if (sscanf(tok, "PCAR%d", &d_out) == 1) { - vt_1 = new PCAMatrix(d, d_out, 0, true); - d = d_out; - } else if (sscanf(tok, "RR%d", &d_out) == 1) { - vt_1 = new RandomRotationMatrix(d, d_out); - d = d_out; - } else if (sscanf(tok, "PCAW%d", &d_out) == 1) { - vt_1 = new PCAMatrix(d, d_out, -0.5, false); - d = d_out; - } else if (sscanf(tok, "PCAWR%d", &d_out) == 1) { - vt_1 = new PCAMatrix(d, d_out, -0.5, true); - d = d_out; - } else if (sscanf(tok, "OPQ%d_%d", &opq_M, &d_out) == 2) { - vt_1 = new OPQMatrix(d, opq_M, d_out); - d = d_out; - } else if (sscanf(tok, "OPQ%d", &opq_M) == 1) { - vt_1 = new OPQMatrix(d, opq_M); - } else if (sscanf(tok, "ITQ%d", &d_out) == 1) { - vt_1 = new ITQTransform(d, d_out, true); - d = d_out; - } else if (stok == "ITQ") { - vt_1 = new ITQTransform(d, d, false); - } else if (sscanf(tok, "Pad%d", &d_out) == 1) { - if (d_out > d) { - vt_1 = new RemapDimensionsTransform(d, d_out, false); - d = d_out; - } - } else if (stok == "L2norm") { - vt_1 = new NormalizationTransform(d, 2.0); - - // coarse quantizers - } else if ( - !coarse_quantizer && - sscanf(tok, "IVF%" PRId64 "_HNSW%d", &ncentroids, &M) == 2) { - coarse_quantizer_1 = new IndexHNSWFlat(d, M, metric); - - } else if ( - !coarse_quantizer && - sscanf(tok, "IVF%" PRId64 "_NSG%d", &ncentroids, &R) == 2) { - coarse_quantizer_1 = new IndexNSGFlat(d, R, metric); - - } else if ( - !coarse_quantizer && - sscanf(tok, "IVF%" PRId64, &ncentroids) == 1) { - if (!parenthesis_ivf.empty()) { - coarse_quantizer_1 = - index_factory(d, parenthesis_ivf.c_str(), metric); - } else if (metric == METRIC_L2) { - coarse_quantizer_1 = new IndexFlatL2(d); - } else { - coarse_quantizer_1 = new IndexFlatIP(d); - } - - } else if (!coarse_quantizer && sscanf(tok, "IMI2x%d", &nbit) == 1) { - FAISS_THROW_IF_NOT_MSG( - metric == METRIC_L2, - "MultiIndex not implemented for inner prod search"); - coarse_quantizer_1 = new MultiIndexQuantizer(d, 2, nbit); - ncentroids = 1 << (2 * nbit); - - } else if ( - !coarse_quantizer && - sscanf(tok, "Residual%dx%d", &M, &nbit) == 2) { - FAISS_THROW_IF_NOT_MSG( - metric == METRIC_L2, - "MultiIndex not implemented for inner prod search"); - coarse_quantizer_1 = new MultiIndexQuantizer(d, M, nbit); - ncentroids = int64_t(1) << (M * nbit); - use_2layer = true; - - } else if ( - std::regex_match( - stok, - std::regex( - "(RQ|RCQ)[0-9]+x[0-9]+(_[0-9]+x[0-9]+)*(_Nnone|_Nfloat|_Nqint8|_Nqint4|_Ncqint8|_Ncqint4)?"))) { - std::vector nbits; - std::smatch sm; - bool is_RCQ = stok.find("RCQ") == 0; - while (std::regex_search( - stok, sm, std::regex("([0-9]+)x([0-9]+)"))) { - int M = std::stoi(sm[1].str()); - int nbit = std::stoi(sm[2].str()); - nbits.resize(nbits.size() + M, nbit); - stok = sm.suffix(); - } - if (is_RCQ) { - index_1 = new ResidualCoarseQuantizer(d, nbits, metric); - } else { - AdditiveQuantizer::Search_type_t st = - parse_AQ_search_type(stok, metric); - if (!coarse_quantizer) { - index_1 = new IndexResidualQuantizer(d, nbits, metric, st); - } else { - index_1 = fix_ivf_fields(new IndexIVFResidualQuantizer( - coarse_quantizer, - d, - ncentroids, - nbits, - metric, - st)); - del_coarse_quantizer.release(); - } - } - } else if ( - std::regex_match( - stok, - std::regex( - "(LSQ|LSCQ)[0-9]+x[0-9]+(_Nnone|_Nfloat|_Nqint8|_Nqint4|_Ncqint8|_Ncqint4)?"))) { - std::vector nbits; - std::smatch sm; - bool is_LSCQ = stok.find("LSCQ") == 0; - std::regex_search(stok, sm, std::regex("([0-9]+)x([0-9]+)")); - int M = std::stoi(sm[1].str()); - int nbit = std::stoi(sm[2].str()); - if (is_LSCQ) { - index_1 = new ResidualCoarseQuantizer(d, nbits, metric); - } else { - AdditiveQuantizer::Search_type_t st = - parse_AQ_search_type(stok, metric); - if (!coarse_quantizer) { - index_1 = new IndexLocalSearchQuantizer( - d, M, nbit, metric, st); - } else { - index_1 = fix_ivf_fields(new IndexIVFLocalSearchQuantizer( - coarse_quantizer, - d, - ncentroids, - M, - nbit, - metric, - st)); - del_coarse_quantizer.release(); - } - } - } else if ( - !coarse_quantizer && - sscanf(tok, "Residual%" PRId64, &ncentroids) == 1) { - coarse_quantizer_1 = new IndexFlatL2(d); - use_2layer = true; - - } else if (stok == "IDMap") { - add_idmap = true; - - // IVFs - } else if (!index && (stok == "Flat" || stok == "FlatDedup")) { - if (coarse_quantizer) { - // if there was an IVF in front, then it is an IVFFlat - IndexIVF* index_ivf = stok == "Flat" - ? new IndexIVFFlat( - coarse_quantizer, d, ncentroids, metric) - : new IndexIVFFlatDedup( - coarse_quantizer, d, ncentroids, metric); - index_1 = fix_ivf_fields(index_ivf); - del_coarse_quantizer.release(); - } else if (hnsw_M > 0) { - index_1 = new IndexHNSWFlat(d, hnsw_M, metric); - } else if (nsg_R > 0) { - index_1 = new IndexNSGFlat(d, nsg_R, metric); - } else { - FAISS_THROW_IF_NOT_MSG( - stok != "FlatDedup", - "dedup supported only for IVFFlat"); - index_1 = new IndexFlat(d, metric); - } - } else if ( - !index && - (stok == "SQ8" || stok == "SQ4" || stok == "SQ6" || - stok == "SQfp16")) { - ScalarQuantizer::QuantizerType qt = stok == "SQ8" - ? ScalarQuantizer::QT_8bit - : stok == "SQ6" ? ScalarQuantizer::QT_6bit - : stok == "SQ4" ? ScalarQuantizer::QT_4bit - : stok == "SQfp16" ? ScalarQuantizer::QT_fp16 - : ScalarQuantizer::QT_4bit; - if (coarse_quantizer) { - FAISS_THROW_IF_NOT(!use_2layer); - IndexIVFScalarQuantizer* index_ivf = - new IndexIVFScalarQuantizer( - coarse_quantizer, d, ncentroids, qt, metric); - index_1 = fix_ivf_fields(index_ivf); - del_coarse_quantizer.release(); - } else if (hnsw_M > 0) { - index_1 = new IndexHNSWSQ(d, qt, hnsw_M, metric); - } else { - index_1 = new IndexScalarQuantizer(d, qt, metric); - } - } else if (!index && sscanf(tok, "PQ%d+%d", &M, &M2) == 2) { - FAISS_THROW_IF_NOT_MSG( - coarse_quantizer, "PQ with + works only with an IVF"); - FAISS_THROW_IF_NOT_MSG( - metric == METRIC_L2, - "IVFPQR not implemented for inner product search"); - IndexIVFPQR* index_ivf = new IndexIVFPQR( - coarse_quantizer, d, ncentroids, M, 8, M2, 8); - index_1 = fix_ivf_fields(index_ivf); - del_coarse_quantizer.release(); - } else if ( - !index && - (sscanf(tok, "PQ%dx4fs_%d", &M, &bbs) == 2 || - (sscanf(tok, "PQ%dx4f%c", &M, &c) == 2 && c == 's') || - (sscanf(tok, "PQ%dx4fs%c", &M, &c) == 2 && c == 'r'))) { - if (bbs == -1) { - bbs = 32; - } - bool by_residual = str_ends_with(stok, "fsr"); - if (coarse_quantizer) { - IndexIVFPQFastScan* index_ivf = new IndexIVFPQFastScan( - coarse_quantizer, d, ncentroids, M, 4, metric, bbs); - index_1 = fix_ivf_fields(index_ivf); - index_ivf->by_residual = by_residual; - del_coarse_quantizer.release(); - } else { - IndexPQFastScan* index_pq = - new IndexPQFastScan(d, M, 4, metric, bbs); - index_1 = index_pq; - } - } else if ( - !index && - (sscanf(tok, "PQ%dx%d", &M, &nbit) == 2 || - sscanf(tok, "PQ%d", &M) == 1 || - sscanf(tok, "PQ%dnp", &M) == 1)) { - bool do_polysemous_training = stok.find("np") == std::string::npos; - if (coarse_quantizer) { - if (!use_2layer) { - IndexIVFPQ* index_ivf = new IndexIVFPQ( - coarse_quantizer, d, ncentroids, M, nbit, metric); - index_1 = fix_ivf_fields(index_ivf); - del_coarse_quantizer.release(); - index_ivf->do_polysemous_training = do_polysemous_training; - } else { - Index2Layer* index_2l = new Index2Layer( - coarse_quantizer, ncentroids, M, nbit); - index_2l->q1.quantizer_trains_alone = - get_trains_alone(coarse_quantizer); - index_2l->q1.own_fields = true; - index_1 = index_2l; - } - } else if (hnsw_M > 0) { - IndexHNSWPQ* ipq = new IndexHNSWPQ(d, M, hnsw_M); - dynamic_cast(ipq->storage)->do_polysemous_training = - do_polysemous_training; - index_1 = ipq; - } else { - IndexPQ* index_pq = new IndexPQ(d, M, nbit, metric); - index_pq->do_polysemous_training = do_polysemous_training; - index_1 = index_pq; - } - } else if ( - !index && - sscanf(tok, "HNSW%d_%d+PQ%d", &M, &ncent, &pq_m) == 3) { - Index* quant = new IndexFlatL2(d); - IndexHNSW2Level* hidx2l = - new IndexHNSW2Level(quant, ncent, pq_m, M); - Index2Layer* idx2l = dynamic_cast(hidx2l->storage); - idx2l->q1.own_fields = true; - index_1 = hidx2l; - } else if ( - !index && - sscanf(tok, "HNSW%d_2x%d+PQ%d", &M, &nbit, &pq_m) == 3) { - Index* quant = new MultiIndexQuantizer(d, 2, nbit); - IndexHNSW2Level* hidx2l = - new IndexHNSW2Level(quant, 1 << (2 * nbit), pq_m, M); - Index2Layer* idx2l = dynamic_cast(hidx2l->storage); - idx2l->q1.own_fields = true; - idx2l->q1.quantizer_trains_alone = 1; - index_1 = hidx2l; - } else if (!index && sscanf(tok, "HNSW%d_PQ%d", &M, &pq_m) == 2) { - index_1 = new IndexHNSWPQ(d, pq_m, M); - } else if ( - !index && sscanf(tok, "HNSW%d_SQ%d", &M, &pq_m) == 2 && - pq_m == 8) { - index_1 = new IndexHNSWSQ(d, ScalarQuantizer::QT_8bit, M); - } else if (!index && sscanf(tok, "HNSW%d", &M) == 1) { - hnsw_M = M; - // here it is unclear what we want: HNSW flat or HNSWx,Y ? - } else if (!index && sscanf(tok, "NSG%d", &R) == 1) { - nsg_R = R; - } else if ( - !index && - (stok == "LSH" || stok == "LSHr" || stok == "LSHrt" || - stok == "LSHt")) { - bool rotate_data = strstr(tok, "r") != nullptr; - bool train_thresholds = strstr(tok, "t") != nullptr; - index_1 = new IndexLSH(d, d, rotate_data, train_thresholds); - } else if ( - !index && - sscanf(tok, "ZnLattice%dx%d_%d", &M, &r2, &nbit) == 3) { - FAISS_THROW_IF_NOT(!coarse_quantizer); - index_1 = new IndexLattice(d, M, nbit, r2); - } else if (stok == "RFlat") { - parenthesis_refine = "Flat"; - } else if (stok == "Refine") { - FAISS_THROW_IF_NOT_MSG( - !parenthesis_refine.empty(), - "Refine index should be provided in parentheses"); - } else { - FAISS_THROW_FMT( - "could not parse token \"%s\" in %s\n", - tok, - description_in); - } - - if (index_1 && add_idmap) { - IndexIDMap* idmap = new IndexIDMap(index_1); - del_index.set(idmap); - idmap->own_fields = true; - index_1 = idmap; - add_idmap = false; - } - - if (vt_1) { - vts.chain.push_back(vt_1); - } - - if (coarse_quantizer_1) { - coarse_quantizer = coarse_quantizer_1; - del_coarse_quantizer.set(coarse_quantizer); - } - - if (index_1) { - index = index_1; - del_index.set(index); - } - } - - if (!index && hnsw_M > 0) { - index = new IndexHNSWFlat(d, hnsw_M, metric); - del_index.set(index); - } else if (!index && nsg_R > 0) { - index = new IndexNSGFlat(d, nsg_R, metric); - del_index.set(index); - } - - FAISS_THROW_IF_NOT_FMT( - index, "description %s did not generate an index", description_in); - - // nothing can go wrong now - del_index.release(); - del_coarse_quantizer.release(); - - if (add_idmap) { - fprintf(stderr, - "index_factory: WARNING: " - "IDMap option not used\n"); - } - - if (vts.chain.size() > 0) { - IndexPreTransform* index_pt = new IndexPreTransform(index); - index_pt->own_fields = true; - // add from back - while (vts.chain.size() > 0) { - index_pt->prepend_transform(vts.chain.back()); - vts.chain.pop_back(); - } - index = index_pt; - } - - if (!parenthesis_refine.empty()) { - Index* refine_index = - index_factory(d_in, parenthesis_refine.c_str(), metric); - IndexRefine* index_rf = new IndexRefine(index, refine_index); - index_rf->own_refine_index = true; - index_rf->own_fields = true; - index = index_rf; - } - - return index; +Index* index_factory(int d, const char* description, MetricType metric) { + return index_factory_sub(d, description, metric).release(); } IndexBinary* index_binary_factory(int d, const char* description) { diff --git a/faiss/index_factory.h b/faiss/index_factory.h index 4df2b2c77..2de8ee0a3 100644 --- a/faiss/index_factory.h +++ b/faiss/index_factory.h @@ -19,6 +19,9 @@ Index* index_factory( const char* description, MetricType metric = METRIC_L2); +/// set to > 0 to get more logs from index_factory +FAISS_API extern int index_factory_verbose; + IndexBinary* index_binary_factory(int d, const char* description); } // namespace faiss diff --git a/tests/test_factory.py b/tests/test_factory.py index 8f7b918a3..c6a20f1ab 100644 --- a/tests/test_factory.py +++ b/tests/test_factory.py @@ -61,14 +61,16 @@ class TestFactory(unittest.TestCase): index = faiss.index_factory(12, "HNSW32_PQ4") assert index.storage.sa_code_size() == 4 + def test_factory_HNSW_newstyle(self): index = faiss.index_factory(12, "HNSW32,Flat") assert index.storage.sa_code_size() == 12 * 4 index = faiss.index_factory(12, "HNSW32,SQ8", faiss.METRIC_INNER_PRODUCT) assert index.storage.sa_code_size() == 12 assert index.metric_type == faiss.METRIC_INNER_PRODUCT - index = faiss.index_factory(12, "HNSW32,PQ4") + index = faiss.index_factory(12, "HNSW,PQ4") assert index.storage.sa_code_size() == 4 + self.assertEqual(index.hnsw.nb_neighbors(1), 32) index = faiss.index_factory(12, "HNSW32,PQ4np") indexpq = faiss.downcast_index(index.storage) assert not indexpq.do_polysemous_training @@ -110,6 +112,8 @@ class TestFactory(unittest.TestCase): self.assertEqual(index.bbs, 64) self.assertEqual(index.nlist, 50) self.assertTrue(index.cp.spherical) + index = faiss.index_factory(56, "IVF50,PQ28x4fsr_64") + self.assertTrue(index.by_residual) index = faiss.index_factory(56, "PQ28x4fs,RFlat") self.assertEqual(index.k_factor, 1.0) @@ -187,3 +191,70 @@ class TestVTDowncast(unittest.TestCase): itqt = faiss.downcast_VectorTransform(codec.chain.at(0)) itqt.pca_then_itq + + +# tests after re-factoring +class TestFactoryV2(unittest.TestCase): + + def test_refine(self): + index = faiss.index_factory(123, "Flat,RFlat") + index.k_factor + + def test_refine_2(self): + index = faiss.index_factory(123, "LSHrt,Refine(Flat)") + index1 = faiss.downcast_index(index.base_index) + self.assertTrue(index1.rotate_data) + self.assertTrue(index1.train_thresholds) + + def test_pre_transform(self): + index = faiss.index_factory(123, "PCAR100,L2Norm,PCAW50,LSHr") + self.assertTrue(index.chain.size() == 3) + + def test_ivf(self): + index = faiss.index_factory(123, "IVF456,Flat") + self.assertEqual(index.__class__, faiss.IndexIVFFlat) + + def test_idmap(self): + index = faiss.index_factory(123, "Flat,IDMap") + self.assertEqual(index.__class__, faiss.IndexIDMap) + + def test_ivf_hnsw(self): + index = faiss.index_factory(123, "IVF100_HNSW,Flat") + quantizer = faiss.downcast_index(index.quantizer) + self.assertEqual(quantizer.hnsw.nb_neighbors(1), 32) + + def test_ivf_parent(self): + index = faiss.index_factory(123, "IVF100(LSHr),Flat") + quantizer = faiss.downcast_index(index.quantizer) + self.assertEqual(quantizer.__class__, faiss.IndexLSH) + + +class TestAdditive(unittest.TestCase): + + def test_rcq(self): + index = faiss.index_factory(12, "IVF256(RCQ2x4),RQ3x4") + self.assertEqual( + faiss.downcast_index(index.quantizer).__class__, + faiss.ResidualCoarseQuantizer + ) + + def test_rq3(self): + index = faiss.index_factory(5, "RQ2x16_3x8_6x4") + + np.testing.assert_array_equal( + faiss.vector_to_array(index.rq.nbits), + np.array([16, 16, 8, 8, 8, 4, 4, 4, 4, 4, 4]) + ) + + def test_norm(self): + index = faiss.index_factory(5, "RQ8x8_Nqint8") + self.assertEqual( + index.rq.search_type, + faiss.AdditiveQuantizer.ST_norm_qint8) + + +class TestSpectralHash(unittest.TestCase): + + def test_sh(self): + index = faiss.index_factory(123, "IVF256,ITQ64,SH1.2") + self.assertEqual(index.__class__, faiss.IndexIVFSpectralHash)