fastembed 0.2.6__tar.gz → 0.2.7__tar.gz

{fastembed-0.2.6 → fastembed-0.2.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: fastembed
-Version: 0.2.6
+Version: 0.2.7
 Summary: Fast, light, accurate library built for retrieval embedding generation
 Home-page: https://github.com/qdrant/fastembed
 License: Apache License
@@ -22,7 +22,7 @@ Requires-Dist: numpy (>=1.26) ; python_version >= "3.12"
 Requires-Dist: onnx (>=1.15.0,<2.0.0)
 Requires-Dist: onnxruntime (>=1.17.0,<2.0.0)
 Requires-Dist: requests (>=2.31,<3.0)
-Requires-Dist: tokenizers (>=0.15.1,<0.16.0)
+Requires-Dist: tokenizers (>=0.15,<0.16)
 Requires-Dist: tqdm (>=4.66,<5.0)
 Project-URL: Repository, https://github.com/qdrant/fastembed
 Description-Content-Type: text/markdown
@@ -31,7 +31,7 @@ Description-Content-Type: text/markdown
 
 FastEmbed is a lightweight, fast, Python library built for embedding generation. We [support popular text models](https://qdrant.github.io/fastembed/examples/Supported_Models/). Please [open a GitHub issue](https://github.com/qdrant/fastembed/issues/new) if you want us to add a new model.
 
-The default text embedding (`TextEmbedding`) model is Flag Embedding, presented in the [MTEB](https://huggingface.co/spaces/mteb/leaderboard) leaderboard. It supports "query" and "passage" prefixes for the input text. Here is an example for [Retrieval Embedding Generation](https://qdrant.github.io/fastembed/examples/Retrieval_with_FastEmbed/) and how to use [FastEmbed with Qdrant](https://qdrant.github.io/fastembed/examples/Usage_With_Qdrant/).
+The default text embedding (`TextEmbedding`) model is Flag Embedding, presented in the [MTEB](https://huggingface.co/spaces/mteb/leaderboard) leaderboard. It supports "query" and "passage" prefixes for the input text. Here is an example for [Retrieval Embedding Generation](https://qdrant.github.io/fastembed/qdrant/Retrieval_with_FastEmbed/) and how to use [FastEmbed with Qdrant](https://qdrant.github.io/fastembed/qdrant/Usage_With_Qdrant/).
 
 ## 📈 Why FastEmbed?
 
@@ -43,12 +43,18 @@ The default text embedding (`TextEmbedding`) model is Flag Embedding, presented
 
 ## 🚀 Installation
 
-To install the FastEmbed library, pip works:
+To install the FastEmbed library, pip works best. You can install it with or without GPU support:
 
 ```bash
 pip install fastembed
 ```
 
+### ⚡️ With GPU
+
+```bash
+pip install fastembed-gpu
+```
+
 ## 📖 Quickstart
 
 ```python
@@ -71,6 +77,23 @@ embeddings_list = list(embedding_model.embed(documents))
 len(embeddings_list[0]) # Vector of 384 dimensions
 ```
 
+### ⚡️ FastEmbed on a GPU
+
+FastEmbed supports running on GPU devices. It requires installation of the `fastembed-gpu` package.
+Make sure not to have the `fastembed` package installed, as it might interfere with the `fastembed-gpu` package.
+
+```bash
+pip install fastembed-gpu
+``` 
+
+```python
+from fastembed import TextEmbedding
+
+embedding_model = TextEmbedding(model_name="BAAI/bge-small-en-v1.5", providers=["CUDAExecutionProvider"])
+print("The model BAAI/bge-small-en-v1.5 is ready to use on a GPU.")
+
+```
+
 ## Usage with Qdrant
 
 Installation with Qdrant Client in Python:
@@ -79,7 +102,13 @@ Installation with Qdrant Client in Python:
 pip install qdrant-client[fastembed]
 ```
 
-You might have to use ```pip install 'qdrant-client[fastembed]'``` on zsh.
+or 
+
+```bash
+pip install qdrant-client[fastembed-gpu]
+```
+
+You might have to use quotes ```pip install 'qdrant-client[fastembed]'``` on zsh.
 
 ```python
 from qdrant_client import QdrantClient
@@ -115,8 +144,3 @@ search_result = client.query(
 )
 print(search_result)
 ```
-
-#### Similar Work
-
-Ilyas M. wrote about using [FlagEmbeddings with Optimum](https://twitter.com/IlysMoutawwakil/status/1705215192425288017) over CUDA.
-

{fastembed-0.2.6 → fastembed-0.2.7}/README.md

@@ -2,7 +2,7 @@
 
 FastEmbed is a lightweight, fast, Python library built for embedding generation. We [support popular text models](https://qdrant.github.io/fastembed/examples/Supported_Models/). Please [open a GitHub issue](https://github.com/qdrant/fastembed/issues/new) if you want us to add a new model.
 
-The default text embedding (`TextEmbedding`) model is Flag Embedding, presented in the [MTEB](https://huggingface.co/spaces/mteb/leaderboard) leaderboard. It supports "query" and "passage" prefixes for the input text. Here is an example for [Retrieval Embedding Generation](https://qdrant.github.io/fastembed/examples/Retrieval_with_FastEmbed/) and how to use [FastEmbed with Qdrant](https://qdrant.github.io/fastembed/examples/Usage_With_Qdrant/).
+The default text embedding (`TextEmbedding`) model is Flag Embedding, presented in the [MTEB](https://huggingface.co/spaces/mteb/leaderboard) leaderboard. It supports "query" and "passage" prefixes for the input text. Here is an example for [Retrieval Embedding Generation](https://qdrant.github.io/fastembed/qdrant/Retrieval_with_FastEmbed/) and how to use [FastEmbed with Qdrant](https://qdrant.github.io/fastembed/qdrant/Usage_With_Qdrant/).
 
 ## 📈 Why FastEmbed?
 
@@ -14,12 +14,18 @@ The default text embedding (`TextEmbedding`) model is Flag Embedding, presented
 
 ## 🚀 Installation
 
-To install the FastEmbed library, pip works:
+To install the FastEmbed library, pip works best. You can install it with or without GPU support:
 
 ```bash
 pip install fastembed
 ```
 
+### ⚡️ With GPU
+
+```bash
+pip install fastembed-gpu
+```
+
 ## 📖 Quickstart
 
 ```python
@@ -42,6 +48,23 @@ embeddings_list = list(embedding_model.embed(documents))
 len(embeddings_list[0]) # Vector of 384 dimensions
 ```
 
+### ⚡️ FastEmbed on a GPU
+
+FastEmbed supports running on GPU devices. It requires installation of the `fastembed-gpu` package.
+Make sure not to have the `fastembed` package installed, as it might interfere with the `fastembed-gpu` package.
+
+```bash
+pip install fastembed-gpu
+``` 
+
+```python
+from fastembed import TextEmbedding
+
+embedding_model = TextEmbedding(model_name="BAAI/bge-small-en-v1.5", providers=["CUDAExecutionProvider"])
+print("The model BAAI/bge-small-en-v1.5 is ready to use on a GPU.")
+
+```
+
 ## Usage with Qdrant
 
 Installation with Qdrant Client in Python:
@@ -50,7 +73,13 @@ Installation with Qdrant Client in Python:
 pip install qdrant-client[fastembed]
 ```
 
-You might have to use ```pip install 'qdrant-client[fastembed]'``` on zsh.
+or 
+
+```bash
+pip install qdrant-client[fastembed-gpu]
+```
+
+You might have to use quotes ```pip install 'qdrant-client[fastembed]'``` on zsh.
 
 ```python
 from qdrant_client import QdrantClient
@@ -85,8 +114,4 @@ search_result = client.query(
     query_text="This is a query document"
 )
 print(search_result)
-```
-
-#### Similar Work
-
-Ilyas M. wrote about using [FlagEmbeddings with Optimum](https://twitter.com/IlysMoutawwakil/status/1705215192425288017) over CUDA.
+```

{fastembed-0.2.6 → fastembed-0.2.7}/fastembed/__init__.py

@@ -3,5 +3,10 @@ import importlib.metadata
 from fastembed.text import TextEmbedding
 from fastembed.sparse import SparseTextEmbedding, SparseEmbedding
 
-__version__ = importlib.metadata.version("fastembed")
+try:
+    version = importlib.metadata.version("fastembed")
+except importlib.metadata.PackageNotFoundError as _:
+    version = importlib.metadata.version("fastembed-gpu")
+
+__version__ = version
 __all__ = ["TextEmbedding", "SparseTextEmbedding", "SparseEmbedding"]

fastembed-0.2.7/fastembed/common/__init__.py

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

{fastembed-0.2.6 → fastembed-0.2.7}/fastembed/common/model_management.py

@@ -11,23 +11,6 @@ from tqdm import tqdm
 from loguru import logger
 
 
-def locate_model_file(model_dir: Path, file_names: List[str]) -> Path:
-    """
-    Find model path for both TransformerJS style `onnx`  subdirectory structure and direct model weights structure used
-    by Optimum and Qdrant
-    """
-    if not model_dir.is_dir():
-        raise ValueError(f"Provided model path '{model_dir}' is not a directory.")
-
-    for file_name in file_names:
-        file_paths = [path for path in model_dir.rglob(file_name) if path.is_file()]
-
-        if file_paths:
-            return file_paths[0]
-
-    raise ValueError(f"Could not find either of {', '.join(file_names)} in {model_dir}")
-
-
 class ModelManagement:
     @classmethod
     def list_supported_models(cls) -> List[Dict[str, Any]]:
@@ -104,21 +87,36 @@ class ModelManagement:
 
     @classmethod
     def download_files_from_huggingface(
-        cls, hf_source_repo: str, cache_dir: Optional[str] = None
+        cls,
+        hf_source_repo: str,
+        cache_dir: Optional[str] = None,
+        extra_patterns: Optional[List[str]] = None,
+        **kwargs,
     ) -> str:
         """
         Downloads a model from HuggingFace Hub.
         Args:
             hf_source_repo (str): Name of the model on HuggingFace Hub, e.g. "qdrant/all-MiniLM-L6-v2-onnx".
             cache_dir (Optional[str]): The path to the cache directory.
+            extra_patterns (Optional[List[str]]): extra patterns to allow in the snapshot download, typically
+                includes the required model files.
         Returns:
             Path: The path to the model directory.
         """
+        allow_patterns = [
+            "config.json",
+            "tokenizer.json",
+            "tokenizer_config.json",
+            "special_tokens_map.json",
+        ]
+        if extra_patterns is not None:
+            allow_patterns.extend(extra_patterns)
 
         return snapshot_download(
             repo_id=hf_source_repo,
-            ignore_patterns=["model.safetensors", "pytorch_model.bin"],
+            allow_patterns=allow_patterns,
             cache_dir=cache_dir,
+            local_files_only=kwargs.get("local_files_only", False),
         )
 
     @classmethod
@@ -175,6 +173,9 @@ class ModelManagement:
 
         model_tar_gz = Path(cache_dir) / f"{fast_model_name}.tar.gz"
 
+        if model_tar_gz.exists():
+            model_tar_gz.unlink()
+
         cls.download_file_from_gcs(
             source_url,
             output_path=str(model_tar_gz),
@@ -190,7 +191,7 @@ class ModelManagement:
         return model_dir
 
     @classmethod
-    def download_model(cls, model: Dict[str, Any], cache_dir: Path) -> Path:
+    def download_model(cls, model: Dict[str, Any], cache_dir: Path, **kwargs) -> Path:
         """
         Downloads a model from HuggingFace Hub or Google Cloud Storage.
 
@@ -219,9 +220,17 @@ class ModelManagement:
         url_source = model.get("sources", {}).get("url")
 
         if hf_source:
+            extra_patterns = [model["model_file"]]
+            extra_patterns.extend(model.get("additional_files", []))
+
             try:
                 return Path(
-                    cls.download_files_from_huggingface(hf_source, cache_dir=str(cache_dir))
+                    cls.download_files_from_huggingface(
+                        hf_source,
+                        cache_dir=str(cache_dir),
+                        extra_patterns=extra_patterns,
+                        local_files_only=kwargs.get("local_files_only", False),
+                    )
                 )
             except (EnvironmentError, RepositoryNotFoundError, ValueError) as e:
                 logger.error(

{fastembed-0.2.6 → fastembed-0.2.7}/fastembed/common/onnx_model.py

@@ -1,19 +1,33 @@
 import os
 from multiprocessing import get_all_start_methods
 from pathlib import Path
-from typing import Any, Dict, Generic, Iterable, List, Optional, Tuple, Type, TypeVar, Union
+from typing import (
+    Any,
+    Dict,
+    Generic,
+    Iterable,
+    List,
+    Optional,
+    Tuple,
+    Type,
+    TypeVar,
+    Union,
+    Sequence,
+)
 
 import numpy as np
 import onnxruntime as ort
 
-from fastembed.common.model_management import locate_model_file
 from fastembed.common.models import load_tokenizer
 from fastembed.common.utils import iter_batch
 from fastembed.parallel_processor import ParallelWorkerPool, Worker
 
+
 # Holds type of the embedding result
 T = TypeVar("T")
 
+OnnxProvider = Union[str, Tuple[str, Dict[Any, Any]]]
+
 
 class OnnxModel(Generic[T]):
     @classmethod
@@ -34,11 +48,26 @@ class OnnxModel(Generic[T]):
         """
         return onnx_input
 
-    def load_onnx_model(self, model_dir: Path, threads: Optional[int], max_length: int) -> None:
-        model_path = locate_model_file(model_dir, ["model.onnx", "model_optimized.onnx"])
+    def load_onnx_model(
+        self,
+        model_dir: Path,
+        model_file: str,
+        threads: Optional[int],
+        providers: Optional[Sequence[OnnxProvider]] = None,
+    ) -> None:
+        model_path = model_dir / model_file
 
         # List of Execution Providers: https://onnxruntime.ai/docs/execution-providers
-        onnx_providers = ["CPUExecutionProvider"]
+
+        onnx_providers = ["CPUExecutionProvider"] if providers is None else list(providers)
+        available_providers = ort.get_available_providers()
+        for provider in onnx_providers:
+            # check providers available
+            provider_name = provider if isinstance(provider, str) else provider[0]
+            if provider_name not in available_providers:
+                raise ValueError(
+                    f"Provider {provider_name} is not available. Available providers: {available_providers}"
+                )
 
         so = ort.SessionOptions()
         so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
@@ -47,7 +76,7 @@ class OnnxModel(Generic[T]):
             so.intra_op_num_threads = threads
             so.inter_op_num_threads = threads
 
-        self.tokenizer = load_tokenizer(model_dir=model_dir, max_length=max_length)
+        self.tokenizer = load_tokenizer(model_dir=model_dir)
         self.model = ort.InferenceSession(
             str(model_path), providers=onnx_providers, sess_options=so
         )

{fastembed-0.2.6 → fastembed-0.2.7}/fastembed/sparse/sparse_embedding_base.py

@@ -32,6 +32,7 @@ class SparseTextEmbeddingBase(ModelManagement):
         self.model_name = model_name
         self.cache_dir = cache_dir
         self.threads = threads
+        self._local_files_only = kwargs.pop("local_files_only", False)
 
     def embed(
         self,

{fastembed-0.2.6 → fastembed-0.2.7}/fastembed/sparse/sparse_text_embedding.py

@@ -1,5 +1,6 @@
-from typing import List, Type, Dict, Any, Union, Iterable, Optional
+from typing import List, Type, Dict, Any, Union, Iterable, Optional, Sequence
 
+from fastembed.common import OnnxProvider
 from fastembed.sparse.sparse_embedding_base import SparseTextEmbeddingBase, SparseEmbedding
 from fastembed.sparse.splade_pp import SpladePP
 
@@ -42,6 +43,7 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
         model_name: str,
         cache_dir: Optional[str] = None,
         threads: Optional[int] = None,
+        providers: Optional[Sequence[OnnxProvider]] = None,
         **kwargs,
     ):
         super().__init__(model_name, cache_dir, threads, **kwargs)
@@ -49,7 +51,9 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
         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):
-                self.model = EMBEDDING_MODEL_TYPE(model_name, cache_dir, threads, **kwargs)
+                self.model = EMBEDDING_MODEL_TYPE(
+                    model_name, cache_dir, threads, providers=providers, **kwargs
+                )
                 return
 
         raise ValueError(

{fastembed-0.2.6 → fastembed-0.2.7}/fastembed/sparse/splade_pp.py

@@ -1,8 +1,8 @@
-from typing import Any, Dict, Iterable, List, Optional, Tuple, Union, Type
+from typing import Any, Dict, Iterable, List, Optional, Tuple, Union, Type, Sequence
 
 import numpy as np
 
-from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel
+from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxProvider
 from fastembed.common.utils import define_cache_dir
 from fastembed.sparse.sparse_embedding_base import SparseEmbedding, SparseTextEmbeddingBase
 
@@ -15,6 +15,7 @@ supported_splade_models = [
         "sources": {
             "hf": "Qdrant/SPLADE_PP_en_v1",
         },
+        "model_file": "model.onnx",
     },
     {
         "model": "prithivida/Splade_PP_en_v1",
@@ -24,6 +25,7 @@ supported_splade_models = [
         "sources": {
             "hf": "Qdrant/SPLADE_PP_en_v1",
         },
+        "model_file": "model.onnx",
     },
 ]
 
@@ -61,6 +63,7 @@ class SpladePP(SparseTextEmbeddingBase, OnnxModel[SparseEmbedding]):
         model_name: str,
         cache_dir: Optional[str] = None,
         threads: Optional[int] = None,
+        providers: Optional[Sequence[OnnxProvider]] = None,
         **kwargs,
     ):
         """
@@ -77,14 +80,19 @@ class SpladePP(SparseTextEmbeddingBase, OnnxModel[SparseEmbedding]):
 
         super().__init__(model_name, cache_dir, threads, **kwargs)
 
-        self.model_name = model_name
-        self._model_description = self._get_model_description(model_name)
+        model_description = self._get_model_description(model_name)
+        cache_dir = define_cache_dir(cache_dir)
 
-        self._cache_dir = define_cache_dir(cache_dir)
-        self._model_dir = self.download_model(self._model_description, self._cache_dir)
-        self._max_length = 512
+        model_dir = self.download_model(
+            model_description, cache_dir, local_files_only=self._local_files_only
+        )
 
-        self.load_onnx_model(self._model_dir, self.threads, self._max_length)
+        self.load_onnx_model(
+            model_dir=model_dir,
+            model_file=model_description["model_file"],
+            threads=threads,
+            providers=providers,
+        )
 
     def embed(
         self,
@@ -110,7 +118,7 @@ class SpladePP(SparseTextEmbeddingBase, OnnxModel[SparseEmbedding]):
         """
         yield from self._embed_documents(
             model_name=self.model_name,
-            cache_dir=str(self._cache_dir),
+            cache_dir=str(self.cache_dir),
             documents=documents,
             batch_size=batch_size,
             parallel=parallel,

{fastembed-0.2.6 → fastembed-0.2.7}/fastembed/text/e5_onnx_embedding.py

@@ -15,6 +15,8 @@ supported_multilingual_e5_models = [
             "url": "https://storage.googleapis.com/qdrant-fastembed/fast-multilingual-e5-large.tar.gz",
             "hf": "qdrant/multilingual-e5-large-onnx",
         },
+        "model_file": "model.onnx",
+        "additional_files": ["model.onnx_data"],
     },
     {
         "model": "sentence-transformers/paraphrase-multilingual-mpnet-base-v2",
@@ -24,6 +26,7 @@ supported_multilingual_e5_models = [
         "sources": {
             "hf": "xenova/paraphrase-multilingual-mpnet-base-v2",
         },
+        "model_file": "onnx/model.onnx",
     },
 ]
 

{fastembed-0.2.6 → fastembed-0.2.7}/fastembed/text/jina_onnx_embedding.py

@@ -13,6 +13,7 @@ supported_jina_models = [
         "description": "English embedding model supporting 8192 sequence length",
         "size_in_GB": 0.52,
         "sources": {"hf": "xenova/jina-embeddings-v2-base-en"},
+        "model_file": "onnx/model.onnx",
     },
     {
         "model": "jinaai/jina-embeddings-v2-small-en",
@@ -20,6 +21,7 @@ supported_jina_models = [
         "description": "English embedding model supporting 8192 sequence length",
         "size_in_GB": 0.12,
         "sources": {"hf": "xenova/jina-embeddings-v2-small-en"},
+        "model_file": "onnx/model.onnx",
     },
 ]
 

{fastembed-0.2.6 → fastembed-0.2.7}/fastembed/text/onnx_embedding.py

@@ -1,8 +1,8 @@
-from typing import Dict, Optional, Tuple, Union, Iterable, Type, List, Any
+from typing import Dict, Optional, Tuple, Union, Iterable, Type, List, Any, Sequence
 
 import numpy as np
 
-from fastembed.common.onnx_model import OnnxModel, EmbeddingWorker
+from fastembed.common.onnx_model import OnnxModel, EmbeddingWorker, OnnxProvider
 from fastembed.common.models import normalize
 from fastembed.common.utils import define_cache_dir
 from fastembed.text.text_embedding_base import TextEmbeddingBase
@@ -16,6 +16,7 @@ supported_onnx_models = [
         "sources": {
             "url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en.tar.gz",
         },
+        "model_file": "model_optimized.onnx",
     },
     {
         "model": "BAAI/bge-base-en-v1.5",
@@ -26,6 +27,7 @@ supported_onnx_models = [
             "url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en-v1.5.tar.gz",
             "hf": "qdrant/bge-base-en-v1.5-onnx-q",
         },
+        "model_file": "model_optimized.onnx",
     },
     {
         "model": "BAAI/bge-large-en-v1.5",
@@ -35,6 +37,7 @@ supported_onnx_models = [
         "sources": {
             "hf": "qdrant/bge-large-en-v1.5-onnx",
         },
+        "model_file": "model.onnx",
     },
     {
         "model": "BAAI/bge-small-en",
@@ -44,18 +47,8 @@ supported_onnx_models = [
         "sources": {
             "url": "https://storage.googleapis.com/qdrant-fastembed/BAAI-bge-small-en.tar.gz",
         },
+        "model_file": "model_optimized.onnx",
     },
-    # {
-    #     "model": "BAAI/bge-small-en",
-    #     "dim": 384,
-    #     "description": "Fast English model",
-    #     "size_in_GB": 0.2,
-    #     "hf_sources": [],
-    #     "compressed_url_sources": [
-    #         "https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-en.tar.gz",
-    #         "https://storage.googleapis.com/qdrant-fastembed/BAAI-bge-small-en.tar.gz"
-    #     ]
-    # },
     {
         "model": "BAAI/bge-small-en-v1.5",
         "dim": 384,
@@ -64,6 +57,7 @@ supported_onnx_models = [
         "sources": {
             "hf": "qdrant/bge-small-en-v1.5-onnx-q",
         },
+        "model_file": "model_optimized.onnx",
     },
     {
         "model": "BAAI/bge-small-zh-v1.5",
@@ -73,6 +67,7 @@ supported_onnx_models = [
         "sources": {
             "url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-zh-v1.5.tar.gz",
         },
+        "model_file": "model_optimized.onnx",
     },
     {
         "model": "sentence-transformers/all-MiniLM-L6-v2",
@@ -83,6 +78,7 @@ supported_onnx_models = [
             "url": "https://storage.googleapis.com/qdrant-fastembed/sentence-transformers-all-MiniLM-L6-v2.tar.gz",
             "hf": "qdrant/all-MiniLM-L6-v2-onnx",
         },
+        "model_file": "model.onnx",
     },
     {
         "model": "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
@@ -92,6 +88,7 @@ supported_onnx_models = [
         "sources": {
             "hf": "qdrant/paraphrase-multilingual-MiniLM-L12-v2-onnx-Q",
         },
+        "model_file": "model_optimized.onnx",
     },
     {
         "model": "nomic-ai/nomic-embed-text-v1",
@@ -101,6 +98,7 @@ supported_onnx_models = [
         "sources": {
             "hf": "nomic-ai/nomic-embed-text-v1",
         },
+        "model_file": "onnx/model.onnx",
     },
     {
         "model": "nomic-ai/nomic-embed-text-v1.5",
@@ -110,6 +108,17 @@ supported_onnx_models = [
         "sources": {
             "hf": "nomic-ai/nomic-embed-text-v1.5",
         },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "nomic-ai/nomic-embed-text-v1.5-Q",
+        "dim": 768,
+        "description": "Quantized 8192 context length english model",
+        "size_in_GB": 0.13,
+        "sources": {
+            "hf": "nomic-ai/nomic-embed-text-v1.5",
+        },
+        "model_file": "onnx/model_quantized.onnx",
     },
     {
         "model": "thenlper/gte-large",
@@ -119,20 +128,8 @@ supported_onnx_models = [
         "sources": {
             "hf": "qdrant/gte-large-onnx",
         },
+        "model_file": "model.onnx",
     },
-    # {
-    #     "model": "sentence-transformers/all-MiniLM-L6-v2",
-    #     "dim": 384,
-    #     "description": "Sentence Transformer model, MiniLM-L6-v2",
-    #     "size_in_GB": 0.09,
-    #     "hf_sources": [
-    #         "qdrant/all-MiniLM-L6-v2-onnx"
-    #     ],
-    #     "compressed_url_sources": [
-    #         "https://storage.googleapis.com/qdrant-fastembed/fast-all-MiniLM-L6-v2.tar.gz",
-    #         "https://storage.googleapis.com/qdrant-fastembed/sentence-transformers-all-MiniLM-L6-v2.tar.gz"
-    #     ]
-    # }
     {
         "model": "mixedbread-ai/mxbai-embed-large-v1",
         "dim": 1024,
@@ -141,6 +138,57 @@ supported_onnx_models = [
         "sources": {
             "hf": "mixedbread-ai/mxbai-embed-large-v1",
         },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "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.",
+        "size_in_GB": 0.09,
+        "sources": {
+            "hf": "snowflake/snowflake-arctic-embed-xs",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "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.",
+        "size_in_GB": 0.13,
+        "sources": {
+            "hf": "snowflake/snowflake-arctic-embed-s",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "snowflake/snowflake-arctic-embed-m",
+        "dim": 768,
+        "description": "Based on intfloat/e5-base-unsupervised model, provides the best retrieval without slowing down inference.",
+        "size_in_GB": 0.43,
+        "sources": {
+            "hf": "Snowflake/snowflake-arctic-embed-m",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "snowflake/snowflake-arctic-embed-m-long",
+        "dim": 768,
+        "description": "Based on nomic-ai/nomic-embed-text-v1-unsupervised model, 8192 context-length model",
+        "size_in_GB": 0.54,
+        "sources": {
+            "hf": "snowflake/snowflake-arctic-embed-m-long",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "snowflake/snowflake-arctic-embed-l",
+        "dim": 1024,
+        "description": "Based on intfloat/e5-large-unsupervised, large model for most accurate retrieval.",
+        "size_in_GB": 1.02,
+        "sources": {
+            "hf": "snowflake/snowflake-arctic-embed-l",
+        },
+        "model_file": "onnx/model.onnx",
     },
 ]
 
@@ -163,6 +211,7 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxModel[np.ndarray]):
         model_name: str = "BAAI/bge-small-en-v1.5",
         cache_dir: Optional[str] = None,
         threads: Optional[int] = None,
+        providers: Optional[Sequence[OnnxProvider]] = None,
         **kwargs,
     ):
         """
@@ -179,14 +228,18 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxModel[np.ndarray]):
 
         super().__init__(model_name, cache_dir, threads, **kwargs)
 
-        self.model_name = model_name
-        self._model_description = self._get_model_description(model_name)
-
-        self._cache_dir = define_cache_dir(cache_dir)
-        self._model_dir = self.download_model(self._model_description, self._cache_dir)
-        self._max_length = 512
+        model_description = self._get_model_description(model_name)
+        cache_dir = define_cache_dir(cache_dir)
+        model_dir = self.download_model(
+            model_description, cache_dir, local_files_only=self._local_files_only
+        )
 
-        self.load_onnx_model(self._model_dir, self.threads, self._max_length)
+        self.load_onnx_model(
+            model_dir=model_dir,
+            model_file=model_description["model_file"],
+            threads=threads,
+            providers=providers,
+        )
 
     def embed(
         self,
@@ -212,7 +265,7 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxModel[np.ndarray]):
         """
         yield from self._embed_documents(
             model_name=self.model_name,
-            cache_dir=str(self._cache_dir),
+            cache_dir=str(self.cache_dir),
             documents=documents,
             batch_size=batch_size,
             parallel=parallel,

{fastembed-0.2.6 → fastembed-0.2.7}/fastembed/text/text_embedding.py

@@ -1,7 +1,8 @@
-from typing import Any, Dict, Iterable, List, Optional, Type, Union
+from typing import Any, Dict, Iterable, List, Optional, Type, Union, Sequence
 
 import numpy as np
 
+from fastembed.common import OnnxProvider
 from fastembed.text.e5_onnx_embedding import E5OnnxEmbedding
 from fastembed.text.jina_onnx_embedding import JinaOnnxEmbedding
 from fastembed.text.onnx_embedding import OnnxTextEmbedding
@@ -49,6 +50,7 @@ class TextEmbedding(TextEmbeddingBase):
         model_name: str = "BAAI/bge-small-en-v1.5",
         cache_dir: Optional[str] = None,
         threads: Optional[int] = None,
+        providers: Optional[Sequence[OnnxProvider]] = None,
         **kwargs,
     ):
         super().__init__(model_name, cache_dir, threads, **kwargs)
@@ -56,7 +58,9 @@ class TextEmbedding(TextEmbeddingBase):
         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):
-                self.model = EMBEDDING_MODEL_TYPE(model_name, cache_dir, threads, **kwargs)
+                self.model = EMBEDDING_MODEL_TYPE(
+                    model_name, cache_dir, threads, providers=providers, **kwargs
+                )
                 return
 
         raise ValueError(

{fastembed-0.2.6 → fastembed-0.2.7}/fastembed/text/text_embedding_base.py

@@ -16,6 +16,7 @@ class TextEmbeddingBase(ModelManagement):
         self.model_name = model_name
         self.cache_dir = cache_dir
         self.threads = threads
+        self._local_files_only = kwargs.pop("local_files_only", False)
 
     def embed(
         self,

{fastembed-0.2.6 → fastembed-0.2.7}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "fastembed"
-version = "0.2.6"
+version = "0.2.7"
 description = "Fast, light, accurate library built for retrieval embedding generation"
 authors = ["NirantK <nirant.bits@gmail.com>"]
 license = "Apache License"
@@ -16,7 +16,7 @@ onnx = "^1.15.0"
 onnxruntime = "^1.17.0"
 tqdm = "^4.66"
 requests = "^2.31"
-tokenizers = "^0.15.1"
+tokenizers = "^0.15"
 huggingface-hub = "^0.20"
 loguru = "^0.7.2"
 numpy = [