truncation to max length

michaelfeil · Feb 27, 2024 · e7905c2 · e7905c2
1 parent d85a1e0
commit e7905c2
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/libs/infinity_emb/infinity_emb/transformer/crossencoder/torch.py b/libs/infinity_emb/infinity_emb/transformer/crossencoder/torch.py
@@ -92,6 +92,6 @@ def tokenize_lengths(self, sentences: List[str]) -> List[int]:
             return_attention_mask=False,
             return_length=False,
             # max_length=self._infinity_tokenizer.model_max_length,
-            # truncation="longest_first",
+            truncation="longest_first",
         ).encodings
         return [len(t.tokens) for t in tks]
diff --git a/libs/infinity_emb/infinity_emb/transformer/embedder/fastembed.py b/libs/infinity_emb/infinity_emb/transformer/embedder/fastembed.py
@@ -72,5 +72,6 @@ def encode_post(self, embedding: np.ndarray) -> EmbeddingReturnType:
     def tokenize_lengths(self, sentences: List[str]) -> List[int]:
         tks = self._infinity_tokenizer.encode_batch(
             sentences,
+            truncation=True,
         )
         return [len(t.tokens) for t in tks]