fix: Fix colbert model shape mismatch (#413)

hh-space-invader · web-flow · commit da30f934d8c5 · 2024-11-28T13:59:52.000+02:00
* fix: Fix colbert model shape mismatch

* refactor: Added the truncation after tokenizer init
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,

Original file line number	Diff line number	Diff line change
`@@ -191,6 +191,9 @@ def load_onnx_model(self) -> None:`
`191`	`191`	`self.tokenizer.encode(symbol, add_special_tokens=False).ids[0]`
`192`	`192`	`for symbol in string.punctuation`
`193`	`193`	`}`
	`194`	`+ current_max_length = self.tokenizer.truncation["max_length"]`
	`195`	`+ # ensure not to overflow after adding document-marker`
	`196`	`+ self.tokenizer.enable_truncation(max_length=current_max_length - 1)`
`194`	`197`
`195`	`198`	`def embed(`
`196`	`199`	`self,`