fastembed-gpu 0.4.0__tar.gz → 0.4.2__tar.gz

fastembed_gpu-0.4.2/NOTICE

@@ -0,0 +1,14 @@
+Copyright 2024 Qdrant
+
+This product includes software developed by Qdrant
+
+This distribution includes the following Jina AI models, each with its respective license:
+- jinaai/jina-colbert-v2
+    - License: cc-by-nc-4.0
+- jinaai/jina-reranker-v2-base-multilingual
+    - License: cc-by-nc-4.0
+
+These models are developed by Jina (https://jina.ai/) and are subject to Jina AI's licensing terms.
+
+Additional Notes:
+This project also includes third-party libraries with their respective licenses. Please refer to the documentation of each library for details regarding its usage and licensing terms.

{fastembed_gpu-0.4.0 → fastembed_gpu-0.4.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: fastembed-gpu
-Version: 0.4.0
+Version: 0.4.2
 Summary: Fast, light, accurate library built for retrieval embedding generation
 Home-page: https://github.com/qdrant/fastembed
 License: Apache License
@@ -15,16 +15,16 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
-Requires-Dist: PyStemmer (>=2.2.0,<3.0.0)
 Requires-Dist: huggingface-hub (>=0.20,<1.0)
 Requires-Dist: loguru (>=0.7.2,<0.8.0)
 Requires-Dist: mmh3 (>=4.1.0,<5.0.0)
-Requires-Dist: numpy (>=1.21,<2) ; python_version < "3.12"
-Requires-Dist: numpy (>=1.26,<2) ; python_version >= "3.12"
-Requires-Dist: onnxruntime-gpu (>=1.17.0,<2.0.0)
+Requires-Dist: numpy (>=1.21) ; python_version < "3.12"
+Requires-Dist: numpy (>=1.26) ; python_version >= "3.12"
+Requires-Dist: onnx (>=1.15.0,<2.0.0)
+Requires-Dist: onnxruntime-gpu (>=1.17.0,<1.20.0)
 Requires-Dist: pillow (>=10.3.0,<11.0.0)
+Requires-Dist: py-rust-stemmers (>=0.1.0,<0.2.0)
 Requires-Dist: requests (>=2.31,<3.0)
-Requires-Dist: snowballstemmer (>=2.2.0,<3.0.0)
 Requires-Dist: tokenizers (>=0.15,<1.0)
 Requires-Dist: tqdm (>=4.66,<5.0)
 Project-URL: Repository, https://github.com/qdrant/fastembed
@@ -169,6 +169,21 @@ embeddings = list(model.embed(images))
 # ]
 ```
 
+### 🔄 Rerankers
+```python
+from typing import List
+from fastembed.rerank.cross_encoder import TextCrossEncoder
+
+query = "Who is maintaining Qdrant?"
+documents: List[str] = [
+    "This is built to be faster and lighter than other embedding libraries e.g. Transformers, Sentence-Transformers, etc.",
+    "fastembed is supported by and maintained by Qdrant.",
+]
+encoder = TextCrossEncoder(model_name="Xenova/ms-marco-MiniLM-L-6-v2")
+scores = list(encoder.rerank(query, documents))
+
+# [-11.48061752319336, 5.472434997558594]
+```
 
 ## ⚡️ FastEmbed on a GPU
 

{fastembed_gpu-0.4.0 → fastembed_gpu-0.4.2}/README.md

@@ -137,6 +137,21 @@ embeddings = list(model.embed(images))
 # ]
 ```
 
+### 🔄 Rerankers
+```python
+from typing import List
+from fastembed.rerank.cross_encoder import TextCrossEncoder
+
+query = "Who is maintaining Qdrant?"
+documents: List[str] = [
+    "This is built to be faster and lighter than other embedding libraries e.g. Transformers, Sentence-Transformers, etc.",
+    "fastembed is supported by and maintained by Qdrant.",
+]
+encoder = TextCrossEncoder(model_name="Xenova/ms-marco-MiniLM-L-6-v2")
+scores = list(encoder.rerank(query, documents))
+
+# [-11.48061752319336, 5.472434997558594]
+```
 
 ## ⚡️ FastEmbed on a GPU
 

{fastembed_gpu-0.4.0 → fastembed_gpu-0.4.2}/fastembed/common/onnx_model.py

@@ -1,17 +1,7 @@
 import warnings
 from dataclasses import dataclass
 from pathlib import Path
-from typing import (
-    Any,
-    Dict,
-    Generic,
-    Iterable,
-    Optional,
-    Sequence,
-    Tuple,
-    Type,
-    TypeVar,
-)
+from typing import Any, Dict, Generic, Iterable, Optional, Sequence, Tuple, Type, TypeVar
 
 import numpy as np
 import onnxruntime as ort

{fastembed_gpu-0.4.0 → fastembed_gpu-0.4.2}/fastembed/late_interaction/colbert.py

@@ -12,6 +12,7 @@ from fastembed.late_interaction.late_interaction_embedding_base import (
 )
 from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker
 
+
 supported_colbert_models = [
     {
         "model": "colbert-ir/colbertv2.0",
@@ -41,7 +42,7 @@ supported_colbert_models = [
 class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
     QUERY_MARKER_TOKEN_ID = 1
     DOCUMENT_MARKER_TOKEN_ID = 2
-    MIN_QUERY_LENGTH = 32
+    MIN_QUERY_LENGTH = 31  # it's 32, we add one additional special token in the beginning
     MASK_TOKEN = "[MASK]"
 
     def _post_process_onnx_output(
@@ -67,15 +68,14 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
         return output.model_output.astype(np.float32)
 
     def _preprocess_onnx_input(
-        self, onnx_input: Dict[str, np.ndarray], is_doc: bool = True
+        self, onnx_input: Dict[str, np.ndarray], is_doc: bool = True, **kwargs: Any
     ) -> Dict[str, np.ndarray]:
-        if is_doc:
-            onnx_input["input_ids"][:, 1] = self.DOCUMENT_MARKER_TOKEN_ID
-        else:
-            onnx_input["input_ids"][:, 1] = self.QUERY_MARKER_TOKEN_ID
+        marker_token = self.DOCUMENT_MARKER_TOKEN_ID if is_doc else self.QUERY_MARKER_TOKEN_ID
+        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)
         return onnx_input
 
-    def tokenize(self, documents: List[str], is_doc: bool = True) -> List[Encoding]:
+    def tokenize(self, documents: List[str], is_doc: bool = True, **kwargs: Any) -> List[Encoding]:
         return (
             self._tokenize_documents(documents=documents)
             if is_doc
@@ -83,9 +83,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
         )
 
     def _tokenize_query(self, query: str) -> List[Encoding]:
-        # ". " is added to a query to be replaced with a special query token
-        query = [f". {query}"]
-        encoded = self.tokenizer.encode_batch(query)
+        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
@@ -96,7 +94,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
                 pad_id=self.mask_token_id,
                 length=self.MIN_QUERY_LENGTH,
             )
-            encoded = self.tokenizer.encode_batch(query)
+            encoded = self.tokenizer.encode_batch([query])
             if prev_padding is None:
                 self.tokenizer.no_padding()
             else:
@@ -104,8 +102,6 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
         return encoded
 
     def _tokenize_documents(self, documents: List[str]) -> List[Encoding]:
-        # ". " is added to a document to be replaced with a special document token
-        documents = [". " + doc for doc in documents]
         encoded = self.tokenizer.encode_batch(documents)
         return encoded
 
@@ -189,7 +185,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
             cuda=self.cuda,
             device_id=self.device_id,
         )
-        self.mask_token_id = self.special_token_to_id["[MASK]"]
+        self.mask_token_id = self.special_token_to_id[self.MASK_TOKEN]
         self.pad_token_id = self.tokenizer.padding["pad_id"]
         self.skip_list = {
             self.tokenizer.encode(symbol, add_special_tokens=False).ids[0]
@@ -230,7 +226,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
             **kwargs,
         )
 
-    def query_embed(self, query: Union[str, List[str]], **kwargs) -> Iterable[np.ndarray]:
+    def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[np.ndarray]:
         if isinstance(query, str):
             query = [query]
 

fastembed_gpu-0.4.2/fastembed/late_interaction/jina_colbert.py

@@ -0,0 +1,62 @@
+from typing import Any, Dict, List, Type
+
+import numpy as np
+
+from fastembed.late_interaction.colbert import Colbert
+from fastembed.text.onnx_text_model import TextEmbeddingWorker
+
+
+supported_jina_colbert_models = [
+    {
+        "model": "jinaai/jina-colbert-v2",
+        "dim": 128,
+        "description": "New model that expands capabilities of colbert-v1 with multilingual and context length of 8192, 2024 year",
+        "license": "cc-by-nc-4.0",
+        "size_in_GB": 2.24,
+        "sources": {
+            "hf": "jinaai/jina-colbert-v2",
+        },
+        "model_file": "onnx/model.onnx",
+        "additional_files": ["onnx/model.onnx_data"],
+    },
+]
+
+
+class JinaColbert(Colbert):
+    QUERY_MARKER_TOKEN_ID = 250002
+    DOCUMENT_MARKER_TOKEN_ID = 250003
+    MIN_QUERY_LENGTH = 31  # it's 32, we add one additional special token in the beginning
+    MASK_TOKEN = "<mask>"
+
+    @classmethod
+    def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
+        return JinaColbertEmbeddingWorker
+
+    @classmethod
+    def list_supported_models(cls) -> List[Dict[str, Any]]:
+        """Lists the supported models.
+
+        Returns:
+            List[Dict[str, Any]]: A list of dictionaries containing the model information.
+        """
+        return supported_jina_colbert_models
+
+    def _preprocess_onnx_input(
+        self, onnx_input: Dict[str, np.ndarray], is_doc: bool = True, **kwargs: Any
+    ) -> Dict[str, np.ndarray]:
+        onnx_input = super()._preprocess_onnx_input(onnx_input, is_doc)
+
+        # the attention mask for jina-colbert-v2 is always 1 in queries
+        if not is_doc:
+            onnx_input["attention_mask"][:] = 1
+        return onnx_input
+
+
+class JinaColbertEmbeddingWorker(TextEmbeddingWorker):
+    def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> JinaColbert:
+        return JinaColbert(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

{fastembed_gpu-0.4.0 → fastembed_gpu-0.4.2}/fastembed/late_interaction/late_interaction_text_embedding.py

@@ -4,15 +4,14 @@ import numpy as np
 
 from fastembed.common import OnnxProvider
 from fastembed.late_interaction.colbert import Colbert
+from fastembed.late_interaction.jina_colbert import JinaColbert
 from fastembed.late_interaction.late_interaction_embedding_base import (
     LateInteractionTextEmbeddingBase,
 )
 
 
 class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
-    EMBEDDINGS_REGISTRY: List[Type[LateInteractionTextEmbeddingBase]] = [
-        Colbert,
-    ]
+    EMBEDDINGS_REGISTRY: List[Type[LateInteractionTextEmbeddingBase]] = [Colbert, JinaColbert]
 
     @classmethod
     def list_supported_models(cls) -> List[Dict[str, Any]]:

{fastembed_gpu-0.4.0 → fastembed_gpu-0.4.2}/fastembed/parallel_processor.py

@@ -24,7 +24,7 @@ class QueueSignals(str, Enum):
 
 class Worker:
     @classmethod
-    def start(cls, **kwargs: Any) -> "Worker":
+    def start(cls, *args: Any, **kwargs: Any) -> "Worker":
         raise NotImplementedError()
 
     def process(self, items: Iterable[Tuple[int, Any]]) -> Iterable[Tuple[int, Any]]:

{fastembed_gpu-0.4.0 → fastembed_gpu-0.4.2}/fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py

@@ -38,6 +38,36 @@ supported_onnx_models = [
         "description": "BGE reranker base model for cross-encoder re-ranking.",
         "license": "mit",
     },
+    {
+        "model": "jinaai/jina-reranker-v1-tiny-en",
+        "size_in_GB": 0.13,
+        "sources": {
+            "hf": "jinaai/jina-reranker-v1-tiny-en",
+        },
+        "model_file": "onnx/model.onnx",
+        "description": "Designed for blazing-fast re-ranking with 8K context length and fewer parameters than jina-reranker-v1-turbo-en.",
+        "license": "apache-2.0",
+    },
+    {
+        "model": "jinaai/jina-reranker-v1-turbo-en",
+        "size_in_GB": 0.15,
+        "sources": {
+            "hf": "jinaai/jina-reranker-v1-turbo-en",
+        },
+        "model_file": "onnx/model.onnx",
+        "description": "Designed for blazing-fast re-ranking with 8K context length.",
+        "license": "apache-2.0",
+    },
+    {
+        "model": "jinaai/jina-reranker-v2-base-multilingual",
+        "size_in_GB": 1.11,
+        "sources": {
+            "hf": "jinaai/jina-reranker-v2-base-multilingual",
+        },
+        "model_file": "onnx/model.onnx",
+        "description": "A multi-lingual reranker model for cross-encoder re-ranking with 1K context length and sliding window",
+        "license": "cc-by-nc-4.0",
+    },
 ]
 
 

{fastembed_gpu-0.4.0 → fastembed_gpu-0.4.2}/fastembed/rerank/cross_encoder/onnx_text_model.py

@@ -4,7 +4,7 @@ from pathlib import Path
 import numpy as np
 from tokenizers import Encoding
 
-from fastembed.common.onnx_model import OnnxModel, OnnxProvider
+from fastembed.common.onnx_model import OnnxModel, OnnxProvider, OnnxOutputContext
 from fastembed.common.preprocessor_utils import load_tokenizer
 from fastembed.common.utils import iter_batch
 
@@ -34,7 +34,7 @@ class OnnxCrossEncoderModel(OnnxModel):
     def tokenize(self, query: str, documents: List[str], **kwargs) -> List[Encoding]:
         return self.tokenizer.encode_batch([(query, doc) for doc in documents])
 
-    def onnx_embed(self, query: str, documents: List[str], **kwargs) -> List[float]:
+    def onnx_embed(self, query: str, documents: List[str], **kwargs) -> OnnxOutputContext:
         tokenized_input = self.tokenize(query, documents, **kwargs)
 
         inputs = {
@@ -51,7 +51,7 @@ class OnnxCrossEncoderModel(OnnxModel):
 
         onnx_input = self._preprocess_onnx_input(inputs, **kwargs)
         outputs = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input)
-        return outputs[0][:, 0].tolist()
+        return OnnxOutputContext(model_output=outputs[0][:, 0].tolist())
 
     def _rerank_documents(
         self, query: str, documents: Iterable[str], batch_size: int, **kwargs
@@ -59,7 +59,7 @@ class OnnxCrossEncoderModel(OnnxModel):
         if not hasattr(self, "model") or self.model is None:
             self.load_onnx_model()
         for batch in iter_batch(documents, batch_size):
-            yield from self.onnx_embed(query, batch, **kwargs)
+            yield from self.onnx_embed(query, batch, **kwargs).model_output
 
     def _preprocess_onnx_input(
         self, onnx_input: Dict[str, np.ndarray], **kwargs

{fastembed_gpu-0.4.0 → fastembed_gpu-0.4.2}/fastembed/sparse/bm25.py

@@ -6,8 +6,7 @@ from typing import Any, Dict, Iterable, List, Optional, Tuple, Type, Union
 
 import mmh3
 import numpy as np
-from snowballstemmer import stemmer as get_stemmer
-
+from py_rust_stemmers import SnowballStemmer
 from fastembed.common.utils import (
     define_cache_dir,
     iter_batch,
@@ -130,7 +129,7 @@ class Bm25(SparseTextEmbeddingBase):
         self.punctuation = set(get_all_punctuation())
         self.stopwords = set(self._load_stopwords(self._model_dir, self.language))
 
-        self.stemmer = get_stemmer(language)
+        self.stemmer = SnowballStemmer(language)
         self.tokenizer = SimpleTokenizer
 
     @classmethod
@@ -235,7 +234,7 @@ class Bm25(SparseTextEmbeddingBase):
             if len(token) > self.token_max_length:
                 continue
 
-            stemmed_token = self.stemmer.stemWord(token.lower())
+            stemmed_token = self.stemmer.stem_word(token.lower())
 
             if stemmed_token:
                 stemmed_tokens.append(stemmed_token)

{fastembed_gpu-0.4.0 → fastembed_gpu-0.4.2}/fastembed/sparse/bm42.py

@@ -5,7 +5,7 @@ from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple, Type, U
 
 import mmh3
 import numpy as np
-from snowballstemmer import stemmer as get_stemmer
+from py_rust_stemmers import SnowballStemmer
 
 from fastembed.common import OnnxProvider
 from fastembed.common.onnx_model import OnnxOutputContext
@@ -119,8 +119,8 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         self.special_tokens = set()
         self.special_tokens_ids = set()
         self.punctuation = set(string.punctuation)
-        self.stopwords = set()
-        self.stemmer = get_stemmer(MODEL_TO_LANGUAGE[model_name])
+        self.stopwords = set(self._load_stopwords(self._model_dir))
+        self.stemmer = SnowballStemmer(MODEL_TO_LANGUAGE[model_name])
         self.alpha = alpha
 
         if not self.lazy_load:
@@ -152,7 +152,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
     def _stem_pair_tokens(self, tokens: List[Tuple[str, Any]]) -> List[Tuple[str, Any]]:
         result = []
         for token, value in tokens:
-            processed_token = self.stemmer.stemWord(token)
+            processed_token = self.stemmer.stem_word(token)
             result.append((processed_token, value))
         return result
 

{fastembed_gpu-0.4.0 → fastembed_gpu-0.4.2}/fastembed/sparse/utils/tokenizer.py

@@ -5,6 +5,7 @@ from typing import List
 
 
 class SimpleTokenizer:
+    @staticmethod
     def tokenize(text: str) -> List[str]:
         text = re.sub(r"[^\w]", " ", text.lower())
         text = re.sub(r"\s+", " ", text)

{fastembed_gpu-0.4.0 → fastembed_gpu-0.4.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "fastembed-gpu"
-version = "0.4.0"
+version = "0.4.2"
 description = "Fast, light, accurate library built for retrieval embedding generation"
 authors = ["Qdrant Team <info@qdrant.tech>", "NirantK <nirant.bits@gmail.com>"]
 license = "Apache License"
@@ -12,20 +12,20 @@ keywords = ["vector", "embedding", "neural", "search", "qdrant", "sentence-trans
 
 [tool.poetry.dependencies]
 python = ">=3.8.0,<3.13"
-onnxruntime-gpu = "^1.17.0"
+onnx = "^1.15.0"
+onnxruntime-gpu = ">=1.17.0,<1.20.0"
 tqdm = "^4.66"
 requests = "^2.31"
 tokenizers = ">=0.15,<1.0"
 huggingface-hub = ">=0.20,<1.0"
 loguru = "^0.7.2"
 numpy = [
-    { version = ">=1.21, <2", python = "<3.12" },
-    { version = ">=1.26, <2", python = ">=3.12" }
+    { version = ">=1.21", python = "<3.12" },
+    { version = ">=1.26", python = ">=3.12" }
 ]
 pillow = "^10.3.0"
-snowballstemmer = "^2.2.0"
-PyStemmer = "^2.2.0"
 mmh3 = "^4.1.0"
+py-rust-stemmers = "^0.1.0"
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^7.4.2"