Init ColBERTv2 managed index (#9656)

* Init ColBERTv2 managed index * cr * cr
run-llama · Dec 21, 2023 · b1adb30 · b1adb30
1 parent fecd172
commit b1adb30
Show file tree

Hide file tree

Showing 6 changed files with 467 additions and 108 deletions.
diff --git a/docs/examples/managed/manage_retrieval_benchmark.ipynb b/docs/examples/managed/manage_retrieval_benchmark.ipynb
diff --git a/examples/paul_graham_essay/ColbertIndex.ipynb b/examples/paul_graham_essay/ColbertIndex.ipynb
diff --git a/llama_index/indices/__init__.py b/llama_index/indices/__init__.py
@@ -31,6 +31,7 @@
     load_index_from_storage,
     load_indices_from_storage,
 )
+from llama_index.indices.managed.colbert_index import ColbertIndex
 from llama_index.indices.managed.vectara import VectaraIndex
 from llama_index.indices.managed.zilliz import ZillizCloudPipelineIndex
 from llama_index.indices.multi_modal import MultiModalVectorStoreIndex
@@ -52,6 +53,7 @@
     "SummaryIndex",
     "TreeIndex",
     "VectaraIndex",
+    "ColbertIndex",
     "ZillizCloudPipelineIndex",
     "DocumentSummaryIndex",
     "KnowledgeGraphIndex",

diff --git a/experimental/colbert_index/__init__.py → ...indices/managed/colbert_index/__init__.py b/experimental/colbert_index/__init__.py → ...indices/managed/colbert_index/__init__.py
diff --git a/experimental/colbert_index/base.py → ...dex/indices/managed/colbert_index/base.py b/experimental/colbert_index/base.py → ...dex/indices/managed/colbert_index/base.py
@@ -47,6 +47,7 @@ def __init__(
         service_context: Optional[ServiceContext] = None,
         storage_context: Optional[StorageContext] = None,
         model_name: str = "colbert-ir/colbertv2.0",
+        index_name: str = "",
         show_progress: bool = False,
         nbits: int = 2,
         gpus: int = 0,
@@ -58,6 +59,7 @@ def __init__(
     ) -> None:
         self.model_name = model_name
         self.index_path = "storage/colbert_index"
+        self.index_name = index_name
         self.nbits = nbits
         self.gpus = gpus
         self.ranks = ranks
@@ -75,6 +77,7 @@ def __init__(
         super().__init__(
             nodes=nodes,
             index_struct=index_struct,
+            index_name=index_name,
             service_context=service_context,
             storage_context=storage_context,
             show_progress=show_progress,
@@ -100,7 +103,7 @@ def _build_index_from_nodes(self, nodes: Sequence[BaseNode]) -> IndexDict:
         """Generate a PLAID index from the ColBERT checkpoint via its hugging face
         model_name.
         """
-        from colbert import Indexer, IndexUpdater, Searcher
+        from colbert import Indexer, Searcher
         from colbert.infra import ColBERTConfig, Run, RunConfig
 
         index_struct = IndexDict()
@@ -121,12 +124,9 @@ def _build_index_from_nodes(self, nodes: Sequence[BaseNode]) -> IndexDict:
                 kmeans_niters=self.kmeans_niters,
             )
             indexer = Indexer(checkpoint=self.model_name, config=config)
-            indexer.index("", collection=docs_list, overwrite=True)
+            indexer.index(name=self.index_name, collection=docs_list, overwrite=True)
             self.store = Searcher(
-                index="", collection=docs_list, checkpoint=self.model_name
-            )
-            self.updater = IndexUpdater(
-                config=config, searcher=self.store, checkpoint=self.model_name
+                index=self.index_name, collection=docs_list, checkpoint=self.model_name
             )
         return index_struct
 

diff --git a/experimental/colbert_index/retriever.py → ...ndices/managed/colbert_index/retriever.py b/experimental/colbert_index/retriever.py → ...ndices/managed/colbert_index/retriever.py