From ab7a99a7488bbada825db8b7cdbeb4e99ff6b25c Mon Sep 17 00:00:00 2001 From: Anush Date: Fri, 26 Apr 2024 20:57:16 +0530 Subject: [PATCH] feat: Quantized models (#201) * feat: Quantized models * refactor: use model_file for GCS * refactoring: refactor model downloading (#209) * refactoring: refactor model downloading * refactor: update docstring Co-authored-by: Anush * Update fastembed/common/model_management.py Co-authored-by: George * fix: model_file for Snowflake models --------- Co-authored-by: George --- fastembed/common/model_management.py | 48 +++++++++---------- fastembed/common/onnx_model.py | 13 +++-- fastembed/sparse/splade_pp.py | 18 ++++--- fastembed/text/e5_onnx_embedding.py | 3 ++ fastembed/text/jina_onnx_embedding.py | 2 + fastembed/text/onnx_embedding.py | 68 ++++++++++++++------------- tests/test_text_onnx_embeddings.py | 3 ++ 7 files changed, 85 insertions(+), 70 deletions(-) diff --git a/fastembed/common/model_management.py b/fastembed/common/model_management.py index 6ccce32..f7a75be 100644 --- a/fastembed/common/model_management.py +++ b/fastembed/common/model_management.py @@ -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( diff --git a/fastembed/common/onnx_model.py b/fastembed/common/onnx_model.py index 190393f..9ddd50b 100644 --- a/fastembed/common/onnx_model.py +++ b/fastembed/common/onnx_model.py @@ -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 ) diff --git a/fastembed/sparse/splade_pp.py b/fastembed/sparse/splade_pp.py index 92f3da8..f59114c 100644 --- a/fastembed/sparse/splade_pp.py +++ b/fastembed/sparse/splade_pp.py @@ -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, diff --git a/fastembed/text/e5_onnx_embedding.py b/fastembed/text/e5_onnx_embedding.py index 9c37174..8df199d 100644 --- a/fastembed/text/e5_onnx_embedding.py +++ b/fastembed/text/e5_onnx_embedding.py @@ -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", }, ] diff --git a/fastembed/text/jina_onnx_embedding.py b/fastembed/text/jina_onnx_embedding.py index 4d2ed96..c0e7acf 100644 --- a/fastembed/text/jina_onnx_embedding.py +++ b/fastembed/text/jina_onnx_embedding.py @@ -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", }, ] diff --git a/fastembed/text/onnx_embedding.py b/fastembed/text/onnx_embedding.py index 190c975..059cce2 100644 --- a/fastembed/text/onnx_embedding.py +++ b/fastembed/text/onnx_embedding.py @@ -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, diff --git a/tests/test_text_onnx_embeddings.py b/tests/test_text_onnx_embeddings.py index da7387f..3ea7685 100644 --- a/tests/test_text_onnx_embeddings.py +++ b/tests/test_text_onnx_embeddings.py @@ -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]),