MOD: index algo

exp_pipeline/pipeline.py

@@ -42,8 +42,11 @@ def run_pipeline(split: str = "train"):
         raise ValueError("Embeddings is empty or not a 2D array. Check input texts and embedding model.")
 
     # 4. 建立FAISS索引
-    index = build_faiss_index(embeddings, texts)
+    index = build_faiss_index(embeddings, texts, index_type="HNSW")
     logger.info("FAISS index built successfully")
+    # 持久化index到./index文件夹
+    index.save("../index/msmarco_hnsw")
+    logger.info("FAISS index saved to ./index/msmarco_hnsw")
     return index
 
 if __name__ == "__main__":

retriever/faiss_index.py

@@ -1,5 +1,5 @@
 # 工厂函数，供pipeline调用
-def build_faiss_index(embeddings, texts, metadata=None, index_type="IVF"):
+def build_faiss_index(embeddings, texts, metadata=None, index_type="HNSW"):
     if embeddings is None or not hasattr(embeddings, 'shape') or len(embeddings.shape) != 2 or embeddings.shape[0] == 0:
         raise ValueError(f"Embeddings is empty or not a 2D array. Got shape: {getattr(embeddings, 'shape', None)}")
     dimension = embeddings.shape[1]
@@ -16,7 +16,7 @@ import logging
 logger = logging.getLogger(__name__)
 
 class FAISSIndex:
-    def __init__(self, dimension: int, index_type: str = "IVF"):
+    def __init__(self, dimension: int, index_type: str = "HNSW"):
         self.dimension = dimension
         self.index_type = index_type
         self.index = None
@@ -29,29 +29,28 @@ class FAISSIndex:
         """Build FAISS index from embeddings"""
         if embeddings.shape[1] != self.dimension:
             raise ValueError(f"Embedding dimension {embeddings.shape[1]} != {self.dimension}")
-        
         # Normalize embeddings for cosine similarity
         faiss.normalize_L2(embeddings)
-        
-        if self.index_type == "IVF":
-            # IVF index for large datasets
+        if self.index_type == "HNSW":
+            # HNSW index for fast approximate search
+            self.index = faiss.IndexHNSWFlat(self.dimension, 32)  # 32 is default M
+            self.index.hnsw.efConstruction = 200
+            self.index.add(embeddings)
+        elif self.index_type == "IVF":
             nlist = min(4096, len(embeddings) // 100)
             quantizer = faiss.IndexFlatIP(self.dimension)
             self.index = faiss.IndexIVFFlat(quantizer, self.dimension, nlist)
             self.index.train(embeddings)
             self.index.add(embeddings)
         else:
-            # Flat index for small datasets
             self.index = faiss.IndexFlatIP(self.dimension)
             self.index.add(embeddings)
-        
         # Store text and metadata
         for i, text in enumerate(texts):
             self.id_to_text[i] = text
             if metadata and i < len(metadata):
                 self.id_to_metadata[i] = metadata[i]
-        
-        logger.info(f"Built FAISS index with {len(embeddings)} vectors")
+        logger.info(f"Built FAISS {self.index_type} index with {len(embeddings)} vectors")
     
     def search(self, query_embeddings: np.ndarray, k: int = 10) -> Tuple[np.ndarray, np.ndarray]:
         """Search for similar vectors"""