fix: Fix colbert model shape mismatch (#413)

* fix: Fix colbert model shape mismatch * refactor: Added the truncation after tokenizer init
qdrant · Nov 28, 2024 · da30f93 · da30f93
1 parent adfc03e
commit da30f93
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/fastembed/late_interaction/colbert.py b/fastembed/late_interaction/colbert.py
@@ -191,6 +191,9 @@ def load_onnx_model(self) -> None:
             self.tokenizer.encode(symbol, add_special_tokens=False).ids[0]
             for symbol in string.punctuation
         }
+        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)
 
     def embed(
         self,