new: gpu package (#224 )

add eofl
sync publih with main
2026-09-22 22:17:49 -05:00 · 2024-05-03 16:35:56 +00:00 · 2024-05-03 16:34:16 +00:00 · 2024-05-03 16:34:16 +00:00 · 2024-05-03 16:34:16 +00:00 · 2024-05-03 16:34:16 +00:00
21 changed files with 463 additions and 3679 deletions
@@ -15,7 +15,6 @@ on:
    tags:
      - 'v*'           # Push events to every version tag

-
 jobs:
  deploy:

@@ -1,8 +1,11 @@
 name: Tests
+run-name: Tests (gpu)

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

 env:
@@ -21,17 +24,15 @@ jobs:
          - '3.12.x'
        os:
          - ubuntu-latest
-          - macos-latest
-          - windows-latest

    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 +40,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?

@@ -14,12 +14,18 @@ The default text embedding (`TextEmbedding`) model is Flag Embedding, presented

 ## 🚀 Installation

-To install the FastEmbed library, pip works:
+To install the FastEmbed library, pip works best. You can install it with or without GPU support:

 ```bash
 pip install fastembed
 ```

+### ⚡️ With GPU
+
+```bash
+pip install fastembed-gpu
+```
+
 ## 📖 Quickstart

 ```python
@@ -42,6 +48,23 @@ embeddings_list = list(embedding_model.embed(documents))
 len(embeddings_list[0]) # Vector of 384 dimensions
 ```

+### ⚡️ FastEmbed on a GPU
+
+FastEmbed supports running on GPU devices. It requires installation of the `fastembed-gpu` package.
+Make sure not to have the `fastembed` package installed, as it might interfere with the `fastembed-gpu` package.
+
+```bash
+pip install fastembed-gpu
+``` 
+
+```python
+from fastembed import TextEmbedding
+
+embedding_model = TextEmbedding(model_name="BAAI/bge-small-en-v1.5", providers=["CUDAExecutionProvider"])
+print("The model BAAI/bge-small-en-v1.5 is ready to use on a GPU.")
+
+```
+
 ## Usage with Qdrant

 Installation with Qdrant Client in Python:
@@ -50,7 +73,13 @@ Installation with Qdrant Client in Python:
 pip install qdrant-client[fastembed]
 ```

-You might have to use ```pip install 'qdrant-client[fastembed]'``` on zsh.
+or 
+
+```bash
+pip install qdrant-client[fastembed-gpu]
+```
+
+You might have to use quotes ```pip install 'qdrant-client[fastembed]'``` on zsh.

 ```python
 from qdrant_client import QdrantClient
@@ -85,8 +114,4 @@ search_result = client.query(
    query_text="This is a query document"
 )
 print(search_result)
-```
-
-#### Similar Work
-
-Ilyas M. wrote about using [FlagEmbeddings with Optimum](https://twitter.com/IlysMoutawwakil/status/1705215192425288017) over CUDA.
+```
@@ -0,0 +1,41 @@
+# Releasing FastEmbed
+
+This is a guide how to release `fastembed` and `fastembed-gpu` packages.
+
+## How to
+
+1. Accumulate changes in the `main` branch.
+2. Bump the version in `pyproject.toml`
+
+3. Rebase the `gpu` branch on `main` and resolve conflicts if occurred:
+
+```bash
+git checkout gpu
+git rebase main
+git push origin gpu
+```
+
+4. Draft release notes
+5. Checkout to `main` and create a tag, e.g.:
+
+```bash
+git checkout main
+git tag -a v0.1.0 -m "Release v0.1.0"
+```
+
+6. Checkout `gpu` and create a tag, e.g.:
+
+```bash
+git checkout gpu
+git tag -a v0.1.0-gpu -m "Release v0.1.0"
+```
+
+7. Push tags:
+
+```bash
+git push --tags
+```
+
+8. Verify that both packages have been published successfully on PyPI. Try installing them and verify imports.
+9. Create a release on GitHub with the written release notes.
+
@@ -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"
    }
@@ -3,5 +3,10 @@ import importlib.metadata
 from fastembed.text import TextEmbedding
 from fastembed.sparse import SparseTextEmbedding, SparseEmbedding

-__version__ = importlib.metadata.version("fastembed")
+try:
+    version = importlib.metadata.version("fastembed")
+except importlib.metadata.PackageNotFoundError as _:
+    version = importlib.metadata.version("fastembed-gpu")
+
+__version__ = version
 __all__ = ["TextEmbedding", "SparseTextEmbedding", "SparseEmbedding"]
@@ -0,0 +1,3 @@
+from fastembed.common.onnx_model import OnnxProvider
+
+__all__ = ["OnnxProvider"]
@@ -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,28 +87,36 @@ class ModelManagement:

    @classmethod
    def download_files_from_huggingface(
-        cls, hf_source_repo: str, cache_dir: Optional[str] = None
+        cls,
+        hf_source_repo: str,
+        cache_dir: Optional[str] = None,
+        extra_patterns: Optional[List[str]] = None,
+        **kwargs,
    ) -> str:
        """
        Downloads a model from HuggingFace Hub.
        Args:
            hf_source_repo (str): Name of the model on HuggingFace Hub, e.g. "qdrant/all-MiniLM-L6-v2-onnx".
            cache_dir (Optional[str]): The path to the cache directory.
+            extra_patterns (Optional[List[str]]): extra patterns to allow in the snapshot download, typically
+                includes the required model files.
        Returns:
            Path: The path to the model directory.
        """
+        allow_patterns = [
+            "config.json",
+            "tokenizer.json",
+            "tokenizer_config.json",
+            "special_tokens_map.json",
+        ]
+        if extra_patterns is not None:
+            allow_patterns.extend(extra_patterns)

        return snapshot_download(
            repo_id=hf_source_repo,
-            allow_patterns=[
-                "*.onnx",
-                "*.onnx_data",
-                "config.json",
-                "tokenizer.json",
-                "tokenizer_config.json",
-                "special_tokens_map.json",
-            ],
+            allow_patterns=allow_patterns,
            cache_dir=cache_dir,
+            local_files_only=kwargs.get("local_files_only", False),
        )

    @classmethod
@@ -200,7 +191,7 @@ class ModelManagement:
        return model_dir

    @classmethod
-    def download_model(cls, model: Dict[str, Any], cache_dir: Path) -> Path:
+    def download_model(cls, model: Dict[str, Any], cache_dir: Path, **kwargs) -> Path:
        """
        Downloads a model from HuggingFace Hub or Google Cloud Storage.

@@ -229,9 +220,17 @@ class ModelManagement:
        url_source = model.get("sources", {}).get("url")

        if hf_source:
+            extra_patterns = [model["model_file"]]
+            extra_patterns.extend(model.get("additional_files", []))
+
            try:
                return Path(
-                    cls.download_files_from_huggingface(hf_source, cache_dir=str(cache_dir))
+                    cls.download_files_from_huggingface(
+                        hf_source,
+                        cache_dir=str(cache_dir),
+                        extra_patterns=extra_patterns,
+                        local_files_only=kwargs.get("local_files_only", False),
+                    )
                )
            except (EnvironmentError, RepositoryNotFoundError, ValueError) as e:
                logger.error(
@@ -1,19 +1,33 @@
 import os
 from multiprocessing import get_all_start_methods
 from pathlib import Path
-from typing import Any, Dict, Generic, Iterable, List, Optional, Tuple, Type, TypeVar, Union
+from typing import (
+    Any,
+    Dict,
+    Generic,
+    Iterable,
+    List,
+    Optional,
+    Tuple,
+    Type,
+    TypeVar,
+    Union,
+    Sequence,
+)

 import numpy as np
 import onnxruntime as ort

-from fastembed.common.model_management import locate_model_file
 from fastembed.common.models import load_tokenizer
 from fastembed.common.utils import iter_batch
 from fastembed.parallel_processor import ParallelWorkerPool, Worker

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

+OnnxProvider = Union[str, Tuple[str, Dict[Any, Any]]]
+

 class OnnxModel(Generic[T]):
    @classmethod
@@ -34,23 +48,35 @@ class OnnxModel(Generic[T]):
        """
        return onnx_input

-    def load_onnx_model(self, model_dir: Path, threads: Optional[int], max_length: int) -> None:
-        model_path = locate_model_file(model_dir, ["model.onnx", "model_optimized.onnx"])
+    def load_onnx_model(
+        self,
+        model_dir: Path,
+        model_file: str,
+        threads: Optional[int],
+        providers: Optional[Sequence[OnnxProvider]] = None,
+    ) -> None:
+        model_path = model_dir / model_file

        # List of Execution Providers: https://onnxruntime.ai/docs/execution-providers
-        onnx_providers = ["CPUExecutionProvider"]
+
+        onnx_providers = ["CPUExecutionProvider"] if providers is None else list(providers)
+        available_providers = ort.get_available_providers()
+        for provider in onnx_providers:
+            # check providers available
+            provider_name = provider if isinstance(provider, str) else provider[0]
+            if provider_name not in available_providers:
+                raise ValueError(
+                    f"Provider {provider_name} is not available. Available providers: {available_providers}"
+                )

        so = ort.SessionOptions()
-        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
        )
@@ -32,6 +32,7 @@ class SparseTextEmbeddingBase(ModelManagement):
        self.model_name = model_name
        self.cache_dir = cache_dir
        self.threads = threads
+        self._local_files_only = kwargs.pop("local_files_only", False)

    def embed(
        self,
@@ -1,5 +1,6 @@
-from typing import List, Type, Dict, Any, Union, Iterable, Optional
+from typing import List, Type, Dict, Any, Union, Iterable, Optional, Sequence

+from fastembed.common import OnnxProvider
 from fastembed.sparse.sparse_embedding_base import SparseTextEmbeddingBase, SparseEmbedding
 from fastembed.sparse.splade_pp import SpladePP

@@ -42,6 +43,7 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
        model_name: str,
        cache_dir: Optional[str] = None,
        threads: Optional[int] = None,
+        providers: Optional[Sequence[OnnxProvider]] = None,
        **kwargs,
    ):
        super().__init__(model_name, cache_dir, threads, **kwargs)
@@ -49,7 +51,9 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
        for EMBEDDING_MODEL_TYPE in self.EMBEDDINGS_REGISTRY:
            supported_models = EMBEDDING_MODEL_TYPE.list_supported_models()
            if any(model_name.lower() == model["model"].lower() for model in supported_models):
-                self.model = EMBEDDING_MODEL_TYPE(model_name, cache_dir, threads, **kwargs)
+                self.model = EMBEDDING_MODEL_TYPE(
+                    model_name, cache_dir, threads, providers=providers, **kwargs
+                )
                return

        raise ValueError(
@@ -1,8 +1,8 @@
-from typing import Any, Dict, Iterable, List, Optional, Tuple, Union, Type
+from typing import Any, Dict, Iterable, List, Optional, Tuple, Union, Type, Sequence

 import numpy as np

-from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel
+from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxProvider
 from fastembed.common.utils import define_cache_dir
 from fastembed.sparse.sparse_embedding_base import SparseEmbedding, SparseTextEmbeddingBase

@@ -15,6 +15,7 @@ supported_splade_models = [
        "sources": {
            "hf": "Qdrant/SPLADE_PP_en_v1",
        },
+        "model_file": "model.onnx",
    },
    {
        "model": "prithivida/Splade_PP_en_v1",
@@ -24,6 +25,7 @@ supported_splade_models = [
        "sources": {
            "hf": "Qdrant/SPLADE_PP_en_v1",
        },
+        "model_file": "model.onnx",
    },
 ]

@@ -61,6 +63,7 @@ class SpladePP(SparseTextEmbeddingBase, OnnxModel[SparseEmbedding]):
        model_name: str,
        cache_dir: Optional[str] = None,
        threads: Optional[int] = None,
+        providers: Optional[Sequence[OnnxProvider]] = None,
        **kwargs,
    ):
        """
@@ -77,14 +80,19 @@ class SpladePP(SparseTextEmbeddingBase, OnnxModel[SparseEmbedding]):

        super().__init__(model_name, cache_dir, threads, **kwargs)

-        self.model_name = model_name
-        self._model_description = self._get_model_description(model_name)
+        model_description = self._get_model_description(model_name)
+        cache_dir = define_cache_dir(cache_dir)

-        self._cache_dir = define_cache_dir(cache_dir)
-        self._model_dir = self.download_model(self._model_description, self._cache_dir)
-        self._max_length = 512
+        model_dir = self.download_model(
+            model_description, cache_dir, local_files_only=self._local_files_only
+        )

-        self.load_onnx_model(self._model_dir, self.threads, self._max_length)
+        self.load_onnx_model(
+            model_dir=model_dir,
+            model_file=model_description["model_file"],
+            threads=threads,
+            providers=providers,
+        )

    def embed(
        self,
@@ -110,7 +118,7 @@ class SpladePP(SparseTextEmbeddingBase, OnnxModel[SparseEmbedding]):
        """
        yield from self._embed_documents(
            model_name=self.model_name,
-            cache_dir=str(self._cache_dir),
+            cache_dir=str(self.cache_dir),
            documents=documents,
            batch_size=batch_size,
            parallel=parallel,
@@ -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",
    },
 ]

@@ -1,8 +1,8 @@
-from typing import Dict, Optional, Tuple, Union, Iterable, Type, List, Any
+from typing import Dict, Optional, Tuple, Union, Iterable, Type, List, Any, Sequence

 import numpy as np

-from fastembed.common.onnx_model import OnnxModel, EmbeddingWorker
+from fastembed.common.onnx_model import OnnxModel, EmbeddingWorker, OnnxProvider
 from fastembed.common.models import normalize
 from fastembed.common.utils import define_cache_dir
 from fastembed.text.text_embedding_base import TextEmbeddingBase
@@ -16,6 +16,7 @@ supported_onnx_models = [
        "sources": {
            "url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en.tar.gz",
        },
+        "model_file": "model_optimized.onnx",
    },
    {
        "model": "BAAI/bge-base-en-v1.5",
@@ -26,6 +27,7 @@ supported_onnx_models = [
            "url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en-v1.5.tar.gz",
            "hf": "qdrant/bge-base-en-v1.5-onnx-q",
        },
+        "model_file": "model_optimized.onnx",
    },
    {
        "model": "BAAI/bge-large-en-v1.5",
@@ -35,6 +37,7 @@ supported_onnx_models = [
        "sources": {
            "hf": "qdrant/bge-large-en-v1.5-onnx",
        },
+        "model_file": "model.onnx",
    },
    {
        "model": "BAAI/bge-small-en",
@@ -44,18 +47,8 @@ supported_onnx_models = [
        "sources": {
            "url": "https://storage.googleapis.com/qdrant-fastembed/BAAI-bge-small-en.tar.gz",
        },
+        "model_file": "model_optimized.onnx",
    },
-    # {
-    #     "model": "BAAI/bge-small-en",
-    #     "dim": 384,
-    #     "description": "Fast English model",
-    #     "size_in_GB": 0.2,
-    #     "hf_sources": [],
-    #     "compressed_url_sources": [
-    #         "https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-en.tar.gz",
-    #         "https://storage.googleapis.com/qdrant-fastembed/BAAI-bge-small-en.tar.gz"
-    #     ]
-    # },
    {
        "model": "BAAI/bge-small-en-v1.5",
        "dim": 384,
@@ -64,6 +57,7 @@ supported_onnx_models = [
        "sources": {
            "hf": "qdrant/bge-small-en-v1.5-onnx-q",
        },
+        "model_file": "model_optimized.onnx",
    },
    {
        "model": "BAAI/bge-small-zh-v1.5",
@@ -73,6 +67,7 @@ supported_onnx_models = [
        "sources": {
            "url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-zh-v1.5.tar.gz",
        },
+        "model_file": "model_optimized.onnx",
    },
    {
        "model": "sentence-transformers/all-MiniLM-L6-v2",
@@ -83,6 +78,7 @@ supported_onnx_models = [
            "url": "https://storage.googleapis.com/qdrant-fastembed/sentence-transformers-all-MiniLM-L6-v2.tar.gz",
            "hf": "qdrant/all-MiniLM-L6-v2-onnx",
        },
+        "model_file": "model.onnx",
    },
    {
        "model": "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
@@ -92,6 +88,7 @@ supported_onnx_models = [
        "sources": {
            "hf": "qdrant/paraphrase-multilingual-MiniLM-L12-v2-onnx-Q",
        },
+        "model_file": "model_optimized.onnx",
    },
    {
        "model": "nomic-ai/nomic-embed-text-v1",
@@ -101,6 +98,7 @@ supported_onnx_models = [
        "sources": {
            "hf": "nomic-ai/nomic-embed-text-v1",
        },
+        "model_file": "onnx/model.onnx",
    },
    {
        "model": "nomic-ai/nomic-embed-text-v1.5",
@@ -110,6 +108,17 @@ supported_onnx_models = [
        "sources": {
            "hf": "nomic-ai/nomic-embed-text-v1.5",
        },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "nomic-ai/nomic-embed-text-v1.5-Q",
+        "dim": 768,
+        "description": "Quantized 8192 context length english model",
+        "size_in_GB": 0.13,
+        "sources": {
+            "hf": "nomic-ai/nomic-embed-text-v1.5",
+        },
+        "model_file": "onnx/model_quantized.onnx",
    },
    {
        "model": "thenlper/gte-large",
@@ -119,20 +128,8 @@ supported_onnx_models = [
        "sources": {
            "hf": "qdrant/gte-large-onnx",
        },
+        "model_file": "model.onnx",
    },
-    # {
-    #     "model": "sentence-transformers/all-MiniLM-L6-v2",
-    #     "dim": 384,
-    #     "description": "Sentence Transformer model, MiniLM-L6-v2",
-    #     "size_in_GB": 0.09,
-    #     "hf_sources": [
-    #         "qdrant/all-MiniLM-L6-v2-onnx"
-    #     ],
-    #     "compressed_url_sources": [
-    #         "https://storage.googleapis.com/qdrant-fastembed/fast-all-MiniLM-L6-v2.tar.gz",
-    #         "https://storage.googleapis.com/qdrant-fastembed/sentence-transformers-all-MiniLM-L6-v2.tar.gz"
-    #     ]
-    # }
    {
        "model": "mixedbread-ai/mxbai-embed-large-v1",
        "dim": 1024,
@@ -141,6 +138,57 @@ supported_onnx_models = [
        "sources": {
            "hf": "mixedbread-ai/mxbai-embed-large-v1",
        },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "snowflake/snowflake-arctic-embed-xs",
+        "dim": 384,
+        "description": "Based on all-MiniLM-L6-v2 model with only 22m parameters, ideal for latency/TCO budgets.",
+        "size_in_GB": 0.09,
+        "sources": {
+            "hf": "snowflake/snowflake-arctic-embed-xs",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "snowflake/snowflake-arctic-embed-s",
+        "dim": 384,
+        "description": "Based on infloat/e5-small-unsupervised, does not trade off retrieval accuracy for its small size.",
+        "size_in_GB": 0.13,
+        "sources": {
+            "hf": "snowflake/snowflake-arctic-embed-s",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "snowflake/snowflake-arctic-embed-m",
+        "dim": 768,
+        "description": "Based on intfloat/e5-base-unsupervised model, provides the best retrieval without slowing down inference.",
+        "size_in_GB": 0.43,
+        "sources": {
+            "hf": "Snowflake/snowflake-arctic-embed-m",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "snowflake/snowflake-arctic-embed-m-long",
+        "dim": 768,
+        "description": "Based on nomic-ai/nomic-embed-text-v1-unsupervised model, 8192 context-length model",
+        "size_in_GB": 0.54,
+        "sources": {
+            "hf": "snowflake/snowflake-arctic-embed-m-long",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "snowflake/snowflake-arctic-embed-l",
+        "dim": 1024,
+        "description": "Based on intfloat/e5-large-unsupervised, large model for most accurate retrieval.",
+        "size_in_GB": 1.02,
+        "sources": {
+            "hf": "snowflake/snowflake-arctic-embed-l",
+        },
+        "model_file": "onnx/model.onnx",
    },
 ]

@@ -163,6 +211,7 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxModel[np.ndarray]):
        model_name: str = "BAAI/bge-small-en-v1.5",
        cache_dir: Optional[str] = None,
        threads: Optional[int] = None,
+        providers: Optional[Sequence[OnnxProvider]] = None,
        **kwargs,
    ):
        """
@@ -179,14 +228,18 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxModel[np.ndarray]):

        super().__init__(model_name, cache_dir, threads, **kwargs)

-        self.model_name = model_name
-        self._model_description = self._get_model_description(model_name)
+        model_description = self._get_model_description(model_name)
+        cache_dir = define_cache_dir(cache_dir)
+        model_dir = self.download_model(
+            model_description, cache_dir, local_files_only=self._local_files_only
+        )

-        self._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,
+            providers=providers,
+        )

    def embed(
        self,
@@ -212,7 +265,7 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxModel[np.ndarray]):
        """
        yield from self._embed_documents(
            model_name=self.model_name,
-            cache_dir=str(self._cache_dir),
+            cache_dir=str(self.cache_dir),
            documents=documents,
            batch_size=batch_size,
            parallel=parallel,
@@ -1,7 +1,8 @@
-from typing import Any, Dict, Iterable, List, Optional, Type, Union
+from typing import Any, Dict, Iterable, List, Optional, Type, Union, Sequence

 import numpy as np

+from fastembed.common import OnnxProvider
 from fastembed.text.e5_onnx_embedding import E5OnnxEmbedding
 from fastembed.text.jina_onnx_embedding import JinaOnnxEmbedding
 from fastembed.text.onnx_embedding import OnnxTextEmbedding
@@ -49,6 +50,7 @@ class TextEmbedding(TextEmbeddingBase):
        model_name: str = "BAAI/bge-small-en-v1.5",
        cache_dir: Optional[str] = None,
        threads: Optional[int] = None,
+        providers: Optional[Sequence[OnnxProvider]] = None,
        **kwargs,
    ):
        super().__init__(model_name, cache_dir, threads, **kwargs)
@@ -56,7 +58,9 @@ class TextEmbedding(TextEmbeddingBase):
        for EMBEDDING_MODEL_TYPE in self.EMBEDDINGS_REGISTRY:
            supported_models = EMBEDDING_MODEL_TYPE.list_supported_models()
            if any(model_name.lower() == model["model"].lower() for model in supported_models):
-                self.model = EMBEDDING_MODEL_TYPE(model_name, cache_dir, threads, **kwargs)
+                self.model = EMBEDDING_MODEL_TYPE(
+                    model_name, cache_dir, threads, providers=providers, **kwargs
+                )
                return

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

    def embed(
        self,
@@ -1,8 +1,8 @@
 [tool.poetry]
-name = "fastembed"
-version = "0.2.6"
+name = "fastembed-gpu"
+version = "0.2.7"
 description = "Fast, light, accurate library built for retrieval embedding generation"
-authors = ["NirantK <nirant.bits@gmail.com>"]
+authors = ["NirantK <nirant.bits@gmail.com>", "Qdrant Team <info@qdrant.tech>"]
 license = "Apache License"
 readme = "README.md"
 packages = [{include = "fastembed"}]
@@ -12,11 +12,10 @@ keywords = ["vector", "embedding", "neural", "search", "qdrant", "sentence-trans

 [tool.poetry.dependencies]
 python = ">=3.8.0,<3.13"
-onnx = "^1.15.0"
-onnxruntime = "^1.17.0"
+onnxruntime-gpu = "^1.17.0"
 tqdm = "^4.66"
 requests = "^2.31"
-tokenizers = "^0.15.1"
+tokenizers = "^0.15"
 huggingface-hub = "^0.20"
 loguru = "^0.7.2"
 numpy = [
@@ -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	880a277640	new: gpu package (#224 )	2024-05-03 16:35:56 +00:00
George Panchuk	e04ac0b335	add eofl	2024-05-03 16:34:16 +00:00
George Panchuk	8991bdc541	sync publih with main	2024-05-03 16:34:16 +00:00
George Panchuk	a58ae07017	fix: workflow dispatch can only be triggered from the default branch	2024-05-03 16:34:16 +00:00
George Panchuk	8bd314795d	alter workflow	2024-05-03 16:34:16 +00:00
George Panchuk	9f4bc3a475	refactoring: alter workflow names	2024-05-03 16:34:16 +00:00
George Panchuk	7c69436237	fix: do not run windows and mac os tests on gpu branch	2024-05-03 16:34:16 +00:00
George Panchuk	e829c89c5c	new: gpu package publish workflow	2024-05-03 16:34:16 +00:00
Andrey Vasnetsov	6ecab6d40d	bump v0.2.7	2024-05-03 16:33:06 +00:00
George Nirant Kasliwal Nirant	d8c592032b	new: allow users to override providers (#214 ) * new: add gpu support, allow users to override providers * fix: update poetry.lock * fix: fix type hint for 3.8 * [readme] Remove similar work * [README] Add GPU support for FastEmbed library * [README] Add device check * fix: revert changes to pyproject and lock, update readme * Update poetry.lock * new: add type alias for providers, add explicit providers to embeddings --------- Co-authored-by: Nirant Kasliwal <nirant.bits@gmail.com> Co-authored-by: Nirant <NirantK@users.noreply.github.com>	2024-05-03 18:31:22 +02:00
GeorgeandAndrey Vasnetsov	da603b8b7d	new: add release instructions (#231 ) * new: add release instructions * review fixes --------- Co-authored-by: Andrey Vasnetsov <andrey@vasnetsov.com>	2024-05-03 18:30:54 +02:00
Andrey Vasnetsov	562d604375	Merge pull request #223 from Waffleboy/main [Bugfix] Allow user to pick local mode only so huggingface does not do a network call and timeout	2024-05-03 18:29:54 +02:00
Andrey Vasnetsov	8b1a98a6a3	Merge pull request #230 from qdrant/update-tokenizers new: update tokenizers	2024-05-03 18:26:09 +02:00
Andrey Vasnetsov	3c9b147e0a	version	2024-05-03 16:21:33 +00:00
George Panchuk	6abd415f4a	fix: add local_files_only to sparse, formatting, refactor	2024-05-03 17:49:10 +02:00
George Panchuk	8184acbb39	new: update tokenizers	2024-05-03 16:57:45 +02:00
George	47cf7f9f92	new: add gpu package into workflow (#228 ) * new: add gpu package into workflow * remove gpu tag	2024-05-03 16:30:23 +02:00
generall	f7896c81f3	do not ship poetry.lock with the repo, as package users wont have it anyway	2024-05-03 13:00:53 +02:00
Thiru	4a59d09248	Allow user to pick local mode only so huggingface does not do a network call and timeout	2024-05-02 00:25:24 +08: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