fix: Fix colbert model shape mismatch

qdrant · Nov 25, 2024 · 0f0c4f4 · 0f0c4f4
1 parent adfc03e
commit 0f0c4f4
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/fastembed/late_interaction/colbert.py b/fastembed/late_interaction/colbert.py
@@ -102,6 +102,9 @@ def _tokenize_query(self, query: str) -> list[Encoding]:
         return encoded
 
     def _tokenize_documents(self, documents: list[str]) -> list[Encoding]:
+        current_max_length = self.tokenizer.truncation["max_length"]
+        # ensure not to overflow after adding document-marker
+        self.tokenizer.enable_truncation(max_length=current_max_length - 1)
         encoded = self.tokenizer.encode_batch(documents)
         return encoded