From d63262dca1fe3c4342795515d3d643ea5998a25e Mon Sep 17 00:00:00 2001 From: hh-space-invader Date: Mon, 11 Nov 2024 13:15:02 +0200 Subject: [PATCH] nit --- fastembed/late_interaction/colbert.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/fastembed/late_interaction/colbert.py b/fastembed/late_interaction/colbert.py index d353d19b..1f33e3c6 100644 --- 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]: