upstream-archive/fastembed
1
0
Fork 0
mirror of https://github.com/qdrant/fastembed.git synced 2026-07-23 11:20:51 -05:00
Code Issues Packages Projects Releases Wiki Activity

feat: Quantized models (#201)

Browse Source 
* 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>
This commit is contained in:
Anush
2024-04-26 20:57:16 +05:30
committed by GitHub
parent 466886a317
commit ab7a99a748
7 changed files with 85 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files

48
fastembed/common/model_management.py
View File

@@ -11,23 +11,6 @@ from tqdm import tqdm
from loguru import logger
def locate_model_file(model_dir: Path, file_names: List[str]) -> Path:
"""
Find model path for both TransformerJS style `onnx` subdirectory structure and direct model weights structure used
by Optimum and Qdrant
"""
if not model_dir.is_dir():
raise ValueError(f"Provided model path '{model_dir}' is not a directory.")
for file_name in file_names:
file_paths = [path for path in model_dir.rglob(file_name) if path.is_file()]
if file_paths:
return file_paths[0]
raise ValueError(f"Could not find either of {', '.join(file_names)} in {model_dir}")
class ModelManagement:
@classmethod
def list_supported_models(cls) -> List[Dict[str, Any]]:
@@ -104,27 +87,33 @@ class ModelManagement:
@classmethod
def download_files_from_huggingface(
cls, hf_source_repo: str, cache_dir: Optional[str] = None
cls,
hf_source_repo: str,
cache_dir: Optional[str] = None,
extra_patterns: Optional[List[str]] = None,
) -> str:
"""
Downloads a model from HuggingFace Hub.
Args:
hf_source_repo (str): Name of the model on HuggingFace Hub, e.g. "qdrant/all-MiniLM-L6-v2-onnx".
cache_dir (Optional[str]): The path to the cache directory.
extra_patterns (Optional[List[str]]): extra patterns to allow in the snapshot download, typically
includes the required model files.
Returns:
Path: The path to the model directory.
"""
allow_patterns = [
"config.json",
"tokenizer.json",
"tokenizer_config.json",
"special_tokens_map.json",
]
if extra_patterns is not None:
allow_patterns.extend(extra_patterns)
return snapshot_download(
repo_id=hf_source_repo,
allow_patterns=[
"*.onnx",
"*.onnx_data",
"config.json",
"tokenizer.json",
"tokenizer_config.json",
"special_tokens_map.json",
],
allow_patterns=allow_patterns,
cache_dir=cache_dir,
)
@@ -229,9 +218,14 @@ class ModelManagement:
url_source = model.get("sources", {}).get("url")
if hf_source:
extra_patterns = [model["model_file"]]
extra_patterns.extend(model.get("additional_files", []))
try:
return Path(
cls.download_files_from_huggingface(hf_source, cache_dir=str(cache_dir))
cls.download_files_from_huggingface(
hf_source, cache_dir=str(cache_dir), extra_patterns=extra_patterns
)
)
except (EnvironmentError, RepositoryNotFoundError, ValueError) as e:
logger.error(

13
fastembed/common/onnx_model.py
View File

@@ -6,11 +6,11 @@ from typing import Any, Dict, Generic, Iterable, List, Optional, Tuple, Type, Ty
import numpy as np
import onnxruntime as ort
from fastembed.common.model_management import locate_model_file
from fastembed.common.models import load_tokenizer
from fastembed.common.utils import iter_batch
from fastembed.parallel_processor import ParallelWorkerPool, Worker
# Holds type of the embedding result
T = TypeVar("T")
@@ -34,8 +34,13 @@ class OnnxModel(Generic[T]):
"""
return onnx_input
def load_onnx_model(self, model_dir: Path, threads: Optional[int], max_length: int) -> None:
model_path = locate_model_file(model_dir, ["model.onnx", "model_optimized.onnx"])
def load_onnx_model(
self,
model_dir: Path,
model_file: str,
threads: Optional[int],
) -> None:
model_path = model_dir / model_file
# List of Execution Providers: https://onnxruntime.ai/docs/execution-providers
onnx_providers = ["CPUExecutionProvider"]
@@ -50,7 +55,7 @@ class OnnxModel(Generic[T]):
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
)

18
fastembed/sparse/splade_pp.py
View File

@@ -15,6 +15,7 @@ supported_splade_models = [
"sources": {
"hf": "Qdrant/SPLADE_PP_en_v1",
},
"model_file": "model.onnx",
},
{
"model": "prithivida/Splade_PP_en_v1",
@@ -24,6 +25,7 @@ supported_splade_models = [
"sources": {
"hf": "Qdrant/SPLADE_PP_en_v1",
},
"model_file": "model.onnx",
},
]
@@ -77,14 +79,16 @@ class SpladePP(SparseTextEmbeddingBase, OnnxModel[SparseEmbedding]):
super().__init__(model_name, cache_dir, threads, **kwargs)
self.model_name = model_name
self._model_description = self._get_model_description(model_name)
model_description = self._get_model_description(model_name)
cache_dir = define_cache_dir(cache_dir)
self._cache_dir = define_cache_dir(cache_dir)
self._model_dir = self.download_model(self._model_description, self._cache_dir)
self._max_length = 512
model_dir = self.download_model(model_description, cache_dir)
self.load_onnx_model(self._model_dir, self.threads, self._max_length)
self.load_onnx_model(
model_dir=model_dir,
model_file=model_description["model_file"],
threads=threads,
)
def embed(
self,
@@ -110,7 +114,7 @@ class SpladePP(SparseTextEmbeddingBase, OnnxModel[SparseEmbedding]):
"""
yield from self._embed_documents(
model_name=self.model_name,
cache_dir=str(self._cache_dir),
cache_dir=str(self.cache_dir),
documents=documents,
batch_size=batch_size,
parallel=parallel,

3
fastembed/text/e5_onnx_embedding.py
View File

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

2
fastembed/text/jina_onnx_embedding.py
View File

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

68
fastembed/text/onnx_embedding.py
View File

@@ -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,7 @@ supported_onnx_models = [
"sources": {
"hf": "mixedbread-ai/mxbai-embed-large-v1",
},
"model_file": "onnx/model.onnx",
},
{
"model": "snowflake/snowflake-arctic-embed-xs",
@@ -150,6 +148,7 @@ supported_onnx_models = [
"sources": {
"hf": "snowflake/snowflake-arctic-embed-xs",
},
"model_file": "onnx/model.onnx",
},
{
"model": "snowflake/snowflake-arctic-embed-s",
@@ -159,6 +158,7 @@ supported_onnx_models = [
"sources": {
"hf": "snowflake/snowflake-arctic-embed-s",
},
"model_file": "onnx/model.onnx",
},
{
"model": "snowflake/snowflake-arctic-embed-m",
@@ -168,6 +168,7 @@ supported_onnx_models = [
"sources": {
"hf": "Snowflake/snowflake-arctic-embed-m",
},
"model_file": "onnx/model.onnx",
},
{
"model": "snowflake/snowflake-arctic-embed-m-long",
@@ -177,6 +178,7 @@ supported_onnx_models = [
"sources": {
"hf": "snowflake/snowflake-arctic-embed-m-long",
},
"model_file": "onnx/model.onnx",
},
{
"model": "snowflake/snowflake-arctic-embed-l",
@@ -186,6 +188,7 @@ supported_onnx_models = [
"sources": {
"hf": "snowflake/snowflake-arctic-embed-l",
},
"model_file": "onnx/model.onnx",
},
]
@@ -224,14 +227,15 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxModel[np.ndarray]):
super().__init__(model_name, cache_dir, threads, **kwargs)
self.model_name = model_name
self._model_description = self._get_model_description(model_name)
model_description = self._get_model_description(model_name)
cache_dir = define_cache_dir(cache_dir)
model_dir = self.download_model(model_description, cache_dir)
self._cache_dir = define_cache_dir(cache_dir)
self._model_dir = self.download_model(self._model_description, self._cache_dir)
self._max_length = 512
self.load_onnx_model(self._model_dir, self.threads, self._max_length)
self.load_onnx_model(
model_dir=model_dir,
model_file=model_description["model_file"],
threads=threads,
)
def embed(
self,
@@ -257,7 +261,7 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxModel[np.ndarray]):
"""
yield from self._embed_documents(
model_name=self.model_name,
cache_dir=str(self._cache_dir),
cache_dir=str(self.cache_dir),
documents=documents,
batch_size=batch_size,
parallel=parallel,

3
tests/test_text_onnx_embeddings.py
View File

@@ -26,6 +26,9 @@ 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]),