fastembed-gpu 0.7.2__tar.gz → 0.7.3__tar.gz

{fastembed_gpu-0.7.2 → fastembed_gpu-0.7.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: fastembed-gpu
-Version: 0.7.2
+Version: 0.7.3
 Summary: Fast, light, accurate library built for retrieval embedding generation
 License: Apache License
 Keywords: vector,embedding,neural,search,qdrant,sentence-transformers

{fastembed_gpu-0.7.2 → fastembed_gpu-0.7.3}/fastembed/late_interaction/colbert.py

@@ -2,8 +2,9 @@ import string
 from typing import Any, Iterable, Optional, Sequence, Type, Union
 
 import numpy as np
-from tokenizers import Encoding
+from tokenizers import Encoding, Tokenizer
 
+from fastembed.common.preprocessor_utils import load_tokenizer
 from fastembed.common.types import NumpyArray
 from fastembed.common import OnnxProvider
 from fastembed.common.onnx_model import OnnxOutputContext
@@ -48,24 +49,24 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
         if not is_doc:
             for embedding in output.model_output:
                 yield embedding
+        else:
+            if output.input_ids is None or output.attention_mask is None:
+                raise ValueError(
+                    "input_ids and attention_mask must be provided for document post-processing"
+                )
 
-        if output.input_ids is None or output.attention_mask is None:
-            raise ValueError(
-                "input_ids and attention_mask must be provided for document post-processing"
-            )
-
-        for i, token_sequence in enumerate(output.input_ids):
-            for j, token_id in enumerate(token_sequence):  # type: ignore
-                if token_id in self.skip_list or token_id == self.pad_token_id:
-                    output.attention_mask[i, j] = 0
+            for i, token_sequence in enumerate(output.input_ids):
+                for j, token_id in enumerate(token_sequence):  # type: ignore
+                    if token_id in self.skip_list or token_id == self.pad_token_id:
+                        output.attention_mask[i, j] = 0
 
-        output.model_output *= np.expand_dims(output.attention_mask, 2)
-        norm = np.linalg.norm(output.model_output, ord=2, axis=2, keepdims=True)
-        norm_clamped = np.maximum(norm, 1e-12)
-        output.model_output /= norm_clamped
+            output.model_output *= np.expand_dims(output.attention_mask, 2)
+            norm = np.linalg.norm(output.model_output, ord=2, axis=2, keepdims=True)
+            norm_clamped = np.maximum(norm, 1e-12)
+            output.model_output /= norm_clamped
 
-        for embedding, attention_mask in zip(output.model_output, output.attention_mask):
-            yield embedding[attention_mask == 1]
+            for embedding, attention_mask in zip(output.model_output, output.attention_mask):
+                yield embedding[attention_mask == 1]
 
     def _preprocess_onnx_input(
         self, onnx_input: dict[str, NumpyArray], is_doc: bool = True, **kwargs: Any
@@ -87,23 +88,8 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
         )
 
     def _tokenize_query(self, query: str) -> list[Encoding]:
-        assert self.tokenizer is not None
-        encoded = self.tokenizer.encode_batch([query])
-        # colbert authors recommend to pad queries with [MASK] tokens for query augmentation to improve performance
-        if len(encoded[0].ids) < self.MIN_QUERY_LENGTH:
-            prev_padding = None
-            if self.tokenizer.padding:
-                prev_padding = self.tokenizer.padding
-            self.tokenizer.enable_padding(
-                pad_token=self.MASK_TOKEN,
-                pad_id=self.mask_token_id,
-                length=self.MIN_QUERY_LENGTH,
-            )
-            encoded = self.tokenizer.encode_batch([query])
-            if prev_padding is None:
-                self.tokenizer.no_padding()
-            else:
-                self.tokenizer.enable_padding(**prev_padding)
+        assert self.query_tokenizer is not None
+        encoded = self.query_tokenizer.encode_batch([query])
         return encoded
 
     def _tokenize_documents(self, documents: list[str]) -> list[Encoding]:
@@ -183,6 +169,8 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
         self.pad_token_id: Optional[int] = None
         self.skip_list: set[int] = set()
 
+        self.query_tokenizer: Optional[Tokenizer] = None
+
         if not self.lazy_load:
             self.load_onnx_model()
 
@@ -195,6 +183,8 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
             cuda=self.cuda,
             device_id=self.device_id,
         )
+        self.query_tokenizer, _ = load_tokenizer(model_dir=self._model_dir)
+
         assert self.tokenizer is not None
         self.mask_token_id = self.special_token_to_id[self.MASK_TOKEN]
         self.pad_token_id = self.tokenizer.padding["pad_id"]
@@ -205,6 +195,12 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
         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)
+        self.query_tokenizer.enable_truncation(max_length=current_max_length - 1)
+        self.query_tokenizer.enable_padding(
+            pad_token=self.MASK_TOKEN,
+            pad_id=self.mask_token_id,
+            length=self.MIN_QUERY_LENGTH,
+        )
 
     def embed(
         self,

{fastembed_gpu-0.7.2 → fastembed_gpu-0.7.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "fastembed-gpu"
-version = "0.7.2"
+version = "0.7.3"
 description = "Fast, light, accurate library built for retrieval embedding generation"
 authors = ["Qdrant Team <info@qdrant.tech>", "NirantK <nirant.bits@gmail.com>"]
 license = "Apache License"