From 44623a69346e8fc48ef56b7264b47ea2e2eecb1e Mon Sep 17 00:00:00 2001 From: Alex Date: Sun, 8 Jun 2025 13:55:47 +0200 Subject: [PATCH 1/2] Add support for FAISS index reuse via faiss_index_path --- model_configs/rag_model.yaml | 1 + src/lighteval/main_accelerate.py | 2 + .../models/transformers/embed_model.py | 206 ++++++++++-------- 3 files changed, 114 insertions(+), 95 deletions(-) diff --git a/model_configs/rag_model.yaml b/model_configs/rag_model.yaml index 1e74f6a4c..070572be9 100644 --- a/model_configs/rag_model.yaml +++ b/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: diff --git a/src/lighteval/main_accelerate.py b/src/lighteval/main_accelerate.py index 5611f696d..07043ef5c 100644 --- a/src/lighteval/main_accelerate.py +++ b/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 diff --git a/src/lighteval/models/transformers/embed_model.py b/src/lighteval/models/transformers/embed_model.py index 41c4352be..ac6099245 100644 --- a/src/lighteval/models/transformers/embed_model.py +++ b/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, From 3b1713a6bf8242e39bd2895043b861f3a3a56247 Mon Sep 17 00:00:00 2001 From: Alex Date: Sun, 8 Jun 2025 14:01:51 +0200 Subject: [PATCH 2/2] Added ReadMe for the lighteval upgrade --- model_configs/README.md | 67 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 67 insertions(+) create mode 100644 model_configs/README.md diff --git a/model_configs/README.md b/model_configs/README.md new file mode 100644 index 000000000..3c72a5a44 --- /dev/null +++ b/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.