def init_index(self, vector_sz: int):
# IndexHNSWFlat supports L2 similarity only
# so we have to apply DOT -> L2 similairy space conversion with the help of an extra dimension
index = faiss.IndexHNSWSQ(vector_sz + 1, faiss.ScalarQuantizer.QT_8bit, self.store_n)
index.hnsw.efSearch = self.ef_search
index.hnsw.efConstruction = self.ef_construction
self.index = index
def _post_process(self, dataset, resources_paths):
if self.config.with_index:
index_file = resources_paths["embeddings_index"]
if os.path.exists(index_file):
dataset.load_faiss_index("embeddings", index_file)
else:
if "embeddings" not in dataset.column_names:
raise ValueError("Couldn't build the index because there are no embeddings.")
import faiss
d = 768
train_size = self.config.index_train_size
logger.info("Building wiki_dpr faiss index")
if self.config.index_name == "exact":
index = faiss.IndexHNSWSQ(d, faiss.ScalarQuantizer.QT_8bit, 128, faiss.METRIC_INNER_PRODUCT)
index.hnsw.efConstruction = 200
index.hnsw.efSearch = 128
dataset.add_faiss_index("embeddings", custom_index=index, train_size=train_size)
else:
quantizer = faiss.IndexHNSWFlat(d, 128, faiss.METRIC_INNER_PRODUCT)
quantizer.hnsw.efConstruction = 200
quantizer.hnsw.efSearch = 128
ivf_index = faiss.IndexIVFPQ(quantizer, d, 4096, 128, 8, faiss.METRIC_INNER_PRODUCT)
ivf_index.nprobe = 64
ivf_index.own_fields = True
quantizer.this.disown()
dataset.add_faiss_index(
"embeddings",
train_size=train_size,
custom_index=ivf_index,
)
logger.info("Saving wiki_dpr faiss index")
dataset.save_faiss_index("embeddings", index_file)
return dataset
# to see progress
index.verbose = True
index.add(xb)
print "search"
for efSearch in 16, 32, 64, 128, 256:
print "efSearch", efSearch,
index.hnsw.efSearch = efSearch
evaluate(index)
if 'hnsw_sq' in todo:
print "Testing HNSW with a scalar quantizer"
# also set M so that the vectors and links both use 128 bytes per
# entry (total 256 bytes)
index = faiss.IndexHNSWSQ(d, faiss.ScalarQuantizer.QT_8bit, 16)
print "training"
# training for the scalar quantizer
index.train(xt)
# this is the default, higher is more accurate and slower to
# construct
index.hnsw.efConstruction = 40
print "add"
# to see progress
index.verbose = True
index.add(xb)
print "search"
def __init__(self) -> None:
self.index = faiss.IndexHNSWSQ(self.d, faiss.ScalarQuantizer.QT_8bit, 64)
self.index.hnsw.efConstruction = 80
self.index.hnsw.efSearch = 64
def build_index_streaming(
cached_embeddings_path,
output_path,
hnsw=False,
sq8_quantization=False,
fp16_quantization=False,
store_n=256,
ef_search=32,
ef_construction=80,
sample_fraction=0.1,
indexing_batch_size=5000000,
):
vector_size = get_vectors_dim(cached_embeddings_path)
if hnsw:
if sq8_quantization:
index = faiss.IndexHNSWSQ(vector_size + 1,
faiss.ScalarQuantizer.QT_8bit, store_n)
elif fp16_quantization:
index = faiss.IndexHNSWSQ(vector_size + 1,
faiss.ScalarQuantizer.QT_fp16, store_n)
else:
index = faiss.IndexHNSWFlat(vector_size + 1, store_n)
index.hnsw.efSearch = ef_search
index.hnsw.efConstruction = ef_construction
else:
if sq8_quantization:
index = faiss.IndexScalarQuantizer(vector_size,
faiss.ScalarQuantizer.QT_8bit,
faiss.METRIC_L2)
elif fp16_quantization:
index = faiss.IndexScalarQuantizer(vector_size,
faiss.ScalarQuantizer.QT_fp16,
faiss.METRIC_L2)
else:
index = faiss.IndexIP(vector_size + 1, store_n)
vector_sample, max_phi, N = get_vector_sample(cached_embeddings_path,
sample_fraction)
if hnsw:
vector_sample = augment_vectors(vector_sample, max_phi)
if sq8_quantization or fp16_quantization: # index requires training
vs = vector_sample.numpy()
logging.info(f'Training Quantizer with matrix of shape {vs.shape}')
index.train(vs)
del vs
del vector_sample
chunks_to_add = []
added = 0
for vector_chunk in parse_vectors_from_directory(cached_embeddings_path,
as_chunks=True):
if hnsw:
vector_chunk = augment_vectors(vector_chunk, max_phi)
chunks_to_add.append(vector_chunk)
if sum(c.shape[0] for c in chunks_to_add) > indexing_batch_size:
logging.info(
f'Adding Vectors {added} -> {added + to_add.shape[0]} of {N}')
to_add = torch.cat(chunks_to_add)
chunks_to_add = []
index.add(to_add)
added += 1
if len(chunks_to_add) > 0:
to_add = torch.cat(chunks_to_add).numpy()
index.add(to_add)
logging.info(
f'Adding Vectors {added} -> {added + to_add.shape[0]} of {N}')
logger.info(f'Index Built, writing index to {output_path}')
faiss.write_index(index, output_path)
logger.info(f'Index dumped')
return index
0, and grouping triplets of 3,5,8 and 4,7,9 which can
blend into one another in some cases.
"""
import umap
from sklearn.datasets import fetch_openml
import matplotlib.pyplot as plt
import seaborn as sns
import os.path
import faiss
import numpy as np
mnist = fetch_openml("mnist_784", version=1)
sns.set(context="paper", style="white")
index = faiss.IndexHNSWSQ(784, faiss.ScalarQuantizer.QT_8bit, 32)
index.verbose = True
faiss_index_file = 'faiss.index'
if os.path.exists(faiss_index_file):
print('load existing index from %s' % faiss_index_file)
index = faiss.read_index(faiss_index_file, faiss.IO_FLAG_MMAP)
index.hnsw.efSearch = 256
else:
# build lossy faiss index
print('build new index and save to %s' % faiss_index_file)
index.hnsw.efConstruction = 40
data = np.ascontiguousarray(mnist.data, dtype=np.float32)
# we no longer need mnist data in its original form
print('train index...')
index.train(data)
print('add vectors to index...')