mirror of
https://github.com/qdrant/fastembed.git
synced 2026-09-22 22:17:49 -05:00
Compare commits
2
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 Anush
|
ff82bf7174 | ||
|
Anush008
|
c0e457305a |
@@ -15,6 +15,7 @@ on:
|
||||
tags:
|
||||
- 'v*' # Push events to every version tag
|
||||
|
||||
|
||||
jobs:
|
||||
deploy:
|
||||
|
||||
|
||||
@@ -2,9 +2,7 @@ name: Tests
|
||||
|
||||
on:
|
||||
push:
|
||||
branches: [ master, main, gpu ]
|
||||
schedule:
|
||||
- cron: 0 0 * * *
|
||||
branches: [ master, main ]
|
||||
pull_request:
|
||||
|
||||
env:
|
||||
@@ -31,9 +29,9 @@ jobs:
|
||||
name: Python ${{ matrix.python-version }} on ${{ matrix.os }} test
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v3
|
||||
- uses: actions/checkout@v2
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v5
|
||||
uses: actions/setup-python@v2
|
||||
with:
|
||||
python-version: ${{ matrix.python-version }}
|
||||
- name: Install dependencies
|
||||
@@ -41,14 +39,8 @@ jobs:
|
||||
python -m pip install poetry
|
||||
poetry config virtualenvs.create false
|
||||
poetry install --no-interaction --no-ansi --without docs
|
||||
|
||||
- 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'
|
||||
- name: Run tests
|
||||
run: |
|
||||
export IS_UBUNTU_CI=$(test "${{ matrix.os }}" = "ubuntu-latest" && echo "true" || echo "false")
|
||||
pytest
|
||||
shell: bash
|
||||
|
||||
@@ -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/qdrant/Retrieval_with_FastEmbed/) and how to use [FastEmbed with Qdrant](https://qdrant.github.io/fastembed/qdrant/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/examples/Retrieval_with_FastEmbed/) and how to use [FastEmbed with Qdrant](https://qdrant.github.io/fastembed/examples/Usage_With_Qdrant/).
|
||||
|
||||
## 📈 Why FastEmbed?
|
||||
|
||||
@@ -14,18 +14,12 @@ The default text embedding (`TextEmbedding`) model is Flag Embedding, presented
|
||||
|
||||
## 🚀 Installation
|
||||
|
||||
To install the FastEmbed library, pip works best. You can install it with or without GPU support:
|
||||
To install the FastEmbed library, pip works:
|
||||
|
||||
```bash
|
||||
pip install fastembed
|
||||
```
|
||||
|
||||
### ⚡️ With GPU
|
||||
|
||||
```bash
|
||||
pip install fastembed-gpu
|
||||
```
|
||||
|
||||
## 📖 Quickstart
|
||||
|
||||
```python
|
||||
@@ -48,23 +42,6 @@ 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:
|
||||
@@ -73,13 +50,7 @@ Installation with Qdrant Client in Python:
|
||||
pip install qdrant-client[fastembed]
|
||||
```
|
||||
|
||||
or
|
||||
|
||||
```bash
|
||||
pip install qdrant-client[fastembed-gpu]
|
||||
```
|
||||
|
||||
You might have to use quotes ```pip install 'qdrant-client[fastembed]'``` on zsh.
|
||||
You might have to use ```pip install 'qdrant-client[fastembed]'``` on zsh.
|
||||
|
||||
```python
|
||||
from qdrant_client import QdrantClient
|
||||
@@ -114,4 +85,8 @@ 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.
|
||||
|
||||
-41
@@ -1,41 +0,0 @@
|
||||
# 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.
|
||||
|
||||
File diff suppressed because one or more lines are too long
@@ -17,7 +17,7 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"execution_count": 2,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
@@ -35,7 +35,7 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 2,
|
||||
"execution_count": 6,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
@@ -89,125 +89,90 @@
|
||||
" </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>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",
|
||||
" <th>4</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>7</th>\n",
|
||||
" <th>5</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>8</th>\n",
|
||||
" <th>6</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>9</th>\n",
|
||||
" <th>7</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>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>12</th>\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>13</th>\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>14</th>\n",
|
||||
" <td>snowflake/snowflake-arctic-embed-m-long</td>\n",
|
||||
" <th>10</th>\n",
|
||||
" <td>jinaai/jina-embeddings-v2-base-en</td>\n",
|
||||
" <td>768</td>\n",
|
||||
" <td>Based on nomic-ai/nomic-embed-text-v1-unsuperv...</td>\n",
|
||||
" <td>0.540</td>\n",
|
||||
" <td>English embedding model supporting 8192 sequen...</td>\n",
|
||||
" <td>0.520</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>15</th>\n",
|
||||
" <th>11</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>16</th>\n",
|
||||
" <th>12</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>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",
|
||||
" <th>13</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>19</th>\n",
|
||||
" <th>14</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>20</th>\n",
|
||||
" <th>15</th>\n",
|
||||
" <td>intfloat/multilingual-e5-large</td>\n",
|
||||
" <td>1024</td>\n",
|
||||
" <td>Multilingual model, e5-large. Recommend using ...</td>\n",
|
||||
@@ -222,50 +187,40 @@
|
||||
"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 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",
|
||||
"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",
|
||||
"\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 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 "
|
||||
"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 "
|
||||
]
|
||||
},
|
||||
"execution_count": 2,
|
||||
"execution_count": 6,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
|
||||
@@ -1,323 +0,0 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "f7f8e989-cdc9-475e-918d-af20530fcfe6",
|
||||
"metadata": {
|
||||
"is_executing": true
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"!pip3 install -q torch transformers optimum pillow"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "e8c9e276-58f2-45a4-af40-6d7bacc30eec",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import os\n",
|
||||
"from pathlib import Path\n",
|
||||
"from typing import Optional, Dict, Union, Tuple\n",
|
||||
"\n",
|
||||
"import torch\n",
|
||||
"import numpy as np\n",
|
||||
"from PIL import Image\n",
|
||||
"from transformers import (\n",
|
||||
" CLIPVisionModelWithProjection,\n",
|
||||
" CLIPTextModelWithProjection,\n",
|
||||
" CLIPImageProcessor,\n",
|
||||
" CLIPTokenizerFast,\n",
|
||||
")\n",
|
||||
"from transformers.models.clip.modeling_clip import (\n",
|
||||
" CLIPTextModelOutput,\n",
|
||||
" CLIPVisionModelOutput,\n",
|
||||
" CLIPModel,\n",
|
||||
")\n",
|
||||
"from optimum.onnxruntime import ORTModelForCustomTasks\n",
|
||||
"from optimum.exporters.onnx.model_configs import CLIPTextWithProjectionOnnxConfig, ViTOnnxConfig\n",
|
||||
"from optimum.exporters.onnx import export_models"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "bbceb40f-22cd-4d92-be6e-fe14f16f7bc2",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"model_id = \"openai/clip-vit-base-patch32\"\n",
|
||||
"output_dir = \"split-clip-onnx\""
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "f96ff7fb-518e-405e-a7e0-46f836ffdec8",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"class CLIPVisionModelWithProjectionOnnxConfig(ViTOnnxConfig):\n",
|
||||
" @property\n",
|
||||
" def outputs(self) -> Dict[str, Dict[int, str]]:\n",
|
||||
" return {\n",
|
||||
" \"image_embeds\": {0: \"batch_size\"},\n",
|
||||
" }"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "18863f0b-6bd5-463f-bebc-38bf40d51b9c",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"class CLIPTextModelWithProjectionAndAttentionOnnxConfig(CLIPTextWithProjectionOnnxConfig):\n",
|
||||
" @property\n",
|
||||
" def inputs(self) -> Dict[str, Dict[int, str]]:\n",
|
||||
" return {\n",
|
||||
" \"input_ids\": {0: \"batch_size\", 1: \"sequence_length\"},\n",
|
||||
" \"attention_mask\": {0: \"batch_size\", 1: \"sequence_length\"},\n",
|
||||
" }"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "5d37b16a-ec30-40e1-8404-0f0d51abfa76",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"class CLIPTextModelWithProjectionNormalized(CLIPTextModelWithProjection):\n",
|
||||
" def forward(\n",
|
||||
" self,\n",
|
||||
" input_ids: Optional[torch.Tensor] = None,\n",
|
||||
" attention_mask: Optional[torch.Tensor] = None,\n",
|
||||
" position_ids: Optional[torch.Tensor] = None,\n",
|
||||
" output_attentions: Optional[bool] = None,\n",
|
||||
" output_hidden_states: Optional[bool] = None,\n",
|
||||
" return_dict: Optional[bool] = None,\n",
|
||||
" ) -> Union[Tuple, CLIPTextModelOutput]:\n",
|
||||
" text_outputs = super().forward(\n",
|
||||
" input_ids,\n",
|
||||
" attention_mask,\n",
|
||||
" position_ids,\n",
|
||||
" output_attentions,\n",
|
||||
" output_hidden_states,\n",
|
||||
" return_dict,\n",
|
||||
" )\n",
|
||||
" normalized_text_embeds = text_outputs.text_embeds / text_outputs.text_embeds.norm(\n",
|
||||
" p=2, dim=-1, keepdim=True\n",
|
||||
" )\n",
|
||||
" return CLIPTextModelOutput(\n",
|
||||
" text_embeds=normalized_text_embeds,\n",
|
||||
" last_hidden_state=text_outputs.last_hidden_state,\n",
|
||||
" hidden_states=text_outputs.hidden_states,\n",
|
||||
" attentions=text_outputs.attentions,\n",
|
||||
" )"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "2cd05d33-5f4f-4b36-aa2c-43fd97d5061d",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"class CLIPVisionModelWithProjectionNormalized(CLIPVisionModelWithProjection):\n",
|
||||
" def forward(\n",
|
||||
" self,\n",
|
||||
" pixel_values: Optional[torch.FloatTensor] = None,\n",
|
||||
" output_attentions: Optional[bool] = None,\n",
|
||||
" output_hidden_states: Optional[bool] = None,\n",
|
||||
" return_dict: Optional[bool] = None,\n",
|
||||
" ) -> Union[Tuple, CLIPVisionModelOutput]:\n",
|
||||
" vision_outputs = super().forward(pixel_values, return_dict)\n",
|
||||
" normalized_image_embeds = vision_outputs.image_embeds / vision_outputs.image_embeds.norm(\n",
|
||||
" p=2, dim=-1, keepdim=True\n",
|
||||
" )\n",
|
||||
" return CLIPVisionModelOutput(\n",
|
||||
" image_embeds=normalized_image_embeds,\n",
|
||||
" last_hidden_state=vision_outputs.last_hidden_state,\n",
|
||||
" hidden_states=vision_outputs.hidden_states,\n",
|
||||
" attentions=vision_outputs.attentions,\n",
|
||||
" )"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "fb5e5617-f8ce-4dcf-9148-68ddd91854c9",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"text_model = CLIPTextModelWithProjectionNormalized.from_pretrained(model_id)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "7f578131-c6bb-460e-8200-d0b7f0aa4135",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"vision_model = CLIPVisionModelWithProjectionNormalized.from_pretrained(model_id)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "c57db080-642f-4b62-96ad-75af9c0ab277",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"text_config = CLIPTextModelWithProjectionAndAttentionOnnxConfig(text_model.config)\n",
|
||||
"vision_config = CLIPVisionModelWithProjectionOnnxConfig(vision_model.config)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "cdb01aab-0dff-4fd5-9297-d16599274fdb",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"text_model.config.save_pretrained(f\"./{output_dir}/text\")\n",
|
||||
"vision_model.config.save_pretrained(f\"./{output_dir}/image\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "133bcd2c-57ae-4132-a691-3d129057f275",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"export_models(\n",
|
||||
" models_and_onnx_configs={\n",
|
||||
" \"text_model\": (text_model, text_config),\n",
|
||||
" \"vision_model\": (vision_model, vision_config),\n",
|
||||
" },\n",
|
||||
" output_dir=Path(f\"./{output_dir}\"),\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "fd9167b9-0d00-4a24-8f5e-41392d923b95",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"os.rename(f\"./{output_dir}/text_model.onnx\", f\"./{output_dir}/text/model.onnx\")\n",
|
||||
"os.rename(f\"./{output_dir}/vision_model.onnx\", f\"./{output_dir}/image/model.onnx\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "2281cc75-1027-4dbe-a7a5-86ee1dad4c3e",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"ort_vision_model = ORTModelForCustomTasks.from_pretrained(\n",
|
||||
" f\"./{output_dir}/image\", config=vision_config\n",
|
||||
")\n",
|
||||
"image_processor = CLIPImageProcessor.from_pretrained(\"openai/clip-vit-base-patch32\")\n",
|
||||
"image_input = image_processor(images=Image.open(\"assets/image.jpeg\"), return_tensors=\"pt\")\n",
|
||||
"\n",
|
||||
"with torch.inference_mode():\n",
|
||||
" image_outputs = ort_vision_model(**image_input)\n",
|
||||
"image_processor.save_pretrained(f\"./{output_dir}/image\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "5ce61045-b1e7-4b62-a47b-4bcb27ac7288",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"ort_text_model = ORTModelForCustomTasks.from_pretrained(f\"./{output_dir}/text\", config=text_config)\n",
|
||||
"text_processor = CLIPTokenizerFast.from_pretrained(\"openai/clip-vit-base-patch32\")\n",
|
||||
"text_input = text_processor(\"What am I using?\", return_tensors=\"pt\")\n",
|
||||
"\n",
|
||||
"with torch.inference_mode():\n",
|
||||
" text_outputs = ort_text_model(**text_input)\n",
|
||||
"text_processor.save_pretrained(f\"./{output_dir}/text\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "ee1e3d13-6884-4aa3-a6ab-aaa41fc17134",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"clip_model = CLIPModel.from_pretrained(\"openai/clip-vit-base-patch32\")\n",
|
||||
"inputs = {**text_input, **image_input}\n",
|
||||
"clip_model.eval()\n",
|
||||
"with torch.inference_mode():\n",
|
||||
" gt_output = clip_model(**inputs)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "ec07ce55-a0e7-42d2-b1f4-ba1370ab45b7",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"print(np.allclose(gt_output.text_embeds.numpy(), text_outputs.text_embeds, atol=1e-6))\n",
|
||||
"print(np.allclose(gt_output.image_embeds.numpy(), image_outputs.image_embeds, atol=1e-6))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "e9b15597-3054-40f1-a080-fea19d378d88",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"token = \"<token>\"\n",
|
||||
"# create_repo(repo_id='Qdrant/clip-ViT-B-32-vision', exist_ok=True, token=token)\n",
|
||||
"# create_repo(repo_id='Qdrant/clip-ViT-B-32-text', exist_ok=True, token=token)\n",
|
||||
"\n",
|
||||
"ort_text_model.push_to_hub(\n",
|
||||
" save_directory=f\"./{output_dir}/text/\",\n",
|
||||
" repository_id=\"Qdrant/clip-ViT-B-32-text\",\n",
|
||||
" use_auth_token=token,\n",
|
||||
")\n",
|
||||
"ort_vision_model.push_to_hub(\n",
|
||||
" save_directory=f\"./{output_dir}/image\",\n",
|
||||
" repository_id=\"Qdrant/clip-ViT-B-32-vision\",\n",
|
||||
" use_auth_token=token,\n",
|
||||
")"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3 (ipykernel)",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.10.12"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 169 KiB |
@@ -3,10 +3,5 @@ import importlib.metadata
|
||||
from fastembed.text import TextEmbedding
|
||||
from fastembed.sparse import SparseTextEmbedding, SparseEmbedding
|
||||
|
||||
try:
|
||||
version = importlib.metadata.version("fastembed")
|
||||
except importlib.metadata.PackageNotFoundError as _:
|
||||
version = importlib.metadata.version("fastembed-gpu")
|
||||
|
||||
__version__ = version
|
||||
__version__ = importlib.metadata.version("fastembed")
|
||||
__all__ = ["TextEmbedding", "SparseTextEmbedding", "SparseEmbedding"]
|
||||
|
||||
@@ -1,3 +0,0 @@
|
||||
from fastembed.common.onnx_model import OnnxProvider
|
||||
|
||||
__all__ = ["OnnxProvider"]
|
||||
|
||||
@@ -1,16 +1,28 @@
|
||||
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]]:
|
||||
@@ -41,157 +53,34 @@ 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,
|
||||
extra_patterns: Optional[List[str]] = None,
|
||||
**kwargs,
|
||||
cls, hf_source_repo: str, cache_dir: Optional[str] = None
|
||||
) -> str:
|
||||
"""
|
||||
Downloads a model from HuggingFace Hub.
|
||||
Args:
|
||||
hf_source_repo (str): Name of the model on HuggingFace Hub, e.g. "qdrant/all-MiniLM-L6-v2-onnx".
|
||||
cache_dir (Optional[str]): The path to the cache directory.
|
||||
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=allow_patterns,
|
||||
allow_patterns=[
|
||||
"*.onnx",
|
||||
"*.onnx_data",
|
||||
"config.json",
|
||||
"tokenizer.json",
|
||||
"tokenizer_config.json",
|
||||
"special_tokens_map.json",
|
||||
],
|
||||
cache_dir=cache_dir,
|
||||
local_files_only=kwargs.get("local_files_only", False),
|
||||
)
|
||||
|
||||
@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, **kwargs) -> Path:
|
||||
def download_model(cls, model: Dict[str, Any], cache_dir: Path) -> Path:
|
||||
"""
|
||||
Downloads a model from HuggingFace Hub or Google Cloud Storage.
|
||||
|
||||
@@ -205,7 +94,6 @@ 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",
|
||||
}
|
||||
}
|
||||
@@ -217,20 +105,11 @@ 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),
|
||||
extra_patterns=extra_patterns,
|
||||
local_files_only=kwargs.get("local_files_only", False),
|
||||
)
|
||||
cls.download_files_from_huggingface(hf_source, cache_dir=str(cache_dir))
|
||||
)
|
||||
except (EnvironmentError, RepositoryNotFoundError, ValueError) as e:
|
||||
logger.error(
|
||||
@@ -238,7 +117,4 @@ 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.")
|
||||
|
||||
@@ -1,33 +1,19 @@
|
||||
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,
|
||||
Sequence,
|
||||
)
|
||||
from typing import Any, Dict, Generic, Iterable, List, Optional, Tuple, Type, TypeVar, Union
|
||||
|
||||
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
|
||||
@@ -48,35 +34,23 @@ class OnnxModel(Generic[T]):
|
||||
"""
|
||||
return onnx_input
|
||||
|
||||
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
|
||||
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"])
|
||||
|
||||
# List of Execution Providers: https://onnxruntime.ai/docs/execution-providers
|
||||
|
||||
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}"
|
||||
)
|
||||
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)
|
||||
self.tokenizer = load_tokenizer(model_dir=model_dir, max_length=max_length)
|
||||
self.model = ort.InferenceSession(
|
||||
str(model_path), providers=onnx_providers, sess_options=so
|
||||
)
|
||||
|
||||
@@ -32,7 +32,6 @@ 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,6 +1,5 @@
|
||||
from typing import List, Type, Dict, Any, Union, Iterable, Optional, Sequence
|
||||
from typing import List, Type, Dict, Any, Union, Iterable, Optional
|
||||
|
||||
from fastembed.common import OnnxProvider
|
||||
from fastembed.sparse.sparse_embedding_base import SparseTextEmbeddingBase, SparseEmbedding
|
||||
from fastembed.sparse.splade_pp import SpladePP
|
||||
|
||||
@@ -43,7 +42,6 @@ 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)
|
||||
@@ -51,9 +49,7 @@ 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, providers=providers, **kwargs
|
||||
)
|
||||
self.model = EMBEDDING_MODEL_TYPE(model_name, cache_dir, threads, **kwargs)
|
||||
return
|
||||
|
||||
raise ValueError(
|
||||
|
||||
@@ -1,8 +1,8 @@
|
||||
from typing import Any, Dict, Iterable, List, Optional, Tuple, Union, Type, Sequence
|
||||
from typing import Any, Dict, Iterable, List, Optional, Tuple, Union, Type
|
||||
|
||||
import numpy as np
|
||||
|
||||
from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxProvider
|
||||
from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel
|
||||
from fastembed.common.utils import define_cache_dir
|
||||
from fastembed.sparse.sparse_embedding_base import SparseEmbedding, SparseTextEmbeddingBase
|
||||
|
||||
@@ -15,7 +15,6 @@ supported_splade_models = [
|
||||
"sources": {
|
||||
"hf": "Qdrant/SPLADE_PP_en_v1",
|
||||
},
|
||||
"model_file": "model.onnx",
|
||||
},
|
||||
{
|
||||
"model": "prithivida/Splade_PP_en_v1",
|
||||
@@ -25,7 +24,6 @@ supported_splade_models = [
|
||||
"sources": {
|
||||
"hf": "Qdrant/SPLADE_PP_en_v1",
|
||||
},
|
||||
"model_file": "model.onnx",
|
||||
},
|
||||
]
|
||||
|
||||
@@ -63,7 +61,6 @@ class SpladePP(SparseTextEmbeddingBase, OnnxModel[SparseEmbedding]):
|
||||
model_name: str,
|
||||
cache_dir: Optional[str] = None,
|
||||
threads: Optional[int] = None,
|
||||
providers: Optional[Sequence[OnnxProvider]] = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
@@ -80,19 +77,14 @@ class SpladePP(SparseTextEmbeddingBase, OnnxModel[SparseEmbedding]):
|
||||
|
||||
super().__init__(model_name, cache_dir, threads, **kwargs)
|
||||
|
||||
model_description = self._get_model_description(model_name)
|
||||
cache_dir = define_cache_dir(cache_dir)
|
||||
self.model_name = model_name
|
||||
self._model_description = self._get_model_description(model_name)
|
||||
|
||||
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(
|
||||
model_dir=model_dir,
|
||||
model_file=model_description["model_file"],
|
||||
threads=threads,
|
||||
providers=providers,
|
||||
)
|
||||
self.load_onnx_model(self._model_dir, self.threads, self._max_length)
|
||||
|
||||
def embed(
|
||||
self,
|
||||
@@ -118,7 +110,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,11 +12,8 @@ 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",
|
||||
@@ -26,7 +23,6 @@ supported_multilingual_e5_models = [
|
||||
"sources": {
|
||||
"hf": "xenova/paraphrase-multilingual-mpnet-base-v2",
|
||||
},
|
||||
"model_file": "onnx/model.onnx",
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
@@ -13,7 +13,6 @@ 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",
|
||||
@@ -21,7 +20,6 @@ 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, Sequence
|
||||
from typing import Dict, Optional, Tuple, Union, Iterable, Type, List, Any
|
||||
|
||||
import numpy as np
|
||||
|
||||
from fastembed.common.onnx_model import OnnxModel, EmbeddingWorker, OnnxProvider
|
||||
from fastembed.common.onnx_model import OnnxModel, EmbeddingWorker
|
||||
from fastembed.common.models import normalize
|
||||
from fastembed.common.utils import define_cache_dir
|
||||
from fastembed.text.text_embedding_base import TextEmbeddingBase
|
||||
@@ -14,9 +14,8 @@ supported_onnx_models = [
|
||||
"description": "Base English model",
|
||||
"size_in_GB": 0.42,
|
||||
"sources": {
|
||||
"url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en.tar.gz",
|
||||
"hf": "yashvardhan7/bge-base-en-onnx",
|
||||
},
|
||||
"model_file": "model_optimized.onnx",
|
||||
},
|
||||
{
|
||||
"model": "BAAI/bge-base-en-v1.5",
|
||||
@@ -24,10 +23,8 @@ 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",
|
||||
@@ -37,7 +34,6 @@ supported_onnx_models = [
|
||||
"sources": {
|
||||
"hf": "qdrant/bge-large-en-v1.5-onnx",
|
||||
},
|
||||
"model_file": "model.onnx",
|
||||
},
|
||||
{
|
||||
"model": "BAAI/bge-small-en",
|
||||
@@ -45,9 +41,8 @@ supported_onnx_models = [
|
||||
"description": "Fast English model",
|
||||
"size_in_GB": 0.13,
|
||||
"sources": {
|
||||
"url": "https://storage.googleapis.com/qdrant-fastembed/BAAI-bge-small-en.tar.gz",
|
||||
"hf": "ggrn/bge-small-en",
|
||||
},
|
||||
"model_file": "model_optimized.onnx",
|
||||
},
|
||||
{
|
||||
"model": "BAAI/bge-small-en-v1.5",
|
||||
@@ -57,7 +52,6 @@ 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",
|
||||
@@ -65,9 +59,8 @@ supported_onnx_models = [
|
||||
"description": "Fast and recommended Chinese model",
|
||||
"size_in_GB": 0.09,
|
||||
"sources": {
|
||||
"url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-zh-v1.5.tar.gz",
|
||||
"hf": "Xenova/bge-small-zh-v1.5",
|
||||
},
|
||||
"model_file": "model_optimized.onnx",
|
||||
},
|
||||
{
|
||||
"model": "sentence-transformers/all-MiniLM-L6-v2",
|
||||
@@ -75,10 +68,8 @@ 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",
|
||||
@@ -88,7 +79,6 @@ 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",
|
||||
@@ -98,7 +88,6 @@ supported_onnx_models = [
|
||||
"sources": {
|
||||
"hf": "nomic-ai/nomic-embed-text-v1",
|
||||
},
|
||||
"model_file": "onnx/model.onnx",
|
||||
},
|
||||
{
|
||||
"model": "nomic-ai/nomic-embed-text-v1.5",
|
||||
@@ -108,17 +97,6 @@ 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",
|
||||
@@ -128,7 +106,6 @@ supported_onnx_models = [
|
||||
"sources": {
|
||||
"hf": "qdrant/gte-large-onnx",
|
||||
},
|
||||
"model_file": "model.onnx",
|
||||
},
|
||||
{
|
||||
"model": "mixedbread-ai/mxbai-embed-large-v1",
|
||||
@@ -138,57 +115,6 @@ 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",
|
||||
},
|
||||
]
|
||||
|
||||
@@ -211,7 +137,6 @@ 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,
|
||||
):
|
||||
"""
|
||||
@@ -228,18 +153,14 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxModel[np.ndarray]):
|
||||
|
||||
super().__init__(model_name, cache_dir, threads, **kwargs)
|
||||
|
||||
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.model_name = model_name
|
||||
self._model_description = self._get_model_description(model_name)
|
||||
|
||||
self.load_onnx_model(
|
||||
model_dir=model_dir,
|
||||
model_file=model_description["model_file"],
|
||||
threads=threads,
|
||||
providers=providers,
|
||||
)
|
||||
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)
|
||||
|
||||
def embed(
|
||||
self,
|
||||
@@ -265,7 +186,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,8 +1,7 @@
|
||||
from typing import Any, Dict, Iterable, List, Optional, Type, Union, Sequence
|
||||
from typing import Any, Dict, Iterable, List, Optional, Type, Union
|
||||
|
||||
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
|
||||
@@ -50,7 +49,6 @@ 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)
|
||||
@@ -58,9 +56,7 @@ 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, providers=providers, **kwargs
|
||||
)
|
||||
self.model = EMBEDDING_MODEL_TYPE(model_name, cache_dir, threads, **kwargs)
|
||||
return
|
||||
|
||||
raise ValueError(
|
||||
|
||||
@@ -16,7 +16,6 @@ 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,
|
||||
|
||||
Generated
+3449
File diff suppressed because it is too large
Load Diff
+2
-3
@@ -1,6 +1,6 @@
|
||||
[tool.poetry]
|
||||
name = "fastembed"
|
||||
version = "0.2.7"
|
||||
version = "0.2.6"
|
||||
description = "Fast, light, accurate library built for retrieval embedding generation"
|
||||
authors = ["NirantK <nirant.bits@gmail.com>"]
|
||||
license = "Apache License"
|
||||
@@ -15,8 +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"
|
||||
tokenizers = "^0.15.1"
|
||||
huggingface-hub = "^0.20"
|
||||
loguru = "^0.7.2"
|
||||
numpy = [
|
||||
|
||||
@@ -26,26 +26,16 @@ 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_ci = os.getenv("CI")
|
||||
is_ubuntu_ci = os.getenv("IS_UBUNTU_CI")
|
||||
|
||||
for model_desc in TextEmbedding.list_supported_models():
|
||||
if not is_ci and model_desc["size_in_GB"] > 1:
|
||||
if is_ubuntu_ci == "false" and model_desc["size_in_GB"] > 1:
|
||||
continue
|
||||
|
||||
dim = model_desc["dim"]
|
||||
|
||||
Reference in New Issue
Block a user