Allow custom FAISS index types

Author: AjayP13

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "DataDreamer"
-version = "0.41.0"
+version = "0.42.0"
 description = "Prompt. Generate Synthetic Data. Train & Align Models."
 license = "MIT"
 authors= [

src/retrievers/embedding_retriever.py

@@ -7,9 +7,10 @@
 
 import numpy as np
 import torch
-from datasets.fingerprint import Hasher
 from sqlitedict import SqliteDict
 
+from datasets.fingerprint import Hasher
+
 from ..datasets import OutputDatasetColumn, OutputIterableDatasetColumn
 from ..embedders.embedder import Embedder
 from ..utils.background_utils import dill_serializer, proxy_resource_in_background
@@ -28,6 +29,7 @@ def __init__(
         texts: OutputDatasetColumn | OutputIterableDatasetColumn,
         embedder: Embedder,
         truncate: bool = False,
+        index_type: str | None = None,
         index_batch_size: int = DEFAULT_BATCH_SIZE,
         index_instruction: None | str = None,
         query_instruction: None | str = None,
@@ -52,6 +54,7 @@ def __init__(
         super().__init__(texts=texts, cache_folder_path=cache_folder_path)
         self.embedder = embedder
         self.truncate = truncate
+        self.index_type = index_type
         self.index_batch_size = index_batch_size
         self.index_instruction = index_instruction
         self.query_instruction = query_instruction
@@ -99,7 +102,12 @@ def __init__(self_resource) -> None:
                     ):
                         self._initialize_retriever_index_folder()
                         index_logger.info("Building index.")
-                        index = faiss.IndexFlatIP(self.embedder.dims)
+                        if self.index_type is not None:  # pragma: no cover
+                            index = faiss.index_factory(
+                                self.embedder.dims, self.index_type
+                            )
+                        else:
+                            index = faiss.IndexFlatIP(self.embedder.dims)
                         if self.device is not None:  # pragma: no cover
                             index = faiss.index_cpu_to_gpus_list(
                                 index=index, gpus=self.device
@@ -137,6 +145,8 @@ def __init__(self_resource) -> None:
                                     ),
                                 )
                             )
+                            if not index.is_trained:  # pragma: no cover
+                                index.train(texts_embedded)
                             index.add(texts_embedded)
                             for id_, text in zip(ids_batch, texts_batch):
                                 index_lookup[id_] = text
@@ -212,7 +222,7 @@ def _batch_lookup(self, indices: list[int]) -> dict[int, dict[str, Any]]:
                 break
             lookup_query = (
                 "SELECT key, value FROM lookup"
-                f' WHERE key IN ({",".join(["?"] * len(indices_batch))})'
+                f" WHERE key IN ({','.join(['?'] * len(indices_batch))})"
             )
             results.update(
                 {