chore: updated model file path

fastembed/common/model_management.py

@@ -1,31 +1,4 @@
-import os
-import shutil
-import tarfile
-from pathlib import Path
-from typing import List, Optional, Dict, Any
-
-import requests
-from huggingface_hub import snapshot_download
-from huggingface_hub.utils import RepositoryNotFoundError
-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}")
+from typing import List, Dict, Any
 
 
 class ModelManagement:
@@ -57,189 +30,3 @@ class ModelManagement:
                 return model
 
         raise ValueError(f"Model {model_name} is not supported in {cls.__name__}.")
-
-    @classmethod
-    def download_file_from_gcs(cls, url: str, output_path: str, show_progress: bool = True) -> str:
-        """
-        Downloads a file from Google Cloud Storage.
-
-        Args:
-            url (str): The URL to download the file from.
-            output_path (str): The path to save the downloaded file to.
-            show_progress (bool, optional): Whether to show a progress bar. Defaults to True.
-
-        Returns:
-            str: The path to the downloaded file.
-        """
-
-        if os.path.exists(output_path):
-            return output_path
-        response = requests.get(url, stream=True)
-
-        # Handle HTTP errors
-        if response.status_code == 403:
-            raise PermissionError(
-                "Authentication Error: You do not have permission to access this resource. "
-                "Please check your credentials."
-            )
-
-        # Get the total size of the file
-        total_size_in_bytes = int(response.headers.get("content-length", 0))
-
-        # Warn if the total size is zero
-        if total_size_in_bytes == 0:
-            print(f"Warning: Content-length header is missing or zero in the response from {url}.")
-
-        show_progress = total_size_in_bytes and show_progress
-
-        with tqdm(
-            total=total_size_in_bytes, unit="iB", unit_scale=True, disable=not show_progress
-        ) as progress_bar:
-            with open(output_path, "wb") as file:
-                for chunk in response.iter_content(chunk_size=1024):
-                    if chunk:  # Filter out keep-alive new chunks
-                        progress_bar.update(len(chunk))
-                        file.write(chunk)
-        return output_path
-
-    @classmethod
-    def download_files_from_huggingface(
-        cls, hf_source_repo: str, cache_dir: Optional[str] = None
-    ) -> 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.
-        Returns:
-            Path: The path to the model directory.
-        """
-
-        return snapshot_download(
-            repo_id=hf_source_repo,
-            allow_patterns=[
-                "*.onnx",
-                "*.onnx_data",
-                "config.json",
-                "tokenizer.json",
-                "tokenizer_config.json",
-                "special_tokens_map.json",
-            ],
-            cache_dir=cache_dir,
-        )
-
-    @classmethod
-    def decompress_to_cache(cls, targz_path: str, cache_dir: str):
-        """
-        Decompresses a .tar.gz file to a cache directory.
-
-        Args:
-            targz_path (str): Path to the .tar.gz file.
-            cache_dir (str): Path to the cache directory.
-
-        Returns:
-            cache_dir (str): Path to the cache directory.
-        """
-        # Check if targz_path exists and is a file
-        if not os.path.isfile(targz_path):
-            raise ValueError(f"{targz_path} does not exist or is not a file.")
-
-        # Check if targz_path is a .tar.gz file
-        if not targz_path.endswith(".tar.gz"):
-            raise ValueError(f"{targz_path} is not a .tar.gz file.")
-
-        try:
-            # 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)
-        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)
-            # and raise the error again
-            if "tmp" in cache_dir:
-                shutil.rmtree(cache_dir)
-            raise ValueError(f"An error occurred while decompressing {targz_path}: {e}")
-
-        return cache_dir
-
-    @classmethod
-    def retrieve_model_gcs(cls, model_name: str, source_url: str, cache_dir: str) -> Path:
-        fast_model_name = f"fast-{model_name.split('/')[-1]}"
-
-        cache_tmp_dir = Path(cache_dir) / "tmp"
-        model_tmp_dir = cache_tmp_dir / fast_model_name
-        model_dir = Path(cache_dir) / fast_model_name
-
-        # check if the model_dir and the model files are both present for macOS
-        if model_dir.exists() and len(list(model_dir.glob("*"))) > 0:
-            return model_dir
-
-        if model_tmp_dir.exists():
-            shutil.rmtree(model_tmp_dir)
-
-        cache_tmp_dir.mkdir(parents=True, exist_ok=True)
-
-        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),
-        )
-
-        cls.decompress_to_cache(targz_path=str(model_tar_gz), cache_dir=str(cache_tmp_dir))
-        assert model_tmp_dir.exists(), f"Could not find {model_tmp_dir} in {cache_tmp_dir}"
-
-        model_tar_gz.unlink()
-        # Rename from tmp to final name is atomic
-        model_tmp_dir.rename(model_dir)
-
-        return model_dir
-
-    @classmethod
-    def download_model(cls, model: Dict[str, Any], cache_dir: Path) -> Path:
-        """
-        Downloads a model from HuggingFace Hub or Google Cloud Storage.
-
-        Args:
-            model (Dict[str, Any]): The model description.
-                Example:
-                ```
-                {
-                    "model": "BAAI/bge-base-en-v1.5",
-                    "dim": 768,
-                    "description": "Base English model, v1.5",
-                    "size_in_GB": 0.44,
-                    "sources": {
-                        "url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en-v1.5.tar.gz",
-                        "hf": "qdrant/bge-base-en-v1.5-onnx-q",
-                    }
-                }
-                ```
-            cache_dir (str): The path to the cache directory.
-
-        Returns:
-            Path: The path to the downloaded model directory.
-        """
-
-        hf_source = model.get("sources", {}).get("hf")
-        url_source = model.get("sources", {}).get("url")
-
-        if hf_source:
-            try:
-                return Path(
-                    cls.download_files_from_huggingface(hf_source, cache_dir=str(cache_dir))
-                )
-            except (EnvironmentError, RepositoryNotFoundError, ValueError) as e:
-                logger.error(
-                    f"Could not download model from HuggingFace: {e}"
-                    "Falling back to other sources."
-                )
-
-        if url_source:
-            return cls.retrieve_model_gcs(model["model"], url_source, str(cache_dir))
-
-        raise ValueError(f"Could not download model {model['model']} from any source.")

fastembed/common/models.py

@@ -3,24 +3,25 @@ from pathlib import Path
 
 import numpy as np
 from tokenizers import Tokenizer, AddedToken
+from huggingface_hub import hf_hub_download
 
 
-def load_tokenizer(model_dir: Path, max_length: int = 512) -> Tokenizer:
-    config_path = model_dir / "config.json"
-    if not config_path.exists():
-        raise ValueError(f"Could not find config.json in {model_dir}")
+def load_tokenizer(repo_id: str, cache_dir: Path, max_length: int = 512) -> Tokenizer:
+    config_path = hf_hub_download(
+        repo_id=repo_id, filename="config.json", cache_dir=str(cache_dir)
+    )
 
-    tokenizer_path = model_dir / "tokenizer.json"
-    if not tokenizer_path.exists():
-        raise ValueError(f"Could not find tokenizer.json in {model_dir}")
+    tokenizer_path = hf_hub_download(
+        repo_id=repo_id, filename="tokenizer.json", cache_dir=str(cache_dir)
+    )
 
-    tokenizer_config_path = model_dir / "tokenizer_config.json"
-    if not tokenizer_config_path.exists():
-        raise ValueError(f"Could not find tokenizer_config.json in {model_dir}")
+    tokenizer_config_path = hf_hub_download(
+        repo_id=repo_id, filename="tokenizer_config.json", cache_dir=str(cache_dir)
+    )
 
-    tokens_map_path = model_dir / "special_tokens_map.json"
-    if not tokens_map_path.exists():
-        raise ValueError(f"Could not find special_tokens_map.json in {model_dir}")
+    tokens_map_path = hf_hub_download(
+        repo_id=repo_id, filename="special_tokens_map.json", cache_dir=str(cache_dir)
+    )
 
     with open(str(config_path)) as config_file:
         config = json.load(config_file)

fastembed/common/onnx_model.py

@@ -6,11 +6,12 @@ from typing import Any, Dict, Generic, Iterable, List, Optional, Tuple, Type, Ty
 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
 
+from huggingface_hub import hf_hub_download
+
 # Holds type of the embedding result
 T = TypeVar("T")
 
@@ -34,8 +35,22 @@ 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_description: dict, threads: Optional[int], cache_dir: Path
+    ) -> None:
+        repo_id = model_description["sources"]["hf"]
+        model_file = model_description.get("model_file", "model.onnx")
+
+        # Some models require additional repo files.
+        # For eg: intfloat/multilingual-e5-large requires the model.onnx_data file.
+        # These can be specified within the "additional_files" option when describing the model properties
+        if additional_files := model_description.get("additional_files"):
+            for file in additional_files:
+                hf_hub_download(repo_id=repo_id, filename=file, cache_dir=str(cache_dir))
+
+        model_path = hf_hub_download(
+            repo_id=repo_id, filename=model_file, cache_dir=str(cache_dir)
+        )
 
         # List of Execution Providers: https://onnxruntime.ai/docs/execution-providers
         onnx_providers = ["CPUExecutionProvider"]
@@ -50,7 +65,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(repo_id, cache_dir)
         self.model = ort.InferenceSession(
             str(model_path), providers=onnx_providers, sess_options=so
         )

fastembed/sparse/splade_pp.py

@@ -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",
     },
 ]
 
@@ -77,14 +79,11 @@ 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)
-
-        self._cache_dir = define_cache_dir(cache_dir)
-        self._model_dir = self.download_model(self._model_description, self._cache_dir)
-        self._max_length = 512
-
-        self.load_onnx_model(self._model_dir, self.threads, self._max_length)
+        self.load_onnx_model(
+            self._get_model_description(model_name),
+            threads,
+            define_cache_dir(cache_dir),
+        )
 
     def embed(
         self,
@@ -110,7 +109,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/text/e5_onnx_embedding.py

@@ -12,9 +12,10 @@ supported_multilingual_e5_models = [
         "description": "Multilingual model, e5-large. Recommend using this model for non-English languages",
         "size_in_GB": 2.24,
         "sources": {
-            "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 +25,7 @@ supported_multilingual_e5_models = [
         "sources": {
             "hf": "xenova/paraphrase-multilingual-mpnet-base-v2",
         },
+        "model_file": "onnx/model.onnx",
     },
 ]
 

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/text/onnx_embedding.py

@@ -14,8 +14,9 @@ supported_onnx_models = [
         "description": "Base English model",
         "size_in_GB": 0.42,
         "sources": {
-            "url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en.tar.gz",
+            "hf": "yashvardhan7/bge-base-en-onnx",
         },
+        "model_file": "onnx/model.onnx",
     },
     {
         "model": "BAAI/bge-base-en-v1.5",
@@ -23,9 +24,9 @@ supported_onnx_models = [
         "description": "Base English model, v1.5",
         "size_in_GB": 0.21,
         "sources": {
-            "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 +36,7 @@ supported_onnx_models = [
         "sources": {
             "hf": "qdrant/bge-large-en-v1.5-onnx",
         },
+        "model_file": "model.onnx",
     },
     {
         "model": "BAAI/bge-small-en",
@@ -42,20 +44,10 @@ supported_onnx_models = [
         "description": "Fast English model",
         "size_in_GB": 0.13,
         "sources": {
-            "url": "https://storage.googleapis.com/qdrant-fastembed/BAAI-bge-small-en.tar.gz",
+            "hf": "ggrn/bge-small-en",
         },
+        "model_file": "onnx/model.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 +56,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",
@@ -71,8 +64,9 @@ supported_onnx_models = [
         "description": "Fast and recommended Chinese model",
         "size_in_GB": 0.09,
         "sources": {
-            "url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-zh-v1.5.tar.gz",
+            "hf": "Xenova/bge-small-zh-v1.5",
         },
+        "model_file": "onnx/model.onnx",
     },
     {
         "model": "sentence-transformers/all-MiniLM-L6-v2",
@@ -80,9 +74,9 @@ supported_onnx_models = [
         "description": "Sentence Transformer model, MiniLM-L6-v2",
         "size_in_GB": 0.09,
         "sources": {
-            "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 +86,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 +96,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 +106,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 +126,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 +136,7 @@ supported_onnx_models = [
         "sources": {
             "hf": "mixedbread-ai/mxbai-embed-large-v1",
         },
+        "model_file": "onnx/model.onnx",
     },
 ]
 
@@ -179,14 +175,11 @@ 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
-
-        self.load_onnx_model(self._model_dir, self.threads, self._max_length)
+        self.load_onnx_model(
+            self._get_model_description(model_name),
+            threads,
+            define_cache_dir(cache_dir),
+        )
 
     def embed(
         self,
@@ -212,7 +205,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,

poetry.lock

@@ -3446,4 +3446,4 @@ testing = ["big-O", "jaraco.functools", "jaraco.itertools", "more-itertools", "p
 [metadata]
 lock-version = "2.0"
 python-versions = ">=3.8.0,<3.13"
-content-hash = "b71c2e05a840cee62f4d1a11c0377e93d735f473e13f684246c021e01c5368c3"
+content-hash = "ace5b40bd629af8ec1a369f1cfb507db7bb6b18c1387357e0d029d9b50a66335"

pyproject.toml

@@ -15,7 +15,6 @@ python = ">=3.8.0,<3.13"
 onnx = "^1.15.0"
 onnxruntime = "^1.17.0"
 tqdm = "^4.66"
-requests = "^2.31"
 tokenizers = "^0.15.1"
 huggingface-hub = "^0.20"
 loguru = "^0.7.2"

tests/test_text_onnx_embeddings.py

@@ -28,6 +28,9 @@ CANONICAL_VECTOR_VALUES = {
     ),
     "thenlper/gte-large": np.array([-0.01920587, 0.00113156, -0.00708992, -0.00632304, -0.04025577]),
     "mixedbread-ai/mxbai-embed-large-v1": np.array([0.02295546, 0.03196154, 0.016512, -0.04031524, -0.0219634]),
+    "nomic-ai/nomic-embed-text-v1.5-Q": np.array(
+        [-0.01554983, 0.0129992, -0.17909265, -0.01062993, 0.00512859]
+    ),
 }