fastembed 0.3.2__tar.gz → 0.3.4__tar.gz

{fastembed-0.3.2 → fastembed-0.3.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: fastembed
-Version: 0.3.2
+Version: 0.3.4
 Summary: Fast, light, accurate library built for retrieval embedding generation
 Home-page: https://github.com/qdrant/fastembed
 License: Apache License
@@ -43,7 +43,7 @@ The default text embedding (`TextEmbedding`) model is Flag Embedding, presented
 
 2. Fast: FastEmbed is designed for speed. We use the ONNX Runtime, which is faster than PyTorch. We also use data-parallelism for encoding large datasets.
 
-3. Accurate: FastEmbed is better than OpenAI Ada-002. We also [supported](https://qdrant.github.io/fastembed/examples/Supported_Models/) an ever expanding set of models, including a few multilingual models.
+3. Accurate: FastEmbed is better than OpenAI Ada-002. We also [support](https://qdrant.github.io/fastembed/examples/Supported_Models/) an ever expanding set of models, including a few multilingual models.
 
 ## 🚀 Installation
 
@@ -243,3 +243,4 @@ search_result = client.query(
 )
 print(search_result)
 ```
+

{fastembed-0.3.2 → fastembed-0.3.4}/README.md

@@ -10,7 +10,7 @@ The default text embedding (`TextEmbedding`) model is Flag Embedding, presented
 
 2. Fast: FastEmbed is designed for speed. We use the ONNX Runtime, which is faster than PyTorch. We also use data-parallelism for encoding large datasets.
 
-3. Accurate: FastEmbed is better than OpenAI Ada-002. We also [supported](https://qdrant.github.io/fastembed/examples/Supported_Models/) an ever expanding set of models, including a few multilingual models.
+3. Accurate: FastEmbed is better than OpenAI Ada-002. We also [support](https://qdrant.github.io/fastembed/examples/Supported_Models/) an ever expanding set of models, including a few multilingual models.
 
 ## 🚀 Installation
 
@@ -209,4 +209,4 @@ search_result = client.query(
     query_text="This is a query document"
 )
 print(search_result)
-```
+```

{fastembed-0.3.2 → fastembed-0.3.4}/fastembed/common/model_management.py

@@ -1,4 +1,5 @@
 import os
+import time
 import shutil
 import tarfile
 from pathlib import Path
@@ -42,9 +43,7 @@ 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:
+    def download_file_from_gcs(cls, url: str, output_path: str, show_progress: bool = True) -> str:
         """
         Downloads a file from Google Cloud Storage.
 
@@ -73,9 +72,7 @@ class ModelManagement:
 
         # 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}."
-            )
+            print(f"Warning: Content-length header is missing or zero in the response from {url}.")
 
         show_progress = total_size_in_bytes and show_progress
 
@@ -163,9 +160,7 @@ class ModelManagement:
         return cache_dir
 
     @classmethod
-    def retrieve_model_gcs(
-        cls, model_name: str, source_url: str, cache_dir: str
-    ) -> Path:
+    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"
@@ -191,12 +186,8 @@ class ModelManagement:
             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}"
+        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
@@ -205,7 +196,7 @@ class ModelManagement:
         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, retries=3, **kwargs) -> Path:
         """
         Downloads a model from HuggingFace Hub or Google Cloud Storage.
 
@@ -225,6 +216,7 @@ class ModelManagement:
                 }
                 ```
             cache_dir (str): The path to the cache directory.
+            retries: (int): The number of times to retry (including the first attempt)
 
         Returns:
             Path: The path to the downloaded model directory.
@@ -233,26 +225,38 @@ 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),
+        sleep = 3.0
+        while retries > 0:
+            retries -= 1
+
+            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),
+                        )
                     )
-                )
-            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))
+                except (EnvironmentError, RepositoryNotFoundError, ValueError) as e:
+                    logger.error(
+                        f"Could not download model from HuggingFace: {e} "
+                        "Falling back to other sources."
+                    )
+            if url_source:
+                try:
+                    return cls.retrieve_model_gcs(model["model"], url_source, str(cache_dir))
+                except Exception:
+                    logger.error(f"Could not download model from url: {url_source}")
+
+            logger.error(
+                f"Could not download model from either source, sleeping for {sleep} seconds, {retries} retries left."
+            )
+            time.sleep(sleep)
+            sleep *= 3
 
         raise ValueError(f"Could not download model {model['model']} from any source.")

{fastembed-0.3.2 → fastembed-0.3.4}/fastembed/image/image_embedding.py

@@ -65,8 +65,8 @@ class ImageEmbedding(ImageEmbeddingBase):
                 return
 
         raise ValueError(
-            f"Model {model_name} is not supported in TextEmbedding."
-            "Please check the supported models using `TextEmbedding.list_supported_models()`"
+            f"Model {model_name} is not supported in ImageEmbedding."
+            "Please check the supported models using `ImageEmbedding.list_supported_models()`"
         )
 
     def embed(

{fastembed-0.3.2 → fastembed-0.3.4}/fastembed/image/onnx_embedding.py

@@ -29,9 +29,28 @@ supported_onnx_models = [
         },
         "model_file": "model.onnx",
     },
+    {
+        "model": "Qdrant/Unicom-ViT-B-16",
+        "dim": 768,
+        "description": "Unicom Unicom-ViT-B-16 from open-metric-learning",
+        "size_in_GB": 0.82,
+        "sources": {
+            "hf": "Qdrant/Unicom-ViT-B-16",
+        },
+        "model_file": "model.onnx",
+    },
+    {
+        "model": "Qdrant/Unicom-ViT-B-32",
+        "dim": 512,
+        "description": "Unicom Unicom-ViT-B-32 from open-metric-learning",
+        "size_in_GB": 0.48,
+        "sources": {
+            "hf": "Qdrant/Unicom-ViT-B-32",
+        },
+        "model_file": "model.onnx",
+    },
 ]
 
-
 class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
     def __init__(
         self,
@@ -122,10 +141,16 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
 
         return onnx_input
 
-    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext
+    ) -> Iterable[np.ndarray]:
         return normalize(output.model_output).astype(np.float32)
 
 
 class OnnxImageEmbeddingWorker(ImageEmbeddingWorker):
-    def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> OnnxImageEmbedding:
-        return OnnxImageEmbedding(model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs)
+    def init_embedding(
+        self, model_name: str, cache_dir: str, **kwargs
+    ) -> OnnxImageEmbedding:
+        return OnnxImageEmbedding(
+            model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs
+        )

{fastembed-0.3.2 → fastembed-0.3.4}/fastembed/sparse/bm25.py

@@ -27,6 +27,7 @@ supported_bm25_models = [
         },
         "model_file": "mock.file",  # bm25 does not require a model, so we just use a mock
         "additional_files": ["stopwords.txt"],
+        "requires_idf": True,
     },
 ]
 
@@ -250,8 +251,8 @@ class Bm25(SparseTextEmbeddingBase):
             tokens = self.tokenizer.tokenize(text)
             stemmed_tokens = self._stem(tokens)
             token_ids = np.array(
-                [self.compute_token_id(token) for token in stemmed_tokens],
-                dtype=np.float32,
+                list(set(self.compute_token_id(token) for token in stemmed_tokens)),
+                dtype=np.int32,
             )
             values = np.ones_like(token_ids)
             yield SparseEmbedding(indices=token_ids, values=values)

{fastembed-0.3.2 → fastembed-0.3.4}/fastembed/sparse/bm42.py

@@ -27,6 +27,7 @@ supported_bm42_models = [
         },
         "model_file": "model.onnx",
         "additional_files": ["stopwords.txt"],
+        "requires_idf": True,
     },
 ]
 

{fastembed-0.3.2 → fastembed-0.3.4}/fastembed/text/onnx_embedding.py

@@ -80,36 +80,6 @@ supported_onnx_models = [
         },
         "model_file": "model_optimized.onnx",
     },
-    {
-        "model": "nomic-ai/nomic-embed-text-v1",
-        "dim": 768,
-        "description": "8192 context length english model",
-        "size_in_GB": 0.52,
-        "sources": {
-            "hf": "nomic-ai/nomic-embed-text-v1",
-        },
-        "model_file": "onnx/model.onnx",
-    },
-    {
-        "model": "nomic-ai/nomic-embed-text-v1.5",
-        "dim": 768,
-        "description": "8192 context length english model",
-        "size_in_GB": 0.52,
-        "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",
         "dim": 1024,
@@ -274,7 +244,9 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
         """
         return onnx_input
 
-    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext
+    ) -> Iterable[np.ndarray]:
         embeddings = output.model_output
         return normalize(embeddings[:, 0]).astype(np.float32)
 
@@ -286,4 +258,6 @@ class OnnxTextEmbeddingWorker(TextEmbeddingWorker):
         cache_dir: str,
         **kwargs,
     ) -> OnnxTextEmbedding:
-        return OnnxTextEmbedding(model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs)
+        return OnnxTextEmbedding(
+            model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs
+        )

fastembed-0.3.4/fastembed/text/pooled_embedding.py

@@ -0,0 +1,87 @@
+from typing import Any, Dict, Iterable, List, Type
+
+import numpy as np
+
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.utils import normalize
+from fastembed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
+from fastembed.text.onnx_text_model import TextEmbeddingWorker
+
+supported_pooled_models = [
+    {
+        "model": "nomic-ai/nomic-embed-text-v1.5",
+        "dim": 768,
+        "description": "8192 context length english model",
+        "size_in_GB": 0.52,
+        "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": "nomic-ai/nomic-embed-text-v1",
+        "dim": 768,
+        "description": "8192 context length english model",
+        "size_in_GB": 0.52,
+        "sources": {
+            "hf": "nomic-ai/nomic-embed-text-v1",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+]
+
+
+class PooledEmbedding(OnnxTextEmbedding):
+    @classmethod
+    def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
+        return PooledEmbeddingWorker
+
+    @classmethod
+    def mean_pooling(
+        cls, model_output: np.ndarray, attention_mask: np.ndarray
+    ) -> np.ndarray:
+        token_embeddings = model_output
+        input_mask_expanded = np.expand_dims(attention_mask, axis=-1)
+        input_mask_expanded = np.tile(
+            input_mask_expanded, (1, 1, token_embeddings.shape[-1])
+        )
+        input_mask_expanded = input_mask_expanded.astype(float)
+        sum_embeddings = np.sum(token_embeddings * input_mask_expanded, axis=1)
+        sum_mask = np.sum(input_mask_expanded, axis=1)
+        pooled_embeddings = sum_embeddings / np.maximum(sum_mask, 1e-9)
+        return pooled_embeddings
+
+    @classmethod
+    def list_supported_models(cls) -> List[Dict[str, Any]]:
+        """Lists the supported models.
+
+        Returns:
+            List[Dict[str, Any]]: A list of dictionaries containing the model information.
+        """
+        return supported_pooled_models
+
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext
+    ) -> Iterable[np.ndarray]:
+        embeddings = output.model_output
+        attn_mask = output.attention_mask
+        return self.mean_pooling(embeddings, attn_mask).astype(np.float32)
+
+
+class PooledEmbeddingWorker(OnnxTextEmbeddingWorker):
+    def init_embedding(
+        self, model_name: str, cache_dir: str, **kwargs
+    ) -> OnnxTextEmbedding:
+        return PooledEmbedding(
+            model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs
+        )

fastembed-0.3.2/fastembed/text/jina_onnx_embedding.py → fastembed-0.3.4/fastembed/text/pooled_normalized_embedding.py

@@ -6,8 +6,20 @@ from fastembed.common.onnx_model import OnnxOutputContext
 from fastembed.common.utils import normalize
 from fastembed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
 from fastembed.text.onnx_text_model import TextEmbeddingWorker
+from fastembed.text.pooled_embedding import PooledEmbedding
 
-supported_jina_models = [
+supported_pooled_normalized_models = [
+    {
+        "model": "sentence-transformers/all-MiniLM-L6-v2",
+        "dim": 384,
+        "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": "jinaai/jina-embeddings-v2-base-en",
         "dim": 768,
@@ -35,20 +47,10 @@ supported_jina_models = [
 ]
 
 
-class JinaOnnxEmbedding(OnnxTextEmbedding):
+class PooledNormalizedEmbedding(PooledEmbedding):
     @classmethod
     def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
-        return JinaEmbeddingWorker
-
-    @classmethod
-    def mean_pooling(cls, model_output, attention_mask) -> np.ndarray:
-        token_embeddings = model_output
-        input_mask_expanded = (np.expand_dims(attention_mask, axis=-1)).astype(float)
-
-        sum_embeddings = np.sum(token_embeddings * input_mask_expanded, axis=1)
-        mask_sum = np.clip(np.sum(input_mask_expanded, axis=1), a_min=1e-9, a_max=None)
-
-        return sum_embeddings / mask_sum
+        return PooledNormalizedEmbeddingWorker
 
     @classmethod
     def list_supported_models(cls) -> List[Dict[str, Any]]:
@@ -57,7 +59,7 @@ class JinaOnnxEmbedding(OnnxTextEmbedding):
         Returns:
             List[Dict[str, Any]]: A list of dictionaries containing the model information.
         """
-        return supported_jina_models
+        return supported_pooled_normalized_models
 
     def _post_process_onnx_output(
         self, output: OnnxOutputContext
@@ -67,10 +69,10 @@ class JinaOnnxEmbedding(OnnxTextEmbedding):
         return normalize(self.mean_pooling(embeddings, attn_mask)).astype(np.float32)
 
 
-class JinaEmbeddingWorker(OnnxTextEmbeddingWorker):
+class PooledNormalizedEmbeddingWorker(OnnxTextEmbeddingWorker):
     def init_embedding(
         self, model_name: str, cache_dir: str, **kwargs
     ) -> OnnxTextEmbedding:
-        return JinaOnnxEmbedding(
+        return PooledNormalizedEmbedding(
             model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs
         )

{fastembed-0.3.2 → fastembed-0.3.4}/fastembed/text/text_embedding.py

@@ -5,8 +5,8 @@ import numpy as np
 from fastembed.common import OnnxProvider
 from fastembed.text.clip_embedding import CLIPOnnxEmbedding
 from fastembed.text.e5_onnx_embedding import E5OnnxEmbedding
-from fastembed.text.jina_onnx_embedding import JinaOnnxEmbedding
-from fastembed.text.mini_lm_embedding import MiniLMOnnxEmbedding
+from fastembed.text.pooled_normalized_embedding import PooledNormalizedEmbedding
+from fastembed.text.pooled_embedding import PooledEmbedding
 from fastembed.text.onnx_embedding import OnnxTextEmbedding
 from fastembed.text.text_embedding_base import TextEmbeddingBase
 
@@ -15,9 +15,9 @@ class TextEmbedding(TextEmbeddingBase):
     EMBEDDINGS_REGISTRY: List[Type[TextEmbeddingBase]] = [
         OnnxTextEmbedding,
         E5OnnxEmbedding,
-        JinaOnnxEmbedding,
         CLIPOnnxEmbedding,
-        MiniLMOnnxEmbedding,
+        PooledNormalizedEmbedding,
+        PooledEmbedding,
     ]
 
     @classmethod

{fastembed-0.3.2 → fastembed-0.3.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "fastembed"
-version = "0.3.2"
+version = "0.3.4"
 description = "Fast, light, accurate library built for retrieval embedding generation"
 authors = ["Qdrant Team <info@qdrant.tech>", "NirantK <nirant.bits@gmail.com>"]
 license = "Apache License"

fastembed-0.3.2/fastembed/text/mini_lm_embedding.py

@@ -1,58 +0,0 @@
-from typing import Any, Dict, Iterable, List, Type
-
-import numpy as np
-
-from fastembed.common.onnx_model import OnnxOutputContext
-from fastembed.common.utils import normalize
-from fastembed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
-from fastembed.text.onnx_text_model import TextEmbeddingWorker
-
-supported_mini_lm_models = [
-    {
-        "model": "sentence-transformers/all-MiniLM-L6-v2",
-        "dim": 384,
-        "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",
-    }
-]
-
-
-class MiniLMOnnxEmbedding(OnnxTextEmbedding):
-    @classmethod
-    def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
-        return MiniLMEmbeddingWorker
-
-    @classmethod
-    def mean_pooling(cls, model_output: np.ndarray, attention_mask: np.ndarray) -> np.ndarray:
-        token_embeddings = model_output
-        input_mask_expanded = np.expand_dims(attention_mask, axis=-1)
-        input_mask_expanded = np.tile(input_mask_expanded, (1, 1, token_embeddings.shape[-1]))
-        input_mask_expanded = input_mask_expanded.astype(float)
-        sum_embeddings = np.sum(token_embeddings * input_mask_expanded, axis=1)
-        sum_mask = np.sum(input_mask_expanded, axis=1)
-        pooled_embeddings = sum_embeddings / np.maximum(sum_mask, 1e-9)
-        return pooled_embeddings
-
-    @classmethod
-    def list_supported_models(cls) -> List[Dict[str, Any]]:
-        """Lists the supported models.
-
-        Returns:
-            List[Dict[str, Any]]: A list of dictionaries containing the model information.
-        """
-        return supported_mini_lm_models
-
-    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
-        embeddings = output.model_output
-        attn_mask = output.attention_mask
-        return normalize(self.mean_pooling(embeddings, attn_mask)).astype(np.float32)
-
-
-class MiniLMEmbeddingWorker(OnnxTextEmbeddingWorker):
-    def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> OnnxTextEmbedding:
-        return MiniLMOnnxEmbedding(model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs)