nit

qdrant · Nov 11, 2024 · d63262d · d63262d
1 parent 8136317
commit d63262d
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/fastembed/late_interaction/colbert.py b/fastembed/late_interaction/colbert.py
@@ -75,8 +75,9 @@ def _preprocess_onnx_input(
         onnx_input["input_ids"] = np.insert(onnx_input["input_ids"], 1, marker_token, axis=1)
         onnx_input["attention_mask"] = np.insert(onnx_input["attention_mask"], 1, 1, axis=1)
 
-        onnx_input["input_ids"] = onnx_input["input_ids"][:, :original_length]
-        onnx_input["attention_mask"] = onnx_input["attention_mask"][:, :original_length]
+        if not is_doc:
+            onnx_input["input_ids"] = onnx_input["input_ids"][:, :original_length]
+            onnx_input["attention_mask"] = onnx_input["attention_mask"][:, :original_length]
         return onnx_input
 
     def tokenize(self, documents: List[str], is_doc: bool = True) -> List[Encoding]: