mirror of
https://github.com/qdrant/fastembed.git
synced 2026-09-22 05:57:51 -05:00
Compare commits
6
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 Anush008
|
24dd472185 | ||
|
Anush008
|
221b554228 | ||
|
Anush
|
fff693b46d | ||
|
Anush008
|
1dbef6b239 | ||
|
Anush
|
ff82bf7174 | ||
|
Anush008
|
c0e457305a |
@@ -3,8 +3,6 @@ name: Tests
|
||||
on:
|
||||
push:
|
||||
branches: [ master, main ]
|
||||
schedule:
|
||||
- cron: 0 0 * * *
|
||||
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
|
||||
|
||||
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,14 +1,4 @@
|
||||
import os
|
||||
import shutil
|
||||
import tarfile
|
||||
from pathlib import Path
|
||||
from typing import List, Optional, Dict, Any
|
||||
|
||||
import requests
|
||||
from huggingface_hub import snapshot_download
|
||||
from huggingface_hub.utils import RepositoryNotFoundError
|
||||
from tqdm import tqdm
|
||||
from loguru import logger
|
||||
from typing import List, Dict, Any
|
||||
|
||||
|
||||
class ModelManagement:
|
||||
@@ -40,200 +30,3 @@ class ModelManagement:
|
||||
return model
|
||||
|
||||
raise ValueError(f"Model {model_name} is not supported in {cls.__name__}.")
|
||||
|
||||
@classmethod
|
||||
def download_file_from_gcs(cls, url: str, output_path: str, show_progress: bool = True) -> str:
|
||||
"""
|
||||
Downloads a file from Google Cloud Storage.
|
||||
|
||||
Args:
|
||||
url (str): The URL to download the file from.
|
||||
output_path (str): The path to save the downloaded file to.
|
||||
show_progress (bool, optional): Whether to show a progress bar. Defaults to True.
|
||||
|
||||
Returns:
|
||||
str: The path to the downloaded file.
|
||||
"""
|
||||
|
||||
if os.path.exists(output_path):
|
||||
return output_path
|
||||
response = requests.get(url, stream=True)
|
||||
|
||||
# Handle HTTP errors
|
||||
if response.status_code == 403:
|
||||
raise PermissionError(
|
||||
"Authentication Error: You do not have permission to access this resource. "
|
||||
"Please check your credentials."
|
||||
)
|
||||
|
||||
# Get the total size of the file
|
||||
total_size_in_bytes = int(response.headers.get("content-length", 0))
|
||||
|
||||
# Warn if the total size is zero
|
||||
if total_size_in_bytes == 0:
|
||||
print(f"Warning: Content-length header is missing or zero in the response from {url}.")
|
||||
|
||||
show_progress = total_size_in_bytes and show_progress
|
||||
|
||||
with tqdm(
|
||||
total=total_size_in_bytes, unit="iB", unit_scale=True, disable=not show_progress
|
||||
) as progress_bar:
|
||||
with open(output_path, "wb") as file:
|
||||
for chunk in response.iter_content(chunk_size=1024):
|
||||
if chunk: # Filter out keep-alive new chunks
|
||||
progress_bar.update(len(chunk))
|
||||
file.write(chunk)
|
||||
return output_path
|
||||
|
||||
@classmethod
|
||||
def download_files_from_huggingface(
|
||||
cls,
|
||||
hf_source_repo: str,
|
||||
cache_dir: Optional[str] = None,
|
||||
extra_patterns: Optional[List[str]] = None,
|
||||
) -> str:
|
||||
"""
|
||||
Downloads a model from HuggingFace Hub.
|
||||
Args:
|
||||
hf_source_repo (str): Name of the model on HuggingFace Hub, e.g. "qdrant/all-MiniLM-L6-v2-onnx".
|
||||
cache_dir (Optional[str]): The path to the cache directory.
|
||||
extra_patterns (Optional[List[str]]): extra patterns to allow in the snapshot download, typically
|
||||
includes the required model files.
|
||||
Returns:
|
||||
Path: The path to the model directory.
|
||||
"""
|
||||
allow_patterns = [
|
||||
"config.json",
|
||||
"tokenizer.json",
|
||||
"tokenizer_config.json",
|
||||
"special_tokens_map.json",
|
||||
]
|
||||
if extra_patterns is not None:
|
||||
allow_patterns.extend(extra_patterns)
|
||||
|
||||
return snapshot_download(
|
||||
repo_id=hf_source_repo,
|
||||
allow_patterns=allow_patterns,
|
||||
cache_dir=cache_dir,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def decompress_to_cache(cls, targz_path: str, cache_dir: str):
|
||||
"""
|
||||
Decompresses a .tar.gz file to a cache directory.
|
||||
|
||||
Args:
|
||||
targz_path (str): Path to the .tar.gz file.
|
||||
cache_dir (str): Path to the cache directory.
|
||||
|
||||
Returns:
|
||||
cache_dir (str): Path to the cache directory.
|
||||
"""
|
||||
# Check if targz_path exists and is a file
|
||||
if not os.path.isfile(targz_path):
|
||||
raise ValueError(f"{targz_path} does not exist or is not a file.")
|
||||
|
||||
# Check if targz_path is a .tar.gz file
|
||||
if not targz_path.endswith(".tar.gz"):
|
||||
raise ValueError(f"{targz_path} is not a .tar.gz file.")
|
||||
|
||||
try:
|
||||
# Open the tar.gz file
|
||||
with tarfile.open(targz_path, "r:gz") as tar:
|
||||
# Extract all files into the cache directory
|
||||
tar.extractall(path=cache_dir)
|
||||
except tarfile.TarError as e:
|
||||
# If any error occurs while opening or extracting the tar.gz file,
|
||||
# delete the cache directory (if it was created in this function)
|
||||
# and raise the error again
|
||||
if "tmp" in cache_dir:
|
||||
shutil.rmtree(cache_dir)
|
||||
raise ValueError(f"An error occurred while decompressing {targz_path}: {e}")
|
||||
|
||||
return cache_dir
|
||||
|
||||
@classmethod
|
||||
def retrieve_model_gcs(cls, model_name: str, source_url: str, cache_dir: str) -> Path:
|
||||
fast_model_name = f"fast-{model_name.split('/')[-1]}"
|
||||
|
||||
cache_tmp_dir = Path(cache_dir) / "tmp"
|
||||
model_tmp_dir = cache_tmp_dir / fast_model_name
|
||||
model_dir = Path(cache_dir) / fast_model_name
|
||||
|
||||
# check if the model_dir and the model files are both present for macOS
|
||||
if model_dir.exists() and len(list(model_dir.glob("*"))) > 0:
|
||||
return model_dir
|
||||
|
||||
if model_tmp_dir.exists():
|
||||
shutil.rmtree(model_tmp_dir)
|
||||
|
||||
cache_tmp_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
model_tar_gz = Path(cache_dir) / f"{fast_model_name}.tar.gz"
|
||||
|
||||
if model_tar_gz.exists():
|
||||
model_tar_gz.unlink()
|
||||
|
||||
cls.download_file_from_gcs(
|
||||
source_url,
|
||||
output_path=str(model_tar_gz),
|
||||
)
|
||||
|
||||
cls.decompress_to_cache(targz_path=str(model_tar_gz), cache_dir=str(cache_tmp_dir))
|
||||
assert model_tmp_dir.exists(), f"Could not find {model_tmp_dir} in {cache_tmp_dir}"
|
||||
|
||||
model_tar_gz.unlink()
|
||||
# Rename from tmp to final name is atomic
|
||||
model_tmp_dir.rename(model_dir)
|
||||
|
||||
return model_dir
|
||||
|
||||
@classmethod
|
||||
def download_model(cls, model: Dict[str, Any], cache_dir: Path) -> Path:
|
||||
"""
|
||||
Downloads a model from HuggingFace Hub or Google Cloud Storage.
|
||||
|
||||
Args:
|
||||
model (Dict[str, Any]): The model description.
|
||||
Example:
|
||||
```
|
||||
{
|
||||
"model": "BAAI/bge-base-en-v1.5",
|
||||
"dim": 768,
|
||||
"description": "Base English model, v1.5",
|
||||
"size_in_GB": 0.44,
|
||||
"sources": {
|
||||
"url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en-v1.5.tar.gz",
|
||||
"hf": "qdrant/bge-base-en-v1.5-onnx-q",
|
||||
}
|
||||
}
|
||||
```
|
||||
cache_dir (str): The path to the cache directory.
|
||||
|
||||
Returns:
|
||||
Path: The path to the downloaded model directory.
|
||||
"""
|
||||
|
||||
hf_source = model.get("sources", {}).get("hf")
|
||||
url_source = model.get("sources", {}).get("url")
|
||||
|
||||
if hf_source:
|
||||
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
|
||||
)
|
||||
)
|
||||
except (EnvironmentError, RepositoryNotFoundError, ValueError) as e:
|
||||
logger.error(
|
||||
f"Could not download model from HuggingFace: {e}"
|
||||
"Falling back to other sources."
|
||||
)
|
||||
|
||||
if url_source:
|
||||
return cls.retrieve_model_gcs(model["model"], url_source, str(cache_dir))
|
||||
|
||||
raise ValueError(f"Could not download model {model['model']} from any source.")
|
||||
|
||||
+14
-13
@@ -3,24 +3,25 @@ from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
from tokenizers import Tokenizer, AddedToken
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
|
||||
def load_tokenizer(model_dir: Path, max_length: int = 512) -> Tokenizer:
|
||||
config_path = model_dir / "config.json"
|
||||
if not config_path.exists():
|
||||
raise ValueError(f"Could not find config.json in {model_dir}")
|
||||
def load_tokenizer(repo_id: str, cache_dir: Path, max_length: int = 512) -> Tokenizer:
|
||||
config_path = hf_hub_download(
|
||||
repo_id=repo_id, filename="config.json", cache_dir=str(cache_dir)
|
||||
)
|
||||
|
||||
tokenizer_path = model_dir / "tokenizer.json"
|
||||
if not tokenizer_path.exists():
|
||||
raise ValueError(f"Could not find tokenizer.json in {model_dir}")
|
||||
tokenizer_path = hf_hub_download(
|
||||
repo_id=repo_id, filename="tokenizer.json", cache_dir=str(cache_dir)
|
||||
)
|
||||
|
||||
tokenizer_config_path = model_dir / "tokenizer_config.json"
|
||||
if not tokenizer_config_path.exists():
|
||||
raise ValueError(f"Could not find tokenizer_config.json in {model_dir}")
|
||||
tokenizer_config_path = hf_hub_download(
|
||||
repo_id=repo_id, filename="tokenizer_config.json", cache_dir=str(cache_dir)
|
||||
)
|
||||
|
||||
tokens_map_path = model_dir / "special_tokens_map.json"
|
||||
if not tokens_map_path.exists():
|
||||
raise ValueError(f"Could not find special_tokens_map.json in {model_dir}")
|
||||
tokens_map_path = hf_hub_download(
|
||||
repo_id=repo_id, filename="special_tokens_map.json", cache_dir=str(cache_dir)
|
||||
)
|
||||
|
||||
with open(str(config_path)) as config_file:
|
||||
config = json.load(config_file)
|
||||
|
||||
@@ -1,19 +1,7 @@
|
||||
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
|
||||
@@ -22,6 +10,7 @@ from fastembed.common.models import load_tokenizer
|
||||
from fastembed.common.utils import iter_batch
|
||||
from fastembed.parallel_processor import ParallelWorkerPool, Worker
|
||||
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
# Holds type of the embedding result
|
||||
T = TypeVar("T")
|
||||
@@ -47,34 +36,36 @@ class OnnxModel(Generic[T]):
|
||||
return onnx_input
|
||||
|
||||
def load_onnx_model(
|
||||
self,
|
||||
model_dir: Path,
|
||||
model_file: str,
|
||||
threads: Optional[int],
|
||||
providers: Optional[Sequence[Union[str, Tuple[str, Dict[Any, Any]]]]] = None,
|
||||
self, model_description: dict, threads: Optional[int], cache_dir: Path
|
||||
) -> None:
|
||||
model_path = model_dir / model_file
|
||||
repo_id = model_description["sources"]["hf"]
|
||||
model_file = model_description.get("model_file", "model.onnx")
|
||||
|
||||
# Some models require additional repo files.
|
||||
# For eg: intfloat/multilingual-e5-large requires the model.onnx_data file.
|
||||
# These can be specified within the "additional_files" option when describing the model properties
|
||||
if additional_files := model_description.get("additional_files"):
|
||||
for file in additional_files:
|
||||
hf_hub_download(repo_id=repo_id, filename=file, cache_dir=str(cache_dir))
|
||||
|
||||
model_path = hf_hub_download(
|
||||
repo_id=repo_id, filename=model_file, cache_dir=str(cache_dir)
|
||||
)
|
||||
|
||||
# List of Execution Providers: https://onnxruntime.ai/docs/execution-providers
|
||||
|
||||
onnx_providers = ["CPUExecutionProvider"] 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(repo_id, cache_dir)
|
||||
self.model = ort.InferenceSession(
|
||||
str(model_path), providers=onnx_providers, sess_options=so
|
||||
)
|
||||
|
||||
@@ -79,15 +79,10 @@ 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)
|
||||
|
||||
model_dir = self.download_model(model_description, cache_dir)
|
||||
|
||||
self.load_onnx_model(
|
||||
model_dir=model_dir,
|
||||
model_file=model_description["model_file"],
|
||||
threads=threads,
|
||||
self._get_model_description(model_name),
|
||||
threads,
|
||||
define_cache_dir(cache_dir),
|
||||
)
|
||||
|
||||
def embed(
|
||||
|
||||
@@ -12,7 +12,6 @@ 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",
|
||||
|
||||
@@ -14,9 +14,9 @@ supported_onnx_models = [
|
||||
"description": "Base English model",
|
||||
"size_in_GB": 0.42,
|
||||
"sources": {
|
||||
"url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en.tar.gz",
|
||||
"hf": "yashvardhan7/bge-base-en-onnx",
|
||||
},
|
||||
"model_file": "model_optimized.onnx",
|
||||
"model_file": "onnx/model.onnx",
|
||||
},
|
||||
{
|
||||
"model": "BAAI/bge-base-en-v1.5",
|
||||
@@ -24,7 +24,6 @@ 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",
|
||||
@@ -45,9 +44,9 @@ 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_file": "onnx/model.onnx",
|
||||
},
|
||||
{
|
||||
"model": "BAAI/bge-small-en-v1.5",
|
||||
@@ -65,9 +64,9 @@ supported_onnx_models = [
|
||||
"description": "Fast and recommended Chinese model",
|
||||
"size_in_GB": 0.09,
|
||||
"sources": {
|
||||
"url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-zh-v1.5.tar.gz",
|
||||
"hf": "Xenova/bge-small-zh-v1.5",
|
||||
},
|
||||
"model_file": "model_optimized.onnx",
|
||||
"model_file": "onnx/model.onnx",
|
||||
},
|
||||
{
|
||||
"model": "sentence-transformers/all-MiniLM-L6-v2",
|
||||
@@ -75,7 +74,6 @@ 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",
|
||||
@@ -140,56 +138,6 @@ supported_onnx_models = [
|
||||
},
|
||||
"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",
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
@@ -227,14 +175,10 @@ 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)
|
||||
|
||||
self.load_onnx_model(
|
||||
model_dir=model_dir,
|
||||
model_file=model_description["model_file"],
|
||||
threads=threads,
|
||||
self._get_model_description(model_name),
|
||||
threads,
|
||||
define_cache_dir(cache_dir),
|
||||
)
|
||||
|
||||
def embed(
|
||||
|
||||
Generated
+386
-406
File diff suppressed because it is too large
Load Diff
+2
-13
@@ -12,18 +12,9 @@ keywords = ["vector", "embedding", "neural", "search", "qdrant", "sentence-trans
|
||||
|
||||
[tool.poetry.dependencies]
|
||||
python = ">=3.8.0,<3.13"
|
||||
onnx = [
|
||||
{version = "^1.15.0", optional = true, markers = "extra != 'gpu'"}
|
||||
]
|
||||
onnxruntime = [
|
||||
{version = "^1.17.0", optional = true, markers = "extra != 'gpu'"}
|
||||
]
|
||||
onnxruntime-gpu = [
|
||||
{ version = "^1.17.0", optional = true, python = "<3.13", markers = "extra == 'gpu'" }
|
||||
]
|
||||
|
||||
onnx = "^1.15.0"
|
||||
onnxruntime = "^1.17.0"
|
||||
tqdm = "^4.66"
|
||||
requests = "^2.31"
|
||||
tokenizers = "^0.15.1"
|
||||
huggingface-hub = "^0.20"
|
||||
loguru = "^0.7.2"
|
||||
@@ -45,8 +36,6 @@ pillow = "^10.2.0"
|
||||
cairosvg = "^2.7.1"
|
||||
mknotebooks = "^0.8.0"
|
||||
|
||||
[tool.poetry.extras]
|
||||
gpu = ["onnxruntime-gpu"]
|
||||
|
||||
[build-system]
|
||||
requires = ["poetry-core"]
|
||||
|
||||
@@ -26,26 +26,19 @@ 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]
|
||||
"nomic-ai/nomic-embed-text-v1.5-Q": np.array(
|
||||
[-0.01554983, 0.0129992, -0.17909265, -0.01062993, 0.00512859]
|
||||
),
|
||||
"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