fastembed-gpu 0.3.4__tar.gz → 0.3.6__tar.gz

{fastembed_gpu-0.3.4 → fastembed_gpu-0.3.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: fastembed-gpu
-Version: 0.3.4
+Version: 0.3.6
 Summary: Fast, light, accurate library built for retrieval embedding generation
 Home-page: https://github.com/qdrant/fastembed
 License: Apache License
@@ -40,9 +40,9 @@ The default text embedding (`TextEmbedding`) model is Flag Embedding, presented
 
 1. Light: FastEmbed is a lightweight library with few external dependencies. We don't require a GPU and don't download GBs of PyTorch dependencies, and instead use the ONNX Runtime. This makes it a great candidate for serverless runtimes like AWS Lambda. 
 
-2. Fast: FastEmbed is designed for speed. We use the ONNX Runtime, which is faster than PyTorch. We also use data-parallelism for encoding large datasets.
+2. Fast: FastEmbed is designed for speed. We use the ONNX Runtime, which is faster than PyTorch. We also use data parallelism for encoding large datasets.
 
-3. Accurate: FastEmbed is better than OpenAI Ada-002. We also [support](https://qdrant.github.io/fastembed/examples/Supported_Models/) an ever expanding set of models, including a few multilingual models.
+3. Accurate: FastEmbed is better than OpenAI Ada-002. We also [support](https://qdrant.github.io/fastembed/examples/Supported_Models/) an ever-expanding set of models, including a few multilingual models.
 
 ## 🚀 Installation
 
@@ -86,7 +86,7 @@ The list of all the available models can be found [here](https://qdrant.github.i
 from fastembed import TextEmbedding
 
 model = TextEmbedding(model_name="BAAI/bge-small-en-v1.5")
-embeddings = list(embedding_model.embed(documents))
+embeddings = list(model.embed(documents))
 
 # [
 #   array([-0.1115,  0.0097,  0.0052,  0.0195, ...], dtype=float32),
@@ -105,7 +105,7 @@ embeddings = list(embedding_model.embed(documents))
 from fastembed import SparseTextEmbedding
 
 model = SparseTextEmbedding(model_name="prithivida/Splade_PP_en_v1")
-embeddings = list(embedding_model.embed(documents))
+embeddings = list(model.embed(documents))
 
 # [
 #   SparseEmbedding(indices=[ 17, 123, 919, ... ], values=[0.71, 0.22, 0.39, ...]),
@@ -120,7 +120,7 @@ embeddings = list(embedding_model.embed(documents))
 from fastembed import SparseTextEmbedding
 
 model = SparseTextEmbedding(model_name="Qdrant/bm42-all-minilm-l6-v2-attentions")
-embeddings = list(embedding_model.embed(documents))
+embeddings = list(model.embed(documents))
 
 # [
 #   SparseEmbedding(indices=[ 17, 123, 919, ... ], values=[0.71, 0.22, 0.39, ...]),
@@ -136,7 +136,7 @@ embeddings = list(embedding_model.embed(documents))
 from fastembed import LateInteractionTextEmbedding
 
 model = LateInteractionTextEmbedding(model_name="colbert-ir/colbertv2.0")
-embeddings = list(embedding_model.embed(documents))
+embeddings = list(model.embed(documents))
 
 # [
 #   array([
@@ -161,7 +161,7 @@ images = [
 ]
 
 model = ImageEmbedding(model_name="Qdrant/clip-ViT-B-32-vision")
-embeddings = list(embedding_model.embed(images))
+embeddings = list(model.embed(images))
 
 # [
 #   array([-0.1115,  0.0097,  0.0052,  0.0195, ...], dtype=float32),
@@ -179,7 +179,7 @@ It requires installation of the `fastembed-gpu` package.
 pip install fastembed-gpu
 ```
 
-Check our [example](https://qdrant.github.io/fastembed/examples/FastEmbed_GPU/) for the detailed instructions and CUDA 12.x support.
+Check our [example](https://qdrant.github.io/fastembed/examples/FastEmbed_GPU/) for detailed instructions and CUDA 12.x support.
 
 ```python
 from fastembed import TextEmbedding

{fastembed_gpu-0.3.4 → fastembed_gpu-0.3.6}/README.md

@@ -8,9 +8,9 @@ The default text embedding (`TextEmbedding`) model is Flag Embedding, presented
 
 1. Light: FastEmbed is a lightweight library with few external dependencies. We don't require a GPU and don't download GBs of PyTorch dependencies, and instead use the ONNX Runtime. This makes it a great candidate for serverless runtimes like AWS Lambda. 
 
-2. Fast: FastEmbed is designed for speed. We use the ONNX Runtime, which is faster than PyTorch. We also use data-parallelism for encoding large datasets.
+2. Fast: FastEmbed is designed for speed. We use the ONNX Runtime, which is faster than PyTorch. We also use data parallelism for encoding large datasets.
 
-3. Accurate: FastEmbed is better than OpenAI Ada-002. We also [support](https://qdrant.github.io/fastembed/examples/Supported_Models/) an ever expanding set of models, including a few multilingual models.
+3. Accurate: FastEmbed is better than OpenAI Ada-002. We also [support](https://qdrant.github.io/fastembed/examples/Supported_Models/) an ever-expanding set of models, including a few multilingual models.
 
 ## 🚀 Installation
 
@@ -54,7 +54,7 @@ The list of all the available models can be found [here](https://qdrant.github.i
 from fastembed import TextEmbedding
 
 model = TextEmbedding(model_name="BAAI/bge-small-en-v1.5")
-embeddings = list(embedding_model.embed(documents))
+embeddings = list(model.embed(documents))
 
 # [
 #   array([-0.1115,  0.0097,  0.0052,  0.0195, ...], dtype=float32),
@@ -73,7 +73,7 @@ embeddings = list(embedding_model.embed(documents))
 from fastembed import SparseTextEmbedding
 
 model = SparseTextEmbedding(model_name="prithivida/Splade_PP_en_v1")
-embeddings = list(embedding_model.embed(documents))
+embeddings = list(model.embed(documents))
 
 # [
 #   SparseEmbedding(indices=[ 17, 123, 919, ... ], values=[0.71, 0.22, 0.39, ...]),
@@ -88,7 +88,7 @@ embeddings = list(embedding_model.embed(documents))
 from fastembed import SparseTextEmbedding
 
 model = SparseTextEmbedding(model_name="Qdrant/bm42-all-minilm-l6-v2-attentions")
-embeddings = list(embedding_model.embed(documents))
+embeddings = list(model.embed(documents))
 
 # [
 #   SparseEmbedding(indices=[ 17, 123, 919, ... ], values=[0.71, 0.22, 0.39, ...]),
@@ -104,7 +104,7 @@ embeddings = list(embedding_model.embed(documents))
 from fastembed import LateInteractionTextEmbedding
 
 model = LateInteractionTextEmbedding(model_name="colbert-ir/colbertv2.0")
-embeddings = list(embedding_model.embed(documents))
+embeddings = list(model.embed(documents))
 
 # [
 #   array([
@@ -129,7 +129,7 @@ images = [
 ]
 
 model = ImageEmbedding(model_name="Qdrant/clip-ViT-B-32-vision")
-embeddings = list(embedding_model.embed(images))
+embeddings = list(model.embed(images))
 
 # [
 #   array([-0.1115,  0.0097,  0.0052,  0.0195, ...], dtype=float32),
@@ -147,7 +147,7 @@ It requires installation of the `fastembed-gpu` package.
 pip install fastembed-gpu
 ```
 
-Check our [example](https://qdrant.github.io/fastembed/examples/FastEmbed_GPU/) for the detailed instructions and CUDA 12.x support.
+Check our [example](https://qdrant.github.io/fastembed/examples/FastEmbed_GPU/) for detailed instructions and CUDA 12.x support.
 
 ```python
 from fastembed import TextEmbedding

fastembed_gpu-0.3.6/fastembed/common/__init__.py

@@ -0,0 +1,3 @@
+from fastembed.common.types import ImageInput, OnnxProvider, PathInput, PilInput
+
+__all__ = ["OnnxProvider", "ImageInput", "PathInput", "PilInput"]

{fastembed_gpu-0.3.4 → fastembed_gpu-0.3.6}/fastembed/common/model_management.py

@@ -148,7 +148,9 @@ class ModelManagement:
             # Open the tar.gz file
             with tarfile.open(targz_path, "r:gz") as tar:
                 # Extract all files into the cache directory
-                tar.extractall(path=cache_dir)
+                tar.extractall(
+                    path=cache_dir,
+                )
         except tarfile.TarError as e:
             # If any error occurs while opening or extracting the tar.gz file,
             # delete the cache directory (if it was created in this function)

{fastembed_gpu-0.3.4 → fastembed_gpu-0.3.6}/fastembed/common/preprocessor_utils.py

@@ -1,7 +1,6 @@
 import json
 from pathlib import Path
 from typing import Tuple
-
 from tokenizers import AddedToken, Tokenizer
 
 from fastembed.image.transform.operators import Compose
@@ -18,7 +17,7 @@ def load_special_tokens(model_dir: Path) -> dict:
     return tokens_map
 
 
-def load_tokenizer(model_dir: Path, max_length: int = 512) -> Tuple[Tokenizer, dict]:
+def load_tokenizer(model_dir: Path) -> Tuple[Tokenizer, dict]:
     config_path = model_dir / "config.json"
     if not config_path.exists():
         raise ValueError(f"Could not find config.json in {model_dir}")
@@ -36,13 +35,20 @@ def load_tokenizer(model_dir: Path, max_length: int = 512) -> Tuple[Tokenizer, d
 
     with open(str(tokenizer_config_path)) as tokenizer_config_file:
         tokenizer_config = json.load(tokenizer_config_file)
+        assert (
+            "model_max_length" in tokenizer_config or "max_length" in tokenizer_config
+        ), "Models without model_max_length or max_length are not supported."
+        if "model_max_length" not in tokenizer_config:
+            max_context = tokenizer_config["max_length"]
+        elif "max_length" not in tokenizer_config:
+            max_context = tokenizer_config["model_max_length"]
+        else:
+            max_context = min(tokenizer_config["model_max_length"], tokenizer_config["max_length"])
 
     tokens_map = load_special_tokens(model_dir)
 
     tokenizer = Tokenizer.from_file(str(tokenizer_path))
-    tokenizer.enable_truncation(
-        max_length=min(tokenizer_config["model_max_length"], max_length)
-    )
+    tokenizer.enable_truncation(max_length=max_context)
     tokenizer.enable_padding(
         pad_id=config.get("pad_token_id", 0), pad_token=tokenizer_config["pad_token"]
     )

{fastembed_gpu-0.3.4 → fastembed_gpu-0.3.6}/fastembed/common/types.py

@@ -1,5 +1,6 @@
 import os
 import sys
+from PIL import Image
 from typing import Any, Dict, Iterable, Tuple, Union
 
 if sys.version_info >= (3, 10):
@@ -9,6 +10,7 @@ else:
 
 
 PathInput: TypeAlias = Union[str, os.PathLike]
-ImageInput: TypeAlias = Union[PathInput, Iterable[PathInput]]
+PilInput: TypeAlias = Union[Image.Image, Iterable[Image.Image]]
+ImageInput: TypeAlias = Union[PathInput, Iterable[PathInput], PilInput]
 
 OnnxProvider: TypeAlias = Union[str, Tuple[str, Dict[Any, Any]]]

{fastembed_gpu-0.3.4 → fastembed_gpu-0.3.6}/fastembed/image/onnx_embedding.py

@@ -12,7 +12,7 @@ supported_onnx_models = [
     {
         "model": "Qdrant/clip-ViT-B-32-vision",
         "dim": 512,
-        "description": "CLIP vision encoder based on ViT-B/32",
+        "description": "Image embeddings, Multimodal (text&image), 2021 year",
         "size_in_GB": 0.34,
         "sources": {
             "hf": "Qdrant/clip-ViT-B-32-vision",
@@ -22,7 +22,7 @@ supported_onnx_models = [
     {
         "model": "Qdrant/resnet50-onnx",
         "dim": 2048,
-        "description": "ResNet-50 from `Deep Residual Learning for Image Recognition <https://arxiv.org/abs/1512.03385>`__.",
+        "description": "Image embeddings, Unimodal (image), 2016 year",
         "size_in_GB": 0.1,
         "sources": {
             "hf": "Qdrant/resnet50-onnx",
@@ -32,7 +32,7 @@ supported_onnx_models = [
     {
         "model": "Qdrant/Unicom-ViT-B-16",
         "dim": 768,
-        "description": "Unicom Unicom-ViT-B-16 from open-metric-learning",
+        "description": "Image embeddings (more detailed than Unicom-ViT-B-32), Multimodal (text&image), 2023 year",
         "size_in_GB": 0.82,
         "sources": {
             "hf": "Qdrant/Unicom-ViT-B-16",
@@ -42,7 +42,7 @@ supported_onnx_models = [
     {
         "model": "Qdrant/Unicom-ViT-B-32",
         "dim": 512,
-        "description": "Unicom Unicom-ViT-B-32 from open-metric-learning",
+        "description": "Image embeddings, Multimodal (text&image), 2023 year",
         "size_in_GB": 0.48,
         "sources": {
             "hf": "Qdrant/Unicom-ViT-B-32",
@@ -51,6 +51,7 @@ supported_onnx_models = [
     },
 ]
 
+
 class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
     def __init__(
         self,
@@ -141,16 +142,10 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
 
         return onnx_input
 
-    def _post_process_onnx_output(
-        self, output: OnnxOutputContext
-    ) -> Iterable[np.ndarray]:
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
         return normalize(output.model_output).astype(np.float32)
 
 
 class OnnxImageEmbeddingWorker(ImageEmbeddingWorker):
-    def init_embedding(
-        self, model_name: str, cache_dir: str, **kwargs
-    ) -> OnnxImageEmbedding:
-        return OnnxImageEmbedding(
-            model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs
-        )
+    def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> OnnxImageEmbedding:
+        return OnnxImageEmbedding(model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs)

{fastembed_gpu-0.3.4 → fastembed_gpu-0.3.6}/fastembed/image/onnx_image_model.py

@@ -7,7 +7,7 @@ from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple, Type
 import numpy as np
 from PIL import Image
 
-from fastembed.common import ImageInput, OnnxProvider, PathInput
+from fastembed.common import ImageInput, OnnxProvider
 from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxOutputContext, T
 from fastembed.common.preprocessor_utils import load_preprocessor
 from fastembed.common.utils import iter_batch
@@ -54,9 +54,12 @@ class OnnxImageModel(OnnxModel[T]):
     def _build_onnx_input(self, encoded: np.ndarray) -> Dict[str, np.ndarray]:
         return {node.name: encoded for node in self.model.get_inputs()}
 
-    def onnx_embed(self, images: List[PathInput], **kwargs) -> OnnxOutputContext:
+    def onnx_embed(self, images: List[ImageInput], **kwargs) -> OnnxOutputContext:
         with contextlib.ExitStack():
-            image_files = [Image.open(image) for image in images]
+            image_files = [
+                Image.open(image) if not isinstance(image, Image.Image) else image
+                for image in images
+            ]
             encoded = self.processor(image_files)
         onnx_input = self._build_onnx_input(encoded)
         onnx_input = self._preprocess_onnx_input(onnx_input)
@@ -75,7 +78,11 @@ class OnnxImageModel(OnnxModel[T]):
     ) -> Iterable[T]:
         is_small = False
 
-        if isinstance(images, str) or isinstance(images, Path):
+        if (
+            isinstance(images, str)
+            or isinstance(images, Path)
+            or (isinstance(images, Image.Image))
+        ):
             images = [images]
             is_small = True
 
@@ -90,9 +97,7 @@ class OnnxImageModel(OnnxModel[T]):
             for batch in iter_batch(images, batch_size):
                 yield from self._post_process_onnx_output(self.onnx_embed(batch))
         else:
-            start_method = (
-                "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
-            )
+            start_method = "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
             params = {"model_name": model_name, "cache_dir": cache_dir, **kwargs}
             pool = ParallelWorkerPool(
                 parallel, self._get_worker_class(), start_method=start_method

{fastembed_gpu-0.3.4 → fastembed_gpu-0.3.6}/fastembed/late_interaction/colbert.py

@@ -22,7 +22,17 @@ supported_colbert_models = [
             "hf": "colbert-ir/colbertv2.0",
         },
         "model_file": "model.onnx",
-    }
+    },
+    {
+        "model": "answerdotai/answerai-colbert-small-v1",
+        "dim": 96,
+        "description": "Text embeddings, Unimodal (text), Multilingual (~100 languages), 512 input tokens truncation, 2024 year",
+        "size_in_GB": 0.13,
+        "sources": {
+            "hf": "answerdotai/answerai-colbert-small-v1",
+        },
+        "model_file": "vespa_colbert.onnx",
+    },
 ]
 
 

{fastembed_gpu-0.3.4 → fastembed_gpu-0.3.6}/fastembed/sparse/bm25.py

@@ -17,6 +17,38 @@ from fastembed.sparse.sparse_embedding_base import (
 )
 from fastembed.sparse.utils.tokenizer import WordTokenizer
 
+supported_languages = [
+    "arabic",
+    "azerbaijani",
+    "basque",
+    "bengali",
+    "catalan",
+    "chinese",
+    "danish",
+    "dutch",
+    "english",
+    "finnish",
+    "french",
+    "german",
+    "greek",
+    "hebrew",
+    "hinglish",
+    "hungarian",
+    "indonesian",
+    "italian",
+    "kazakh",
+    "nepali",
+    "norwegian",
+    "portuguese",
+    "romanian",
+    "russian",
+    "slovene",
+    "spanish",
+    "swedish",
+    "tajik",
+    "turkish",
+]
+
 supported_bm25_models = [
     {
         "model": "Qdrant/bm25",
@@ -26,15 +58,11 @@ supported_bm25_models = [
             "hf": "Qdrant/bm25",
         },
         "model_file": "mock.file",  # bm25 does not require a model, so we just use a mock
-        "additional_files": ["stopwords.txt"],
+        "additional_files": [f"{lang}.txt" for lang in supported_languages],
         "requires_idf": True,
     },
 ]
 
-MODEL_TO_LANGUAGE = {
-    "Qdrant/bm25": "english",
-}
-
 
 class Bm25(SparseTextEmbeddingBase):
     """Implements traditional BM25 in a form of sparse embeddings.
@@ -71,10 +99,16 @@ class Bm25(SparseTextEmbeddingBase):
         k: float = 1.2,
         b: float = 0.75,
         avg_len: float = 256.0,
+        language: str = "english",
         **kwargs,
     ):
         super().__init__(model_name, cache_dir, **kwargs)
 
+        if language not in supported_languages:
+            raise ValueError(f"{language} language is not supported")
+        else:
+            self.language = language
+
         self.k = k
         self.b = b
         self.avg_len = avg_len
@@ -87,8 +121,8 @@ class Bm25(SparseTextEmbeddingBase):
         )
 
         self.punctuation = set(string.punctuation)
-        self.stopwords = set(self._load_stopwords(model_dir))
-        self.stemmer = get_stemmer(MODEL_TO_LANGUAGE[model_name])
+        self.stopwords = set(self._load_stopwords(model_dir, self.language))
+        self.stemmer = get_stemmer(language)
         self.tokenizer = WordTokenizer
 
     @classmethod
@@ -101,8 +135,8 @@ class Bm25(SparseTextEmbeddingBase):
         return supported_bm25_models
 
     @classmethod
-    def _load_stopwords(cls, model_dir: Path) -> List[str]:
-        stopwords_path = model_dir / "stopwords.txt"
+    def _load_stopwords(cls, model_dir: Path, language: str) -> List[str]:
+        stopwords_path = model_dir / f"{language}.txt"
         if not stopwords_path.exists():
             return []
 

{fastembed_gpu-0.3.4 → fastembed_gpu-0.3.6}/fastembed/sparse/sparse_text_embedding.py

@@ -8,6 +8,7 @@ from fastembed.sparse.sparse_embedding_base import (
     SparseTextEmbeddingBase,
 )
 from fastembed.sparse.splade_pp import SpladePP
+import warnings
 
 
 class SparseTextEmbedding(SparseTextEmbeddingBase):
@@ -50,13 +51,18 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
         **kwargs,
     ):
         super().__init__(model_name, cache_dir, threads, **kwargs)
+        if model_name == "prithvida/Splade_PP_en_v1":
+            warnings.warn(
+                "The right spelling is prithivida/Splade_PP_en_v1. "
+                "Support of this name will be removed soon, please fix the model_name",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            model_name = "prithivida/Splade_PP_en_v1"
 
         for EMBEDDING_MODEL_TYPE in self.EMBEDDINGS_REGISTRY:
             supported_models = EMBEDDING_MODEL_TYPE.list_supported_models()
-            if any(
-                model_name.lower() == model["model"].lower()
-                for model in supported_models
-            ):
+            if any(model_name.lower() == model["model"].lower() for model in supported_models):
                 self.model = EMBEDDING_MODEL_TYPE(
                     model_name,
                     cache_dir,
@@ -95,9 +101,7 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
         """
         yield from self.model.embed(documents, batch_size, parallel, **kwargs)
 
-    def query_embed(
-        self, query: Union[str, Iterable[str]], **kwargs
-    ) -> Iterable[SparseEmbedding]:
+    def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[SparseEmbedding]:
         """
         Embeds queries
 

{fastembed_gpu-0.3.4 → fastembed_gpu-0.3.6}/fastembed/sparse/splade_pp.py

@@ -1,7 +1,6 @@
 from typing import Any, Dict, Iterable, List, Optional, Sequence, Type, Union
 
 import numpy as np
-
 from fastembed.common import OnnxProvider
 from fastembed.common.onnx_model import OnnxOutputContext
 from fastembed.common.utils import define_cache_dir
@@ -13,9 +12,9 @@ from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker
 
 supported_splade_models = [
     {
-        "model": "prithvida/Splade_PP_en_v1",
+        "model": "prithivida/Splade_PP_en_v1",
         "vocab_size": 30522,
-        "description": "Misspelled version of the model. Retained for backward compatibility. Independent Implementation of SPLADE++ Model for English",
+        "description": "Independent Implementation of SPLADE++ Model for English",
         "size_in_GB": 0.532,
         "sources": {
             "hf": "Qdrant/SPLADE_PP_en_v1",
@@ -23,7 +22,7 @@ supported_splade_models = [
         "model_file": "model.onnx",
     },
     {
-        "model": "prithivida/Splade_PP_en_v1",
+        "model": "prithvida/Splade_PP_en_v1",
         "vocab_size": 30522,
         "description": "Independent Implementation of SPLADE++ Model for English",
         "size_in_GB": 0.532,
@@ -78,7 +77,6 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         Raises:
             ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
         """
-
         super().__init__(model_name, cache_dir, threads, **kwargs)
 
         model_description = self._get_model_description(model_name)

{fastembed_gpu-0.3.4 → fastembed_gpu-0.3.6}/fastembed/text/clip_embedding.py

@@ -10,7 +10,7 @@ supported_clip_models = [
     {
         "model": "Qdrant/clip-ViT-B-32-text",
         "dim": 512,
-        "description": "CLIP text encoder",
+        "description": "Text embeddings, Multimodal (text&image), English, 77 input tokens truncation, Prefixes for queries/documents: not necessary, 2021 year",
         "size_in_GB": 0.25,
         "sources": {
             "hf": "Qdrant/clip-ViT-B-32-text",

{fastembed_gpu-0.3.4 → fastembed_gpu-0.3.6}/fastembed/text/e5_onnx_embedding.py

@@ -9,7 +9,7 @@ supported_multilingual_e5_models = [
     {
         "model": "intfloat/multilingual-e5-large",
         "dim": 1024,
-        "description": "Multilingual model, e5-large. Recommend using this model for non-English languages",
+        "description": "Text embeddings, Unimodal (text), Multilingual (~100 languages), 512 input tokens truncation, Prefixes for queries/documents: necessary, 2024 year",
         "size_in_GB": 2.24,
         "sources": {
             "url": "https://storage.googleapis.com/qdrant-fastembed/fast-multilingual-e5-large.tar.gz",
@@ -21,7 +21,7 @@ supported_multilingual_e5_models = [
     {
         "model": "sentence-transformers/paraphrase-multilingual-mpnet-base-v2",
         "dim": 768,
-        "description": "Sentence-transformers model for tasks like clustering or semantic search",
+        "description": "Text embeddings, Unimodal (text), Multilingual (~50 languages), 384 input tokens truncation, Prefixes for queries/documents: not necessary, 2021 year",
         "size_in_GB": 1.00,
         "sources": {
             "hf": "xenova/paraphrase-multilingual-mpnet-base-v2",

{fastembed_gpu-0.3.4 → fastembed_gpu-0.3.6}/fastembed/text/onnx_embedding.py

@@ -12,7 +12,7 @@ supported_onnx_models = [
     {
         "model": "BAAI/bge-base-en",
         "dim": 768,
-        "description": "Base English model",
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: necessary, 2023 year",
         "size_in_GB": 0.42,
         "sources": {
             "url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en.tar.gz",
@@ -22,7 +22,7 @@ supported_onnx_models = [
     {
         "model": "BAAI/bge-base-en-v1.5",
         "dim": 768,
-        "description": "Base English model, v1.5",
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: not so necessary, 2023 year",
         "size_in_GB": 0.21,
         "sources": {
             "url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en-v1.5.tar.gz",
@@ -33,7 +33,7 @@ supported_onnx_models = [
     {
         "model": "BAAI/bge-large-en-v1.5",
         "dim": 1024,
-        "description": "Large English model, v1.5",
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: not so necessary, 2023 year",
         "size_in_GB": 1.20,
         "sources": {
             "hf": "qdrant/bge-large-en-v1.5-onnx",
@@ -43,7 +43,7 @@ supported_onnx_models = [
     {
         "model": "BAAI/bge-small-en",
         "dim": 384,
-        "description": "Fast English model",
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: necessary, 2023 year",
         "size_in_GB": 0.13,
         "sources": {
             "url": "https://storage.googleapis.com/qdrant-fastembed/BAAI-bge-small-en.tar.gz",
@@ -53,7 +53,7 @@ supported_onnx_models = [
     {
         "model": "BAAI/bge-small-en-v1.5",
         "dim": 384,
-        "description": "Fast and Default English model",
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: not so necessary, 2023 year",
         "size_in_GB": 0.067,
         "sources": {
             "hf": "qdrant/bge-small-en-v1.5-onnx-q",
@@ -63,7 +63,7 @@ supported_onnx_models = [
     {
         "model": "BAAI/bge-small-zh-v1.5",
         "dim": 512,
-        "description": "Fast and recommended Chinese model",
+        "description": "Text embeddings, Unimodal (text), Chinese, 512 input tokens truncation, Prefixes for queries/documents: not so necessary, 2023 year",
         "size_in_GB": 0.09,
         "sources": {
             "url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-zh-v1.5.tar.gz",
@@ -73,7 +73,7 @@ supported_onnx_models = [
     {
         "model": "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
         "dim": 384,
-        "description": "Sentence Transformer model, paraphrase-multilingual-MiniLM-L12-v2",
+        "description": "Text embeddings, Unimodal (text), Multilingual (~50 languages), 512 input tokens truncation, Prefixes for queries/documents: not necessary, 2019 year",
         "size_in_GB": 0.22,
         "sources": {
             "hf": "qdrant/paraphrase-multilingual-MiniLM-L12-v2-onnx-Q",
@@ -83,7 +83,7 @@ supported_onnx_models = [
     {
         "model": "thenlper/gte-large",
         "dim": 1024,
-        "description": "Large general text embeddings model",
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: not necessary, 2023 year",
         "size_in_GB": 1.20,
         "sources": {
             "hf": "qdrant/gte-large-onnx",
@@ -93,7 +93,7 @@ supported_onnx_models = [
     {
         "model": "mixedbread-ai/mxbai-embed-large-v1",
         "dim": 1024,
-        "description": "MixedBread Base sentence embedding model, does well on MTEB",
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: necessary, 2024 year",
         "size_in_GB": 0.64,
         "sources": {
             "hf": "mixedbread-ai/mxbai-embed-large-v1",
@@ -103,7 +103,7 @@ supported_onnx_models = [
     {
         "model": "snowflake/snowflake-arctic-embed-xs",
         "dim": 384,
-        "description": "Based on all-MiniLM-L6-v2 model with only 22m parameters, ideal for latency/TCO budgets.",
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: necessary, 2024 year",
         "size_in_GB": 0.09,
         "sources": {
             "hf": "snowflake/snowflake-arctic-embed-xs",
@@ -113,7 +113,7 @@ supported_onnx_models = [
     {
         "model": "snowflake/snowflake-arctic-embed-s",
         "dim": 384,
-        "description": "Based on infloat/e5-small-unsupervised, does not trade off retrieval accuracy for its small size.",
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: necessary, 2024 year",
         "size_in_GB": 0.13,
         "sources": {
             "hf": "snowflake/snowflake-arctic-embed-s",
@@ -123,7 +123,7 @@ supported_onnx_models = [
     {
         "model": "snowflake/snowflake-arctic-embed-m",
         "dim": 768,
-        "description": "Based on intfloat/e5-base-unsupervised model, provides the best retrieval without slowing down inference.",
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: necessary, 2024 year",
         "size_in_GB": 0.43,
         "sources": {
             "hf": "Snowflake/snowflake-arctic-embed-m",
@@ -133,7 +133,7 @@ supported_onnx_models = [
     {
         "model": "snowflake/snowflake-arctic-embed-m-long",
         "dim": 768,
-        "description": "Based on nomic-ai/nomic-embed-text-v1-unsupervised model, 8192 context-length model",
+        "description": "Text embeddings, Unimodal (text), English, 2048 input tokens truncation, Prefixes for queries/documents: necessary, 2024 year",
         "size_in_GB": 0.54,
         "sources": {
             "hf": "snowflake/snowflake-arctic-embed-m-long",
@@ -143,7 +143,7 @@ supported_onnx_models = [
     {
         "model": "snowflake/snowflake-arctic-embed-l",
         "dim": 1024,
-        "description": "Based on intfloat/e5-large-unsupervised, large model for most accurate retrieval.",
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: necessary, 2024 year",
         "size_in_GB": 1.02,
         "sources": {
             "hf": "snowflake/snowflake-arctic-embed-l",
@@ -244,9 +244,7 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
         """
         return onnx_input
 
-    def _post_process_onnx_output(
-        self, output: OnnxOutputContext
-    ) -> Iterable[np.ndarray]:
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
         embeddings = output.model_output
         return normalize(embeddings[:, 0]).astype(np.float32)
 
@@ -258,6 +256,4 @@ class OnnxTextEmbeddingWorker(TextEmbeddingWorker):
         cache_dir: str,
         **kwargs,
     ) -> OnnxTextEmbedding:
-        return OnnxTextEmbedding(
-            model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs
-        )
+        return OnnxTextEmbedding(model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs)

{fastembed_gpu-0.3.4 → fastembed_gpu-0.3.6}/fastembed/text/pooled_embedding.py

@@ -11,7 +11,7 @@ supported_pooled_models = [
     {
         "model": "nomic-ai/nomic-embed-text-v1.5",
         "dim": 768,
-        "description": "8192 context length english model",
+        "description": "Text embeddings, Multimodal (text, image), English, 8192 input tokens truncation, Prefixes for queries/documents: necessary, 2024 year",
         "size_in_GB": 0.52,
         "sources": {
             "hf": "nomic-ai/nomic-embed-text-v1.5",
@@ -21,7 +21,7 @@ supported_pooled_models = [
     {
         "model": "nomic-ai/nomic-embed-text-v1.5-Q",
         "dim": 768,
-        "description": "Quantized 8192 context length english model",
+        "description": "Text embeddings, Multimodal (text, image), English, 8192 input tokens truncation, Prefixes for queries/documents: necessary, 2024 year",
         "size_in_GB": 0.13,
         "sources": {
             "hf": "nomic-ai/nomic-embed-text-v1.5",
@@ -31,7 +31,7 @@ supported_pooled_models = [
     {
         "model": "nomic-ai/nomic-embed-text-v1",
         "dim": 768,
-        "description": "8192 context length english model",
+        "description": "Text embeddings, Multimodal (text, image), English, 8192 input tokens truncation, Prefixes for queries/documents: necessary, 2024 year",
         "size_in_GB": 0.52,
         "sources": {
             "hf": "nomic-ai/nomic-embed-text-v1",

{fastembed_gpu-0.3.4 → fastembed_gpu-0.3.6}/fastembed/text/pooled_normalized_embedding.py

@@ -12,7 +12,7 @@ supported_pooled_normalized_models = [
     {
         "model": "sentence-transformers/all-MiniLM-L6-v2",
         "dim": 384,
-        "description": "Sentence Transformer model, MiniLM-L6-v2",
+        "description": "Text embeddings, Unimodal (text), English, 256 input tokens truncation, Prefixes for queries/documents: not necessary, 2021 year",
         "size_in_GB": 0.09,
         "sources": {
             "url": "https://storage.googleapis.com/qdrant-fastembed/sentence-transformers-all-MiniLM-L6-v2.tar.gz",
@@ -23,7 +23,7 @@ supported_pooled_normalized_models = [
     {
         "model": "jinaai/jina-embeddings-v2-base-en",
         "dim": 768,
-        "description": "English embedding model supporting 8192 sequence length",
+        "description": "Text embeddings, Unimodal (text), English, 8192 input tokens truncation, Prefixes for queries/documents: not necessary, 2023 year",
         "size_in_GB": 0.52,
         "sources": {"hf": "xenova/jina-embeddings-v2-base-en"},
         "model_file": "onnx/model.onnx",
@@ -31,7 +31,7 @@ supported_pooled_normalized_models = [
     {
         "model": "jinaai/jina-embeddings-v2-small-en",
         "dim": 512,
-        "description": "English embedding model supporting 8192 sequence length",
+        "description": "Text embeddings, Unimodal (text), English, 8192 input tokens truncation, Prefixes for queries/documents: not necessary, 2023 year",
         "size_in_GB": 0.12,
         "sources": {"hf": "xenova/jina-embeddings-v2-small-en"},
         "model_file": "onnx/model.onnx",
@@ -39,11 +39,19 @@ supported_pooled_normalized_models = [
     {
         "model": "jinaai/jina-embeddings-v2-base-de",
         "dim": 768,
-        "description": "German embedding model supporting 8192 sequence length",
+        "description": "Text embeddings, Unimodal (text), Multilingual (German, English), 8192 input tokens truncation, Prefixes for queries/documents: not necessary, 2024 year",
         "size_in_GB": 0.32,
         "sources": {"hf": "jinaai/jina-embeddings-v2-base-de"},
         "model_file": "onnx/model_fp16.onnx",
     },
+    {
+        "model": "jinaai/jina-embeddings-v2-base-code",
+        "dim": 768,
+        "description": "Text embeddings, Unimodal (text), Multilingual (English, 30 programming languages), 8192 input tokens truncation, Prefixes for queries/documents: not necessary, 2024 year",
+        "size_in_GB": 0.64,
+        "sources": {"hf": "jinaai/jina-embeddings-v2-base-code"},
+        "model_file": "onnx/model.onnx",
+    },
 ]
 
 

{fastembed_gpu-0.3.4 → fastembed_gpu-0.3.6}/pyproject.toml

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

fastembed_gpu-0.3.4/fastembed/common/__init__.py

@@ -1,3 +0,0 @@
-from fastembed.common.types import ImageInput, OnnxProvider, PathInput
-
-__all__ = ["OnnxProvider", "ImageInput", "PathInput"]