update examples

LibVQ/base_index/BaseIndex.py

@@ -76,8 +76,8 @@ def hard_negative(self,
         score, search_results = self.search(query_embeddings, topk=topk, batch_size=batch_size, nprobe=nprobe)
         query2hardneg = {}
         for qid, neighbors in enumerate(search_results):
-            neg = list(filter(lambda x: x not in ground_truths[qid], neighbors))
-            query2hardneg[qid] = [x for x in neg if x != -1]
+            neg = list(filter(lambda x: x not in ground_truths[qid] and x != -1, neighbors))
+            query2hardneg[qid] = neg
         return query2hardneg
 
     def generate_virtual_traindata(self,

examples/MSMARCO/README.md

@@ -32,9 +32,9 @@ pleaser refer to [dataset.README.md](../../LibVQ/dataset/README.md)
 python ./basic_index/faiss_index.py  \
 --preprocess_dir ./data/passage/preprocess \
 --embeddings_dir ./data/passage/evaluate/co-condenser \
---index_method ivf_pq \
+--index_method ivf_opq \
 --ivf_centers_num 10000 \
---subvector_num 32 \
+--subvector_num 64 \
 --subvector_bits 8 \
 --nprobe 100
 ```

examples/MSMARCO/basic_index/faiss_index.py

@@ -10,7 +10,7 @@
 from LibVQ.base_index import FaissIndex
 from LibVQ.dataset.dataset import load_rel
 
-faiss.omp_set_num_threads(32)
+
 
 if __name__ == '__main__':
     parser = HfArgumentParser((IndexArguments, DataArguments, ModelArguments, TrainingArguments))

examples/MSMARCO/learnable_index/train_index.py

@@ -2,6 +2,7 @@
 sys.path.append('./')
 import os
 import pickle
+import gc
 
 import faiss
 import numpy as np
@@ -62,6 +63,8 @@
                                                            batch_size=64,
                                                            nprobe=min(index_args.ivf_centers_num, 500))
         pickle.dump(trainquery2hardneg, open(neg_file, 'wb'))
+        del trainquery2hardneg
+        gc.collect()
 
     data_args.save_ckpt_dir = f'./saved_ckpts/{training_args.training_mode}_{index_args.index_method}/'
 
@@ -131,6 +134,8 @@
                 pickle.dump(query2neg,
                             open(os.path.join(data_args.embeddings_dir, f"train-queries-virtual_hardneg.pickle"), 'wb'))
 
+                del query2neg, query2pos
+                gc.collect()
 
     # distill with no label data
     if training_args.training_mode == 'distill_index_nolabel':

examples/MSMARCO/learnable_index/train_index_and_encoder.py

@@ -2,6 +2,7 @@
 sys.path.append('./')
 import os
 import pickle
+import gc
 
 import faiss
 import numpy as np
@@ -73,6 +74,9 @@
                                                            nprobe=min(index_args.ivf_centers_num, 500))
         pickle.dump(trainquery2hardneg, open(neg_file, 'wb'))
 
+        del trainquery2hardneg
+        gc.collect()
+
     data_args.save_ckpt_dir = f'./saved_ckpts/{training_args.training_mode}_{index_args.index_method}/'
 
     # contrastive learning
@@ -192,6 +196,9 @@
             pickle.dump(query2neg,
                         open(os.path.join(data_args.embeddings_dir, f"train-queries-virtual_hardneg.pickle"), 'wb'))
 
+            del query2pos, query2neg
+            gc.collect()
+
 
     # distill with no label data
     if training_args.training_mode == 'distill_index-and-query-encoder_nolabel':

examples/NQ/learnable_index/train_index.py

@@ -2,6 +2,7 @@
 sys.path.append('./')
 import os
 import pickle
+import gc
 
 import faiss
 import numpy as np
@@ -64,6 +65,9 @@
                                                            nprobe=index_args.ivf_centers_num)
         pickle.dump(trainquery2hardneg, open(neg_file, 'wb'))
 
+        del trainquery2hardneg
+        gc.collect()
+
     data_args.save_ckpt_dir = f'./saved_ckpts/{training_args.training_mode}_{index_args.index_method}/'
 
     # contrastive learning
@@ -131,6 +135,9 @@
             pickle.dump(query2neg,
                         open(os.path.join(data_args.embeddings_dir, f"train-queries-virtual_hardneg.pickle"), 'wb'))
 
+            del query2neg, query2pos
+            gc.collect()
+
 
     # distill with no label data
     if training_args.training_mode == 'distill_index_nolabel':

examples/NQ/learnable_index/train_index_and_encoder.py

@@ -2,6 +2,7 @@
 sys.path.append('./')
 import os
 import pickle
+import gc
 
 import faiss
 import numpy as np
@@ -74,6 +75,8 @@
                                                            batch_size=64,
                                                            nprobe=index_args.ivf_centers_num)
         pickle.dump(trainquery2hardneg, open(neg_file, 'wb'))
+        del trainquery2hardneg
+        gc.collect()
 
     data_args.save_ckpt_dir = f'./saved_ckpts/{training_args.training_mode}_{index_args.index_method}/'
 
@@ -173,6 +176,9 @@
             pickle.dump(query2neg,
                         open(os.path.join(data_args.embeddings_dir, f"train-queries-virtual_hardneg.pickle"), 'wb'))
 
+            del query2pos, query2neg
+            gc.collect()
+
 
     # distill with no label data
     if training_args.training_mode == 'distill_index-and-query-encoder_nolabel':