import faiss 
import torch 
import numpy as np

N = 10_000
M = 16

res = faiss.StandardGpuResources()
res.setTempMemory(1024*1024*512)
co = faiss.GpuClonerOptions()
co.useFloat16 = False 

corpus_embeds = np.random.random((N, 768)).astype('float32')
corpus_embeds[:, 0] += np.arange(N) / 1000.
text_ids = np.arange(N)

faiss.omp_set_num_threads(32)
index = faiss.IndexResidualQuantizer(768, M, 8, faiss.METRIC_INNER_PRODUCT)

index.verbose = True
index = faiss.index_cpu_to_gpu(res, 0, index, co)    
index.train(corpus_embeds)
index.add(corpus_embeds)
index = faiss.index_gpu_to_cpu(index)

rq = index.rq

for i in range(0, N, 8):
    batch_ids = text_ids[i:i+8]
    batch_embeds = corpus_embeds[i:i+8]

    nn_scores, nn_ids = index.search(batch_embeds, 10)
    codes = rq.compute_codes(batch_embeds)

    for query_id, docids in zip(batch_ids, nn_ids):
        print("query_id: ", query_id, "docids: ", docids)
    print(nn_scores)
    #print(codes.shape, codes, codes.dtype)
    break

pure virtual method called                                                                                                                                                │·····
terminate called without an active exception                                                                                                                              │·····
Aborted