debug: fix workflow

new: disable windows and macos tests on gpu and its child branches
fix links (#215 )
2026-09-22 14:07:51 -05:00 · 2024-05-02 00:06:21 +02:00 · 2024-05-01 23:00:31 +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
11 changed files with 330 additions and 204 deletions
@@ -2,7 +2,9 @@ name: Tests

 on:
  push:
-    branches: [ master, main ]
+    branches: [ main ]
+  schedule:
+    - cron: 0 0 * * *
  pull_request:

 env:
@@ -21,17 +23,28 @@ jobs:
          - '3.12.x'
        os:
          - ubuntu-latest
-          - macos-latest
-          - windows-latest
+        include:
+          - os: macos-latest
+            python-version: '3.8.x'
+            if: >-
+              ${{ !contains(github.ref, 'refs/heads/gpu') &&
+                  github.event_name != 'pull_request' ||
+                  github.event.pull_request.base.ref != 'gpu' }}
+          - os: windows-latest
+            python-version: '3.8.x'
+            if: >-
+              ${{ !contains(github.ref, 'refs/heads/gpu') &&
+                  github.event_name != 'pull_request' ||
+                  github.event.pull_request.base.ref != 'gpu' }}

    runs-on: ${{ matrix.os }}

    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 +52,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"
    }
@@ -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,27 +87,33 @@ 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,
    ) -> 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,
        )

@@ -229,9 +218,14 @@ 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
+                    )
                )
            except (EnvironmentError, RepositoryNotFoundError, ValueError) as e:
                logger.error(
@@ -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,
@@ -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",
    },
 ]

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

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

@@ -179,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,
@@ -212,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,
@@ -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	494f21fca9	debug: fix workflow	2024-05-02 00:06:21 +02:00
George Panchuk	edf1d52e4b	new: disable windows and macos tests on gpu and its child branches	2024-05-01 23:00:31 +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