update-colbert-ranker

raphaelsty · Jun 6, 2024 · e0e75f9 · e0e75f9
1 parent 955af03
commit e0e75f9
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 7 deletions.
diff --git a/neural_cherche/__version__.py b/neural_cherche/__version__.py
@@ -1,3 +1,3 @@
-VERSION = (1, 4, 2)
+VERSION = (1, 4, 3)
 
 __version__ = ".".join(map(str, VERSION))
diff --git a/neural_cherche/rank/colbert.py b/neural_cherche/rank/colbert.py
@@ -138,11 +138,12 @@ def encode_documents(
         if not documents:
             return {}
 
+        documents_content = [
+            " ".join([document[field] for field in self.on]) for document in documents
+        ]
+
         embeddings = self.encode_queries(
-            queries=[
-                " ".join([document[field] for field in self.on])
-                for document in documents
-            ],
+            queries=documents_content,
             batch_size=batch_size,
             tqdm_bar=tqdm_bar,
             query_mode=query_mode,
@@ -152,8 +153,8 @@ def encode_documents(
         )
 
         return {
-            document[self.key]: embedding
-            for document, embedding in zip(documents, embeddings.values())
+            document[self.key]: embeddings[content]
+            for document, content in zip(documents, documents_content)
         }
 
     def encode_candidates_documents(