new: update dev workflow

fix links (#215 )
new: remove slurm environment variables (#213 )
2026-09-22 22:17:49 -05:00 · 2024-05-02 20:34:47 +02:00 · 2024-04-27 22:24:29 +05:30 · 2024-04-26 21:45:02 +02:00 · 2024-04-26 17:27:16 +02:00 · 2024-04-26 10:28:37 +05:30
13 changed files with 449 additions and 183 deletions
@@ -1,8 +1,9 @@
 name: Tests
+run-name: Tests (dev)

 on:
  push:
-    branches: [ master, main ]
+    branches: [ dev ]
  pull_request:

 env:
@@ -10,7 +11,6 @@ env:

 jobs:
  test:
-
    strategy:
      matrix:
        python-version:
@@ -29,9 +29,9 @@ jobs:
    name: Python ${{ matrix.python-version }} on ${{ matrix.os }} test

    steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v3
      - name: Set up Python
-        uses: actions/setup-python@v2
+        uses: actions/setup-python@v5
        with:
          python-version: ${{ matrix.python-version }}
      - name: Install dependencies
@@ -39,8 +39,14 @@ jobs:
          python -m pip install poetry
          poetry config virtualenvs.create false
          poetry install --no-interaction --no-ansi --without docs
-      - name: Run tests
-        run: |
-          export IS_UBUNTU_CI=$(test "${{ matrix.os }}" = "ubuntu-latest" && echo "true" || echo "false")
-          pytest
-        shell: bash
+
+      - name: Install Test Dependencies
+        run: pip install pytest pytest-md pytest-emoji
+
+      - name: Run pytest
+        uses: pavelzw/pytest-action@v2
+        with:
+            verbose: true
+            emoji: true
+            job-summary: true
+            report-title: 'FastEmbed Test Report'
@@ -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?

@@ -17,7 +17,7 @@
  },
  {
   "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
@@ -35,7 +35,7 @@
  },
  {
   "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
@@ -89,90 +89,125 @@
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
+       "      <td>snowflake/snowflake-arctic-embed-xs</td>\n",
+       "      <td>384</td>\n",
+       "      <td>Based on all-MiniLM-L6-v2 model with only 22m ...</td>\n",
+       "      <td>0.090</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
       "      <td>jinaai/jina-embeddings-v2-small-en</td>\n",
       "      <td>512</td>\n",
       "      <td>English embedding model supporting 8192 sequen...</td>\n",
       "      <td>0.120</td>\n",
       "    </tr>\n",
       "    <tr>\n",
-       "      <th>4</th>\n",
+       "      <th>5</th>\n",
+       "      <td>snowflake/snowflake-arctic-embed-s</td>\n",
+       "      <td>384</td>\n",
+       "      <td>Based on infloat/e5-small-unsupervised, does n...</td>\n",
+       "      <td>0.130</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6</th>\n",
       "      <td>BAAI/bge-small-en</td>\n",
       "      <td>384</td>\n",
       "      <td>Fast English model</td>\n",
       "      <td>0.130</td>\n",
       "    </tr>\n",
       "    <tr>\n",
-       "      <th>5</th>\n",
+       "      <th>7</th>\n",
       "      <td>BAAI/bge-base-en-v1.5</td>\n",
       "      <td>768</td>\n",
       "      <td>Base English model, v1.5</td>\n",
       "      <td>0.210</td>\n",
       "    </tr>\n",
       "    <tr>\n",
-       "      <th>6</th>\n",
+       "      <th>8</th>\n",
       "      <td>sentence-transformers/paraphrase-multilingual-...</td>\n",
       "      <td>384</td>\n",
       "      <td>Sentence Transformer model, paraphrase-multili...</td>\n",
       "      <td>0.220</td>\n",
       "    </tr>\n",
       "    <tr>\n",
-       "      <th>7</th>\n",
+       "      <th>9</th>\n",
       "      <td>BAAI/bge-base-en</td>\n",
       "      <td>768</td>\n",
       "      <td>Base English model</td>\n",
       "      <td>0.420</td>\n",
       "    </tr>\n",
       "    <tr>\n",
-       "      <th>8</th>\n",
-       "      <td>nomic-ai/nomic-embed-text-v1</td>\n",
-       "      <td>768</td>\n",
-       "      <td>8192 context length english model</td>\n",
-       "      <td>0.520</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>9</th>\n",
-       "      <td>nomic-ai/nomic-embed-text-v1.5</td>\n",
-       "      <td>768</td>\n",
-       "      <td>8192 context length english model</td>\n",
-       "      <td>0.520</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
       "      <th>10</th>\n",
+       "      <td>snowflake/snowflake-arctic-embed-m</td>\n",
+       "      <td>768</td>\n",
+       "      <td>Based on intfloat/e5-base-unsupervised model, ...</td>\n",
+       "      <td>0.430</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>11</th>\n",
       "      <td>jinaai/jina-embeddings-v2-base-en</td>\n",
       "      <td>768</td>\n",
       "      <td>English embedding model supporting 8192 sequen...</td>\n",
       "      <td>0.520</td>\n",
       "    </tr>\n",
       "    <tr>\n",
-       "      <th>11</th>\n",
+       "      <th>12</th>\n",
+       "      <td>nomic-ai/nomic-embed-text-v1</td>\n",
+       "      <td>768</td>\n",
+       "      <td>8192 context length english model</td>\n",
+       "      <td>0.520</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>13</th>\n",
+       "      <td>nomic-ai/nomic-embed-text-v1.5</td>\n",
+       "      <td>768</td>\n",
+       "      <td>8192 context length english model</td>\n",
+       "      <td>0.520</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>14</th>\n",
+       "      <td>snowflake/snowflake-arctic-embed-m-long</td>\n",
+       "      <td>768</td>\n",
+       "      <td>Based on nomic-ai/nomic-embed-text-v1-unsuperv...</td>\n",
+       "      <td>0.540</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>15</th>\n",
       "      <td>mixedbread-ai/mxbai-embed-large-v1</td>\n",
       "      <td>1024</td>\n",
       "      <td>MixedBread Base sentence embedding model, does...</td>\n",
       "      <td>0.640</td>\n",
       "    </tr>\n",
       "    <tr>\n",
-       "      <th>12</th>\n",
+       "      <th>16</th>\n",
       "      <td>sentence-transformers/paraphrase-multilingual-...</td>\n",
       "      <td>768</td>\n",
       "      <td>Sentence-transformers model for tasks like clu...</td>\n",
       "      <td>1.000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
-       "      <th>13</th>\n",
+       "      <th>17</th>\n",
+       "      <td>snowflake/snowflake-arctic-embed-l</td>\n",
+       "      <td>1024</td>\n",
+       "      <td>Based on intfloat/e5-large-unsupervised, large...</td>\n",
+       "      <td>1.020</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>18</th>\n",
       "      <td>BAAI/bge-large-en-v1.5</td>\n",
       "      <td>1024</td>\n",
       "      <td>Large English model, v1.5</td>\n",
       "      <td>1.200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
-       "      <th>14</th>\n",
+       "      <th>19</th>\n",
       "      <td>thenlper/gte-large</td>\n",
       "      <td>1024</td>\n",
       "      <td>Large general text embeddings model</td>\n",
       "      <td>1.200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
-       "      <th>15</th>\n",
+       "      <th>20</th>\n",
       "      <td>intfloat/multilingual-e5-large</td>\n",
       "      <td>1024</td>\n",
       "      <td>Multilingual model, e5-large. Recommend using ...</td>\n",
@@ -187,40 +222,50 @@
       "0                              BAAI/bge-small-en-v1.5   384   \n",
       "1                              BAAI/bge-small-zh-v1.5   512   \n",
       "2              sentence-transformers/all-MiniLM-L6-v2   384   \n",
-       "3                  jinaai/jina-embeddings-v2-small-en   512   \n",
-       "4                                   BAAI/bge-small-en   384   \n",
-       "5                               BAAI/bge-base-en-v1.5   768   \n",
-       "6   sentence-transformers/paraphrase-multilingual-...   384   \n",
-       "7                                    BAAI/bge-base-en   768   \n",
-       "8                        nomic-ai/nomic-embed-text-v1   768   \n",
-       "9                      nomic-ai/nomic-embed-text-v1.5   768   \n",
-       "10                  jinaai/jina-embeddings-v2-base-en   768   \n",
-       "11                 mixedbread-ai/mxbai-embed-large-v1  1024   \n",
-       "12  sentence-transformers/paraphrase-multilingual-...   768   \n",
-       "13                             BAAI/bge-large-en-v1.5  1024   \n",
-       "14                                 thenlper/gte-large  1024   \n",
-       "15                     intfloat/multilingual-e5-large  1024   \n",
+       "3                 snowflake/snowflake-arctic-embed-xs   384   \n",
+       "4                  jinaai/jina-embeddings-v2-small-en   512   \n",
+       "5                  snowflake/snowflake-arctic-embed-s   384   \n",
+       "6                                   BAAI/bge-small-en   384   \n",
+       "7                               BAAI/bge-base-en-v1.5   768   \n",
+       "8   sentence-transformers/paraphrase-multilingual-...   384   \n",
+       "9                                    BAAI/bge-base-en   768   \n",
+       "10                 snowflake/snowflake-arctic-embed-m   768   \n",
+       "11                  jinaai/jina-embeddings-v2-base-en   768   \n",
+       "12                       nomic-ai/nomic-embed-text-v1   768   \n",
+       "13                     nomic-ai/nomic-embed-text-v1.5   768   \n",
+       "14            snowflake/snowflake-arctic-embed-m-long   768   \n",
+       "15                 mixedbread-ai/mxbai-embed-large-v1  1024   \n",
+       "16  sentence-transformers/paraphrase-multilingual-...   768   \n",
+       "17                 snowflake/snowflake-arctic-embed-l  1024   \n",
+       "18                             BAAI/bge-large-en-v1.5  1024   \n",
+       "19                                 thenlper/gte-large  1024   \n",
+       "20                     intfloat/multilingual-e5-large  1024   \n",
       "\n",
       "                                          description  size_in_GB  \n",
       "0                      Fast and Default English model       0.067  \n",
       "1                  Fast and recommended Chinese model       0.090  \n",
       "2            Sentence Transformer model, MiniLM-L6-v2       0.090  \n",
-       "3   English embedding model supporting 8192 sequen...       0.120  \n",
-       "4                                  Fast English model       0.130  \n",
-       "5                            Base English model, v1.5       0.210  \n",
-       "6   Sentence Transformer model, paraphrase-multili...       0.220  \n",
-       "7                                  Base English model       0.420  \n",
-       "8                   8192 context length english model       0.520  \n",
-       "9                   8192 context length english model       0.520  \n",
-       "10  English embedding model supporting 8192 sequen...       0.520  \n",
-       "11  MixedBread Base sentence embedding model, does...       0.640  \n",
-       "12  Sentence-transformers model for tasks like clu...       1.000  \n",
-       "13                          Large English model, v1.5       1.200  \n",
-       "14                Large general text embeddings model       1.200  \n",
-       "15  Multilingual model, e5-large. Recommend using ...       2.240  "
+       "3   Based on all-MiniLM-L6-v2 model with only 22m ...       0.090  \n",
+       "4   English embedding model supporting 8192 sequen...       0.120  \n",
+       "5   Based on infloat/e5-small-unsupervised, does n...       0.130  \n",
+       "6                                  Fast English model       0.130  \n",
+       "7                            Base English model, v1.5       0.210  \n",
+       "8   Sentence Transformer model, paraphrase-multili...       0.220  \n",
+       "9                                  Base English model       0.420  \n",
+       "10  Based on intfloat/e5-base-unsupervised model, ...       0.430  \n",
+       "11  English embedding model supporting 8192 sequen...       0.520  \n",
+       "12                  8192 context length english model       0.520  \n",
+       "13                  8192 context length english model       0.520  \n",
+       "14  Based on nomic-ai/nomic-embed-text-v1-unsuperv...       0.540  \n",
+       "15  MixedBread Base sentence embedding model, does...       0.640  \n",
+       "16  Sentence-transformers model for tasks like clu...       1.000  \n",
+       "17  Based on intfloat/e5-large-unsupervised, large...       1.020  \n",
+       "18                          Large English model, v1.5       1.200  \n",
+       "19                Large general text embeddings model       1.200  \n",
+       "20  Multilingual model, e5-large. Recommend using ...       2.240  "
      ]
     },
-     "execution_count": 6,
+     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
@@ -1,28 +1,16 @@
+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}")
-
-
 class ModelManagement:
    @classmethod
    def list_supported_models(cls) -> List[Dict[str, Any]]:
@@ -53,32 +41,153 @@ class ModelManagement:

        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
+        cls,
+        hf_source_repo: str,
+        cache_dir: Optional[str] = None,
+        extra_patterns: Optional[List[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.
+            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,
-            allow_patterns=[
-                "*.onnx",
-                "*.onnx_data",
-                "config.json",
-                "tokenizer.json",
-                "tokenizer_config.json",
-                "special_tokens_map.json",
-            ],
+            allow_patterns=allow_patterns,
            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:
        """
@@ -94,6 +203,7 @@ class ModelManagement:
                    "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",
                    }
                }
@@ -105,11 +215,17 @@ class ModelManagement:
        """

        hf_source = model.get("sources", {}).get("hf")
+        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
+                    )
                )
            except (EnvironmentError, RepositoryNotFoundError, ValueError) as e:
                logger.error(
@@ -117,4 +233,7 @@ class ModelManagement:
                    "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.")
@@ -6,11 +6,11 @@ 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

+
 # Holds type of the embedding result
 T = TypeVar("T")

@@ -34,23 +34,25 @@ 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],
+    ) -> None:
+        model_path = model_dir / model_file

        # List of Execution Providers: https://onnxruntime.ai/docs/execution-providers
        onnx_providers = ["CPUExecutionProvider"]

        so = ort.SessionOptions()
-        if os.getenv("SLURM_JOB_ID") is not None:
-            so.intra_op_num_threads = int(os.getenv("SLURM_CPUS_ON_NODE"))
-            so.inter_op_num_threads = int(os.getenv("SLURM_CPUS_ON_NODE"))
        so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL

        if threads is not None:
            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
        )
@@ -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,16 @@ 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)

-        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,
+        )

    def embed(
        self,
@@ -110,7 +114,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,
@@ -12,8 +12,11 @@ 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",
@@ -23,6 +26,7 @@ supported_multilingual_e5_models = [
        "sources": {
            "hf": "xenova/paraphrase-multilingual-mpnet-base-v2",
        },
+        "model_file": "onnx/model.onnx",
    },
 ]

@@ -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",
    },
 ]

@@ -14,8 +14,9 @@ supported_onnx_models = [
        "description": "Base English model",
        "size_in_GB": 0.42,
        "sources": {
-            "hf": "yashvardhan7/bge-base-en-onnx",
+            "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",
@@ -23,8 +24,10 @@ 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",
@@ -34,6 +37,7 @@ supported_onnx_models = [
        "sources": {
            "hf": "qdrant/bge-large-en-v1.5-onnx",
        },
+        "model_file": "model.onnx",
    },
    {
        "model": "BAAI/bge-small-en",
@@ -41,8 +45,9 @@ supported_onnx_models = [
        "description": "Fast English model",
        "size_in_GB": 0.13,
        "sources": {
-            "hf": "ggrn/bge-small-en",
+            "url": "https://storage.googleapis.com/qdrant-fastembed/BAAI-bge-small-en.tar.gz",
        },
+        "model_file": "model_optimized.onnx",
    },
    {
        "model": "BAAI/bge-small-en-v1.5",
@@ -52,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",
@@ -59,8 +65,9 @@ supported_onnx_models = [
        "description": "Fast and recommended Chinese model",
        "size_in_GB": 0.09,
        "sources": {
-            "hf": "Xenova/bge-small-zh-v1.5",
+            "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",
@@ -68,8 +75,10 @@ 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",
@@ -79,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",
@@ -88,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",
@@ -97,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",
@@ -106,6 +128,7 @@ supported_onnx_models = [
        "sources": {
            "hf": "qdrant/gte-large-onnx",
        },
+        "model_file": "model.onnx",
    },
    {
        "model": "mixedbread-ai/mxbai-embed-large-v1",
@@ -115,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",
    },
 ]

@@ -153,14 +227,15 @@ 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)
+        model_description = self._get_model_description(model_name)
+        cache_dir = define_cache_dir(cache_dir)
+        model_dir = self.download_model(model_description, 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
-
-        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,
+        )

    def embed(
        self,
@@ -186,7 +261,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,
@@ -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 = "ace5b40bd629af8ec1a369f1cfb507db7bb6b18c1387357e0d029d9b50a66335"
+content-hash = "b71c2e05a840cee62f4d1a11c0377e93d735f473e13f684246c021e01c5368c3"
@@ -15,6 +15,7 @@ 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"
@@ -26,16 +26,26 @@ CANONICAL_VECTOR_VALUES = {
    "nomic-ai/nomic-embed-text-v1.5": np.array(
        [-1.6531514e-02, 8.5380634e-05, -1.8171231e-01, -3.9333291e-03, 1.2763254e-02]
    ),
+    "nomic-ai/nomic-embed-text-v1.5-Q": np.array(
+        [-0.01554983,  0.0129992 , -0.17909265, -0.01062993,  0.00512859]
+    ),
    "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]),
+    "snowflake/snowflake-arctic-embed-xs": np.array([0.0092, 0.0619, 0.0196, 0.009, -0.0114]),
+    "snowflake/snowflake-arctic-embed-s": np.array([-0.0416, -0.0867, 0.0209, 0.0554, -0.0272]),
+    "snowflake/snowflake-arctic-embed-m": np.array([-0.0329, 0.0364, 0.0481, 0.0016, 0.0328]),
+    "snowflake/snowflake-arctic-embed-m-long": np.array(
+        [0.0080, -0.0266, -0.0335, 0.0282, 0.0143]
+    ),
+    "snowflake/snowflake-arctic-embed-l": np.array([0.0189, -0.0673, 0.0183, 0.0124, 0.0146]),
 }


 def test_embedding():
-    is_ubuntu_ci = os.getenv("IS_UBUNTU_CI")
+    is_ci = os.getenv("CI")

    for model_desc in TextEmbedding.list_supported_models():
-        if is_ubuntu_ci == "false" and model_desc["size_in_GB"] > 1:
+        if not is_ci and model_desc["size_in_GB"] > 1:
            continue

        dim = model_desc["dim"]
Author	SHA1	Message	Date
George Panchuk	5afa07b921	new: update dev workflow	2024-05-02 20:34:47 +02:00
Arun	432da42c11	fix links (#215 )	2024-04-27 22:24:29 +05:30
George	5cde2898bc	new: remove slurm environment variables (#213 )	2024-04-26 21:45:02 +02:00
AnushandGeorge	ab7a99a748	feat: Quantized models (#201 ) * feat: Quantized models * refactor: use model_file for GCS * refactoring: refactor model downloading (#209) * refactoring: refactor model downloading * refactor: update docstring Co-authored-by: Anush <anushshetty90@gmail.com> * Update fastembed/common/model_management.py Co-authored-by: George <george.panchuk@qdrant.tech> * fix: model_file for Snowflake models --------- Co-authored-by: George <george.panchuk@qdrant.tech>	2024-04-26 17:27:16 +02:00
Anush	466886a317	ci: Schedule python-tests.yml (#211 ) * ci: Schedule python-tests.yml * ci: use emojis * ci: Bump action versions python-tests.yml * ci: python-tests.yml	2024-04-26 10:28:37 +05:30
Anush	cc4112d859	feat: Snowflake models (#207 ) * feat: Snowflake models * Added snowflake/snowflake-arctic-embed-m * docs: snowflake/snowflake-arctic-embed-m	2024-04-19 19:00:10 +05:30
Nirant	864217a7d9	Re-run of identical hardware and generate graphs (#205 )	2024-04-18 11:13:18 +05:30