Feature/faiss index caching

model_configs/README.md

@@ -0,0 +1,67 @@
+## 🔍 FAISS Index Reuse for RAG Evaluation
+
+This fork adds support for **reusing or saving FAISS indexes** during RAG evaluations in LightEval, dramatically improving speed and avoiding redundant recomputation.
+
+### ✨ Why This Matters
+
+By default, LightEval rebuilds the FAISS index **every run**, even if your dataset and embedding model are unchanged. With this update, you can:
+- Load a previously computed FAISS index from disk
+- Save the index the first time it's built
+- Avoid recomputing embeddings and chunk processing unnecessarily
+
+---
+
+### 🔧 How to Use
+
+To activate FAISS index reuse, modify your `rag_model.yaml` file as follows:
+
+```yaml
+rag_params:
+    embedding_model: "thenlper/gte-small" 
+    docs_name_or_path: "m-ric/huggingface_doc" 
+    similarity_fn: cosine 
+    top_k: 20
+    num_chunks: 100000
+    faiss_index_path: "./faiss_cache/my_index.index"
+```
+
+#### ✅ `faiss_index_path`
+- **If this path exists**, the index will be loaded from disk using `FAISS.load_local(...)`.
+- **If it doesn't exist**, the index will be built from `docs_name_or_path`, then saved to this path using `vector_db.save_local(...)`.
+
+> Example:
+```yaml
+faiss_index_path: "indexes/wiki_stem_faiss_index"
+```
+
+This will create `wiki_stem_faiss_index/index.faiss` and `index.pkl`.
+
+---
+
+### 🛠️ What Was Changed
+
+#### 1. `embed_model.py`
+- Added a `faiss_index_path` parameter in `EmbeddingModelConfig`.
+- Modified `_build_knowledge_base(...)` to check if the FAISS index should be loaded or built and saved.
+
+#### 2. `main_accelerate.py`
+- Extracts `faiss_index_path` from the `rag_params` config block and passes it to the model config.
+
+---
+
+### ✅ Benefits
+
+- ⏱️ **Faster**: Skips re-embedding and chunking for static datasets
+- 💾 **Cacheable**: Can be reused across multiple model evaluations
+- 🧪 **Modular**: Keeps FAISS logic inside the embedding model where it belongs
+
+---
+
+### 💡 Tips
+
+- You must ensure the embedding model and dataset used match the FAISS index — otherwise, results will be unreliable.
+- The `faiss_index_path` directory will be created if it doesn't exist.
+
+---
+
+This addition makes RAG evaluation workflows faster and more reproducible.

model_configs/rag_model.yaml

@@ -15,6 +15,7 @@ model:
     similarity_fn: cosine # Select the similarity function to use, options are: cosine, dot_product, max_inner_product, jaccard
     top_k: 20 # Choose the number of documents to retrieve
     num_chunks: 100000 # This is the limit we set for the number of chunks to retrieve from the document where each chunk is 512 tokens
+    faiss_index_path: "./faiss_cache/my_index.index"
 
   # Ignore this section, do not modify!
   merged_weights:

src/lighteval/main_accelerate.py

@@ -217,6 +217,8 @@ def accelerate(  # noqa C901
             rag_model_config.top_k = config["rag_params"]["top_k"]
             rag_model_config.similarity_fn = config["rag_params"]["similarity_fn"]
             rag_model_config.num_chunks = config["rag_params"]["num_chunks"]
+            rag_model_config.faiss_index_path = config["rag_params"].get("faiss_index_path", None)
+
     else:
         model_args_dict: dict = {k.split("=")[0]: k.split("=")[1] if "=" in k else True for k in model_args.split(",")}
         model_args_dict["accelerator"] = accelerator

src/lighteval/models/transformers/embed_model.py

@@ -68,6 +68,7 @@
 class EmbeddingModelConfig(TransformersModelConfig):
     """Configuration for the embedding model."""
 
+    faiss_index_path: Optional[str] = None
     num_chunks: int = 100000
     similarity_fn: str = "cosine"
     top_k: int = 5
@@ -117,7 +118,8 @@ def __init__(
         self.top_k = config.top_k
         self.num_chunks = config.num_chunks
         self.similarity_fn = config.similarity_fn
-        self.vector_db = self._build_knowledge_base()
+        self.vector_db = self._build_knowledge_base(config.faiss_index_path)
+
 
     @property
     def tokenizer(self):
@@ -347,101 +349,115 @@ def _split_documents(
 
         return docs_processed_unique
 
-    def _build_knowledge_base(self) -> FAISS:
-        ds = load_dataset(self.docs_name_or_path, split="train")
-        knowledge_base = [
-            LangchainDocument(
-                page_content=doc["text"],
-                metadata={"source": doc["source"]}) for doc in tqdm(ds, desc="Loading KB")
-        ]
-
-        docs_processed = self._split_documents(
-             self._max_length, knowledge_base)
-
-        texts = [d.page_content for d in docs_processed]
-        metadatas = [d.metadata for d in docs_processed]
-
-        logger.info(f"Embedding {len(texts)} documents...")
-
-        embeddings = self.model.embed_documents(texts)
-        if not embeddings:
-             raise ValueError("Embedding process yielded no vectors!")
-        dim = len(embeddings[0])
-
-        logger.info(f"Building FAISS IVF-PQ index (dim={dim})...")
-
-        N = len(texts)
-
-        if N == 0:
-            raise ValueError("Cannot build an index with 0 vectors.")
-        elif N < 8:
-            # For very small N, IVF might be overkill or unstable.
-            # A small fixed nlist or even falling back to Flat might be better.
-            # Let's use a minimum of 4 lists if N is sufficient, otherwise just 1.
-            nlist = min(4, N)
-            logger.warning(f"Very small N ({N}). Using nlist={nlist}. Consider using IndexFlatL2 instead.")
-        else:
-            sqrt_n = math.sqrt(N)
-            log2_sqrt_n = math.log2(sqrt_n)
-
-            pow2_low = 2**math.floor(log2_sqrt_n)
-            pow2_high = 2**math.ceil(log2_sqrt_n)
-
-            if sqrt_n - pow2_low < pow2_high - sqrt_n: # Check distance, prefer lower if equidistant
-                closest_pow2 = pow2_low
-            else:
-                closest_pow2 = pow2_high
-
-            nlist = 4 * closest_pow2
-            nlist = max(4, nlist)
-            nlist = min(nlist, N)
-
-        nlist = int(nlist)
-        logger.info(f"Using N={N}, calculated nlist = {nlist} (based on 2 * closest power of 2 to sqrt(N))")
-
-        pq_m = min(32, dim)
-        index_key = f"IVF{nlist},PQ{pq_m}"
-        try:
-            index = faiss.index_factory(dim, index_key)
-        except Exception as e:
-             logger.error(f"Failed to create index factory for {index_key}. Dim={dim}. Error: {e}")
-             # Fallback to FlatL2 if IVF-PQ fails (e.g., dim < pq_m or too few vectors for nlist)
-             logger.warning("Falling back to IndexFlatL2 due to index_factory error.")
-             index = faiss.IndexFlatL2(dim)
-             index_key = "IndexFlatL2"
-
-        # Check if training is required (IVF needs training, Flat does not)
-        if hasattr(index, 'is_trained') and not index.is_trained:
-            logger.info(f"Training FAISS index ({index_key})...")
-            embeddings_np = np.asarray(embeddings, dtype="float32")
-            if embeddings_np.shape[0] < nlist and index_key.startswith("IVF"):
-                 logger.warning(f"Insufficient vectors ({embeddings_np.shape[0]}) to train {nlist} clusters. Reducing nlist.")
-                 # Adjust nlist or fallback - simple fallback shown here
-                 index = faiss.IndexFlatL2(dim)
-                 index_key = "IndexFlatL2 (fallback)"
-                 logger.info(f"Switched to {index_key} due to insufficient training data.")
-            elif index_key.startswith("IVF"): # Only train if IVF and enough data
-                index.train(embeddings_np)
-
-        logger.info(f"Adding vectors to FAISS index ({index_key})...")
-        index.add(np.asarray(embeddings, dtype="float32"))
-        index.nprobe = min(8, nlist) if index_key.startswith("IVF") else 1 # recall/speed knob for IVF
-
-        ids = list(map(str, range(len(texts))))
-        docstore = InMemoryDocstore(dict(zip(ids, docs_processed)))
-        index_to_docstore_id = {i: ids[i] for i in range(len(ids))}
-
-        logger.info(f"Knowledge-base built with {index_key} (nprobe = {getattr(index, 'nprobe', 'N/A')})")
-
-        vector_db = FAISS(
-            embedding_function=self.model.embed_query,
-            index=index,
-            docstore=docstore,
-            index_to_docstore_id=index_to_docstore_id,
-            distance_strategy=distance_strategy_mapping[self.similarity_fn],
-            normalize_L2=(self.similarity_fn == "cosine")
+    def _build_knowledge_base(self, faiss_index_path: Optional[str] = None) -> FAISS:
+      if faiss_index_path and os.path.exists(faiss_index_path):
+        logger.info(f"Loading prebuilt FAISS index from {faiss_index_path}...")
+        return FAISS.load_local(
+            faiss_index_path,
+            self.model.embed_query,
+            allow_dangerous_deserialization=True
         )
-        return vector_db
+
+      ds = load_dataset(self.docs_name_or_path, split="train")
+      knowledge_base = [
+          LangchainDocument(
+              page_content=doc["text"],
+              metadata={"source": doc["source"]}) for doc in tqdm(ds, desc="Loading KB")
+      ]
+
+      docs_processed = self._split_documents(
+            self._max_length, knowledge_base)
+
+      texts = [d.page_content for d in docs_processed]
+      metadatas = [d.metadata for d in docs_processed]
+
+      logger.info(f"Embedding {len(texts)} documents...")
+
+      embeddings = self.model.embed_documents(texts)
+      if not embeddings:
+            raise ValueError("Embedding process yielded no vectors!")
+      dim = len(embeddings[0])
+
+      logger.info(f"Building FAISS IVF-PQ index (dim={dim})...")
+
+      N = len(texts)
+
+      if N == 0:
+          raise ValueError("Cannot build an index with 0 vectors.")
+      elif N < 8:
+          # For very small N, IVF might be overkill or unstable.
+          # A small fixed nlist or even falling back to Flat might be better.
+          # Let's use a minimum of 4 lists if N is sufficient, otherwise just 1.
+          nlist = min(4, N)
+          logger.warning(f"Very small N ({N}). Using nlist={nlist}. Consider using IndexFlatL2 instead.")
+      else:
+          sqrt_n = math.sqrt(N)
+          log2_sqrt_n = math.log2(sqrt_n)
+
+          pow2_low = 2**math.floor(log2_sqrt_n)
+          pow2_high = 2**math.ceil(log2_sqrt_n)
+
+          if sqrt_n - pow2_low < pow2_high - sqrt_n: # Check distance, prefer lower if equidistant
+              closest_pow2 = pow2_low
+          else:
+              closest_pow2 = pow2_high
+
+          nlist = 4 * closest_pow2
+          nlist = max(4, nlist)
+          nlist = min(nlist, N)
+
+      nlist = int(nlist)
+      logger.info(f"Using N={N}, calculated nlist = {nlist} (based on 2 * closest power of 2 to sqrt(N))")
+
+      pq_m = min(32, dim)
+      index_key = f"IVF{nlist},PQ{pq_m}"
+      try:
+          index = faiss.index_factory(dim, index_key)
+      except Exception as e:
+            logger.error(f"Failed to create index factory for {index_key}. Dim={dim}. Error: {e}")
+            # Fallback to FlatL2 if IVF-PQ fails (e.g., dim < pq_m or too few vectors for nlist)
+            logger.warning("Falling back to IndexFlatL2 due to index_factory error.")
+            index = faiss.IndexFlatL2(dim)
+            index_key = "IndexFlatL2"
+
+      # Check if training is required (IVF needs training, Flat does not)
+      if hasattr(index, 'is_trained') and not index.is_trained:
+          logger.info(f"Training FAISS index ({index_key})...")
+          embeddings_np = np.asarray(embeddings, dtype="float32")
+          if embeddings_np.shape[0] < nlist and index_key.startswith("IVF"):
+                logger.warning(f"Insufficient vectors ({embeddings_np.shape[0]}) to train {nlist} clusters. Reducing nlist.")
+                # Adjust nlist or fallback - simple fallback shown here
+                index = faiss.IndexFlatL2(dim)
+                index_key = "IndexFlatL2 (fallback)"
+                logger.info(f"Switched to {index_key} due to insufficient training data.")
+          elif index_key.startswith("IVF"): # Only train if IVF and enough data
+              index.train(embeddings_np)
+
+      logger.info(f"Adding vectors to FAISS index ({index_key})...")
+      index.add(np.asarray(embeddings, dtype="float32"))
+      index.nprobe = min(8, nlist) if index_key.startswith("IVF") else 1 # recall/speed knob for IVF
+
+      ids = list(map(str, range(len(texts))))
+      docstore = InMemoryDocstore(dict(zip(ids, docs_processed)))
+      index_to_docstore_id = {i: ids[i] for i in range(len(ids))}
+
+      logger.info(f"Knowledge-base built with {index_key} (nprobe = {getattr(index, 'nprobe', 'N/A')})")
+
+      vector_db = FAISS(
+          embedding_function=self.model.embed_query,
+          index=index,
+          docstore=docstore,
+          index_to_docstore_id=index_to_docstore_id,
+          distance_strategy=distance_strategy_mapping[self.similarity_fn],
+          normalize_L2=(self.similarity_fn == "cosine")
+      )
+
+      if faiss_index_path:
+        logger.info(f"Saving FAISS index to {faiss_index_path}...")
+        os.makedirs(os.path.dirname(faiss_index_path), exist_ok=True)
+        vector_db.save_local(faiss_index_path)
+
+      return vector_db
 
     def run_model(
         self,