fastembed 0.2.4__tar.gz → 0.2.6__tar.gz

{fastembed-0.2.4 → fastembed-0.2.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: fastembed
-Version: 0.2.4
+Version: 0.2.6
 Summary: Fast, light, accurate library built for retrieval embedding generation
 Home-page: https://github.com/qdrant/fastembed
 License: Apache License
@@ -31,7 +31,7 @@ Description-Content-Type: text/markdown
 
 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, the top model in the [MTEB](https://huggingface.co/spaces/mteb/leaderboard) leaderboard. It supports "query" and "passage" prefixes for the input text. Here is an example for [Retrieval Embedding Generation](https://qdrant.github.io/fastembed/examples/Retrieval_with_FastEmbed/) and how to use [FastEmbed with Qdrant](https://qdrant.github.io/fastembed/examples/Usage_With_Qdrant/).
+The default text embedding (`TextEmbedding`) model is Flag Embedding, presented in the [MTEB](https://huggingface.co/spaces/mteb/leaderboard) leaderboard. It supports "query" and "passage" prefixes for the input text. Here is an example for [Retrieval Embedding Generation](https://qdrant.github.io/fastembed/examples/Retrieval_with_FastEmbed/) and how to use [FastEmbed with Qdrant](https://qdrant.github.io/fastembed/examples/Usage_With_Qdrant/).
 
 ## 📈 Why FastEmbed?
 
@@ -52,7 +52,6 @@ pip install fastembed
 ## 📖 Quickstart
 
 ```python
-import numpy as np
 from fastembed import TextEmbedding
 from typing import List
 

{fastembed-0.2.4 → fastembed-0.2.6}/README.md

@@ -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, the top model in the [MTEB](https://huggingface.co/spaces/mteb/leaderboard) leaderboard. It supports "query" and "passage" prefixes for the input text. Here is an example for [Retrieval Embedding Generation](https://qdrant.github.io/fastembed/examples/Retrieval_with_FastEmbed/) and how to use [FastEmbed with Qdrant](https://qdrant.github.io/fastembed/examples/Usage_With_Qdrant/).
+The default text embedding (`TextEmbedding`) model is Flag Embedding, presented in the [MTEB](https://huggingface.co/spaces/mteb/leaderboard) leaderboard. It supports "query" and "passage" prefixes for the input text. Here is an example for [Retrieval Embedding Generation](https://qdrant.github.io/fastembed/examples/Retrieval_with_FastEmbed/) and how to use [FastEmbed with Qdrant](https://qdrant.github.io/fastembed/examples/Usage_With_Qdrant/).
 
 ## 📈 Why FastEmbed?
 
@@ -23,7 +23,6 @@ pip install fastembed
 ## 📖 Quickstart
 
 ```python
-import numpy as np
 from fastembed import TextEmbedding
 from typing import List
 

fastembed-0.2.6/fastembed/__init__.py

@@ -0,0 +1,7 @@
+import importlib.metadata
+
+from fastembed.text import TextEmbedding
+from fastembed.sparse import SparseTextEmbedding, SparseEmbedding
+
+__version__ = importlib.metadata.version("fastembed")
+__all__ = ["TextEmbedding", "SparseTextEmbedding", "SparseEmbedding"]

{fastembed-0.2.4 → fastembed-0.2.6}/fastembed/common/model_management.py

@@ -53,7 +53,7 @@ class ModelManagement:
             Dict[str, Any]: The model description.
         """
         for model in cls.list_supported_models():
-            if model_name == model["model"]:
+            if model_name.lower() == model["model"].lower():
                 return model
 
         raise ValueError(f"Model {model_name} is not supported in {cls.__name__}.")
@@ -92,7 +92,9 @@ class ModelManagement:
 
         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 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
@@ -101,7 +103,9 @@ class ModelManagement:
         return output_path
 
     @classmethod
-    def download_files_from_huggingface(cls, hf_source_repo: str, cache_dir: Optional[str] = None) -> str:
+    def download_files_from_huggingface(
+        cls, hf_source_repo: str, cache_dir: Optional[str] = None
+    ) -> str:
         """
         Downloads a model from HuggingFace Hub.
         Args:
@@ -216,9 +220,14 @@ class ModelManagement:
 
         if hf_source:
             try:
-                return Path(cls.download_files_from_huggingface(hf_source, cache_dir=str(cache_dir)))
+                return Path(
+                    cls.download_files_from_huggingface(hf_source, cache_dir=str(cache_dir))
+                )
             except (EnvironmentError, RepositoryNotFoundError, ValueError) as e:
-                logger.error(f"Could not download model from HuggingFace: {e}" "Falling back to other sources.")
+                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))

{fastembed-0.2.4 → fastembed-0.2.6}/fastembed/common/models.py

@@ -33,7 +33,9 @@ def load_tokenizer(model_dir: Path, max_length: int = 512) -> Tokenizer:
 
     tokenizer = Tokenizer.from_file(str(tokenizer_path))
     tokenizer.enable_truncation(max_length=min(tokenizer_config["model_max_length"], max_length))
-    tokenizer.enable_padding(pad_id=config.get("pad_token_id", 0), pad_token=tokenizer_config["pad_token"])
+    tokenizer.enable_padding(
+        pad_id=config.get("pad_token_id", 0), pad_token=tokenizer_config["pad_token"]
+    )
 
     for token in tokens_map.values():
         if isinstance(token, str):

{fastembed-0.2.4 → fastembed-0.2.6}/fastembed/common/onnx_model.py

@@ -48,7 +48,9 @@ class OnnxModel(Generic[T]):
             so.inter_op_num_threads = threads
 
         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)
+        self.model = ort.InferenceSession(
+            str(model_path), providers=onnx_providers, sess_options=so
+        )
 
     def onnx_embed(self, documents: List[str]) -> Tuple[np.ndarray, np.ndarray]:
         encoded = self.tokenizer.encode_batch(documents)
@@ -58,7 +60,9 @@ class OnnxModel(Generic[T]):
         onnx_input = {
             "input_ids": np.array(input_ids, dtype=np.int64),
             "attention_mask": np.array(attention_mask, dtype=np.int64),
-            "token_type_ids": np.array([np.zeros(len(e), dtype=np.int64) for e in input_ids], dtype=np.int64),
+            "token_type_ids": np.array(
+                [np.zeros(len(e), dtype=np.int64) for e in input_ids], dtype=np.int64
+            ),
         }
 
         onnx_input = self._preprocess_onnx_input(onnx_input)
@@ -97,7 +101,9 @@ class OnnxModel(Generic[T]):
                 "model_name": model_name,
                 "cache_dir": cache_dir,
             }
-            pool = ParallelWorkerPool(parallel, self._get_worker_class(), start_method=start_method)
+            pool = ParallelWorkerPool(
+                parallel, self._get_worker_class(), start_method=start_method
+            )
             for batch in pool.ordered_map(iter_batch(documents, batch_size), **params):
                 yield from self._post_process_onnx_output(batch)
 

{fastembed-0.2.4 → fastembed-0.2.6}/fastembed/embedding.py

@@ -2,10 +2,11 @@ from typing import Optional
 
 from loguru import logger
 
-from fastembed.text.text_embedding import TextEmbedding
+from fastembed import TextEmbedding
 
 logger.warning(
-    "DefaultEmbedding, FlagEmbedding, JinaEmbedding are deprecated." "Use from fastembed import TextEmbedding instead."
+    "DefaultEmbedding, FlagEmbedding, JinaEmbedding are deprecated."
+    "Use from fastembed import TextEmbedding instead."
 )
 
 DefaultEmbedding = TextEmbedding

{fastembed-0.2.4 → fastembed-0.2.6}/fastembed/parallel_processor.py

@@ -128,7 +128,9 @@ class ParallelWorkerPool:
                 yield buffer.pop(next_expected)
                 next_expected += 1
 
-    def semi_ordered_map(self, stream: Iterable[Any], *args: Any, **kwargs: Any) -> Iterable[Tuple[int, Any]]:
+    def semi_ordered_map(
+        self, stream: Iterable[Any], *args: Any, **kwargs: Any
+    ) -> Iterable[Tuple[int, Any]]:
         try:
             self.start(**kwargs)
 

fastembed-0.2.6/fastembed/sparse/__init__.py

@@ -0,0 +1,4 @@
+from fastembed.sparse.sparse_embedding_base import SparseEmbedding
+from fastembed.sparse.sparse_text_embedding import SparseTextEmbedding
+
+__all__ = ["SparseEmbedding", "SparseTextEmbedding"]

{fastembed-0.2.4 → fastembed-0.2.6}/fastembed/sparse/sparse_embedding_base.py

@@ -22,7 +22,13 @@ class SparseEmbedding:
 
 
 class SparseTextEmbeddingBase(ModelManagement):
-    def __init__(self, model_name: str, cache_dir: Optional[str] = None, threads: Optional[int] = None, **kwargs):
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: Optional[str] = None,
+        threads: Optional[int] = None,
+        **kwargs,
+    ):
         self.model_name = model_name
         self.cache_dir = cache_dir
         self.threads = threads

{fastembed-0.2.4 → fastembed-0.2.6}/fastembed/sparse/sparse_text_embedding.py

@@ -48,7 +48,7 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
 
         for EMBEDDING_MODEL_TYPE in self.EMBEDDINGS_REGISTRY:
             supported_models = EMBEDDING_MODEL_TYPE.list_supported_models()
-            if any(model_name == model["model"] for model in supported_models):
+            if any(model_name.lower() == model["model"].lower() for model in supported_models):
                 self.model = EMBEDDING_MODEL_TYPE(model_name, cache_dir, threads, **kwargs)
                 return
 

{fastembed-0.2.4 → fastembed-0.2.6}/fastembed/sparse/splade_pp.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
+from typing import Any, Dict, Iterable, List, Optional, Tuple, Union, Type
 
 import numpy as np
 
@@ -8,7 +8,16 @@ from fastembed.sparse.sparse_embedding_base import SparseEmbedding, SparseTextEm
 
 supported_splade_models = [
     {
-        "model": "prithvida/SPLADE_PP_en_v1",
+        "model": "prithvida/Splade_PP_en_v1",
+        "vocab_size": 30522,
+        "description": "Misspelled version of the model. Retained for backward compatibility. Independent Implementation of SPLADE++ Model for English",
+        "size_in_GB": 0.532,
+        "sources": {
+            "hf": "Qdrant/SPLADE_PP_en_v1",
+        },
+    },
+    {
+        "model": "prithivida/Splade_PP_en_v1",
         "vocab_size": 30522,
         "description": "Independent Implementation of SPLADE++ Model for English",
         "size_in_GB": 0.532,
@@ -21,7 +30,9 @@ supported_splade_models = [
 
 class SpladePP(SparseTextEmbeddingBase, OnnxModel[SparseEmbedding]):
     @classmethod
-    def _post_process_onnx_output(cls, output: Tuple[np.ndarray, np.ndarray]) -> Iterable[SparseEmbedding]:
+    def _post_process_onnx_output(
+        cls, output: Tuple[np.ndarray, np.ndarray]
+    ) -> Iterable[SparseEmbedding]:
         logits, attention_mask = output
         relu_log = np.log(1 + np.maximum(logits, 0))
 
@@ -105,6 +116,10 @@ class SpladePP(SparseTextEmbeddingBase, OnnxModel[SparseEmbedding]):
             parallel=parallel,
         )
 
+    @classmethod
+    def _get_worker_class(cls) -> Type[EmbeddingWorker]:
+        return SpladePPEmbeddingWorker
+
 
 class SpladePPEmbeddingWorker(EmbeddingWorker):
     def init_embedding(

{fastembed-0.2.4 → fastembed-0.2.6}/fastembed/text/e5_onnx_embedding.py

@@ -20,7 +20,7 @@ supported_multilingual_e5_models = [
         "model": "sentence-transformers/paraphrase-multilingual-mpnet-base-v2",
         "dim": 768,
         "description": "Sentence-transformers model for tasks like clustering or semantic search",
-        "size_in_GB": 1.11,
+        "size_in_GB": 1.00,
         "sources": {
             "hf": "xenova/paraphrase-multilingual-mpnet-base-v2",
         },

{fastembed-0.2.4 → fastembed-0.2.6}/fastembed/text/jina_onnx_embedding.py

@@ -11,14 +11,14 @@ supported_jina_models = [
         "model": "jinaai/jina-embeddings-v2-base-en",
         "dim": 768,
         "description": "English embedding model supporting 8192 sequence length",
-        "size_in_GB": 0.55,
+        "size_in_GB": 0.52,
         "sources": {"hf": "xenova/jina-embeddings-v2-base-en"},
     },
     {
         "model": "jinaai/jina-embeddings-v2-small-en",
         "dim": 512,
         "description": "English embedding model supporting 8192 sequence length",
-        "size_in_GB": 0.13,
+        "size_in_GB": 0.12,
         "sources": {"hf": "xenova/jina-embeddings-v2-small-en"},
     },
 ]
@@ -49,7 +49,9 @@ class JinaOnnxEmbedding(OnnxTextEmbedding):
         return supported_jina_models
 
     @classmethod
-    def _post_process_onnx_output(cls, output: Tuple[np.ndarray, np.ndarray]) -> Iterable[np.ndarray]:
+    def _post_process_onnx_output(
+        cls, output: Tuple[np.ndarray, np.ndarray]
+    ) -> Iterable[np.ndarray]:
         embeddings, attn_mask = output
         return normalize(cls.mean_pooling(embeddings, attn_mask)).astype(np.float32)
 

{fastembed-0.2.4 → fastembed-0.2.6}/fastembed/text/onnx_embedding.py

@@ -12,7 +12,7 @@ supported_onnx_models = [
         "model": "BAAI/bge-base-en",
         "dim": 768,
         "description": "Base English model",
-        "size_in_GB": 0.5,
+        "size_in_GB": 0.42,
         "sources": {
             "url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en.tar.gz",
         },
@@ -21,26 +21,17 @@ supported_onnx_models = [
         "model": "BAAI/bge-base-en-v1.5",
         "dim": 768,
         "description": "Base English model, v1.5",
-        "size_in_GB": 0.44,
+        "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": "BAAI/bge-large-en-v1.5-quantized",
-        "dim": 1024,
-        "description": "Large English model, v1.5",
-        "size_in_GB": 1.34,
-        "sources": {
-            "hf": "qdrant/bge-large-en-v1.5-onnx-q",
-        },
-    },
     {
         "model": "BAAI/bge-large-en-v1.5",
         "dim": 1024,
         "description": "Large English model, v1.5",
-        "size_in_GB": 1.34,
+        "size_in_GB": 1.20,
         "sources": {
             "hf": "qdrant/bge-large-en-v1.5-onnx",
         },
@@ -49,7 +40,7 @@ supported_onnx_models = [
         "model": "BAAI/bge-small-en",
         "dim": 384,
         "description": "Fast English model",
-        "size_in_GB": 0.2,
+        "size_in_GB": 0.13,
         "sources": {
             "url": "https://storage.googleapis.com/qdrant-fastembed/BAAI-bge-small-en.tar.gz",
         },
@@ -69,9 +60,8 @@ supported_onnx_models = [
         "model": "BAAI/bge-small-en-v1.5",
         "dim": 384,
         "description": "Fast and Default English model",
-        "size_in_GB": 0.13,
+        "size_in_GB": 0.067,
         "sources": {
-            "url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-en-v1.5.tar.gz",
             "hf": "qdrant/bge-small-en-v1.5-onnx-q",
         },
     },
@@ -79,12 +69,12 @@ supported_onnx_models = [
         "model": "BAAI/bge-small-zh-v1.5",
         "dim": 512,
         "description": "Fast and recommended Chinese model",
-        "size_in_GB": 0.1,
+        "size_in_GB": 0.09,
         "sources": {
             "url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-zh-v1.5.tar.gz",
         },
     },
-    {  # todo: it is not a flag embedding
+    {
         "model": "sentence-transformers/all-MiniLM-L6-v2",
         "dim": 384,
         "description": "Sentence Transformer model, MiniLM-L6-v2",
@@ -98,7 +88,7 @@ supported_onnx_models = [
         "model": "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
         "dim": 384,
         "description": "Sentence Transformer model, paraphrase-multilingual-MiniLM-L12-v2",
-        "size_in_GB": 0.46,
+        "size_in_GB": 0.22,
         "sources": {
             "hf": "qdrant/paraphrase-multilingual-MiniLM-L12-v2-onnx-Q",
         },
@@ -107,7 +97,7 @@ supported_onnx_models = [
         "model": "nomic-ai/nomic-embed-text-v1",
         "dim": 768,
         "description": "8192 context length english model",
-        "size_in_GB": 0.54,
+        "size_in_GB": 0.52,
         "sources": {
             "hf": "nomic-ai/nomic-embed-text-v1",
         },
@@ -116,7 +106,7 @@ supported_onnx_models = [
         "model": "nomic-ai/nomic-embed-text-v1.5",
         "dim": 768,
         "description": "8192 context length english model",
-        "size_in_GB": 0.54,
+        "size_in_GB": 0.52,
         "sources": {
             "hf": "nomic-ai/nomic-embed-text-v1.5",
         },
@@ -125,7 +115,7 @@ supported_onnx_models = [
         "model": "thenlper/gte-large",
         "dim": 1024,
         "description": "Large general text embeddings model",
-        "size_in_GB": 1.34,
+        "size_in_GB": 1.20,
         "sources": {
             "hf": "qdrant/gte-large-onnx",
         },
@@ -143,6 +133,15 @@ supported_onnx_models = [
     #         "https://storage.googleapis.com/qdrant-fastembed/sentence-transformers-all-MiniLM-L6-v2.tar.gz"
     #     ]
     # }
+    {
+        "model": "mixedbread-ai/mxbai-embed-large-v1",
+        "dim": 1024,
+        "description": "MixedBread Base sentence embedding model, does well on MTEB",
+        "size_in_GB": 0.64,
+        "sources": {
+            "hf": "mixedbread-ai/mxbai-embed-large-v1",
+        },
+    },
 ]
 
 
@@ -230,7 +229,9 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxModel[np.ndarray]):
         return onnx_input
 
     @classmethod
-    def _post_process_onnx_output(cls, output: Tuple[np.ndarray, np.ndarray]) -> Iterable[np.ndarray]:
+    def _post_process_onnx_output(
+        cls, output: Tuple[np.ndarray, np.ndarray]
+    ) -> Iterable[np.ndarray]:
         embeddings, _ = output
         return normalize(embeddings[:, 0]).astype(np.float32)
 

{fastembed-0.2.4 → fastembed-0.2.6}/fastembed/text/text_embedding.py

@@ -55,7 +55,7 @@ class TextEmbedding(TextEmbeddingBase):
 
         for EMBEDDING_MODEL_TYPE in self.EMBEDDINGS_REGISTRY:
             supported_models = EMBEDDING_MODEL_TYPE.list_supported_models()
-            if any(model_name == model["model"] for model in supported_models):
+            if any(model_name.lower() == model["model"].lower() for model in supported_models):
                 self.model = EMBEDDING_MODEL_TYPE(model_name, cache_dir, threads, **kwargs)
                 return
 

{fastembed-0.2.4 → fastembed-0.2.6}/fastembed/text/text_embedding_base.py

@@ -6,7 +6,13 @@ from fastembed.common.model_management import ModelManagement
 
 
 class TextEmbeddingBase(ModelManagement):
-    def __init__(self, model_name: str, cache_dir: Optional[str] = None, threads: Optional[int] = None, **kwargs):
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: Optional[str] = None,
+        threads: Optional[int] = None,
+        **kwargs,
+    ):
         self.model_name = model_name
         self.cache_dir = cache_dir
         self.threads = threads

{fastembed-0.2.4 → fastembed-0.2.6}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "fastembed"
-version = "0.2.4"
+version = "0.2.6"
 description = "Fast, light, accurate library built for retrieval embedding generation"
 authors = ["NirantK <nirant.bits@gmail.com>"]
 license = "Apache License"
@@ -26,7 +26,7 @@ numpy = [
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^7.4.2"
-ruff = "^0.2.2"
+ruff = "^0.3.1"
 notebook = ">=7.0.2"
 pre-commit = {version = "^3.6.2", python = ">=3.9,<3.12" }
 
@@ -43,4 +43,4 @@ requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"
 
 [tool.ruff]
-line-length = 120
+line-length = 99