fastembed-gpu 0.2.7__py3-none-any.whl

fastembed/__init__.py

@@ -0,0 +1,12 @@
+import importlib.metadata
+
+from fastembed.text import TextEmbedding
+from fastembed.sparse import SparseTextEmbedding, SparseEmbedding
+
+try:
+    version = importlib.metadata.version("fastembed")
+except importlib.metadata.PackageNotFoundError as _:
+    version = importlib.metadata.version("fastembed-gpu")
+
+__version__ = version
+__all__ = ["TextEmbedding", "SparseTextEmbedding", "SparseEmbedding"]

fastembed/common/__init__.py

@@ -0,0 +1,3 @@
+from fastembed.common.onnx_model import OnnxProvider
+
+__all__ = ["OnnxProvider"]

fastembed/common/model_management.py

@@ -0,0 +1,244 @@
+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
+
+
+class ModelManagement:
+    @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.
+        """
+        raise NotImplementedError()
+
+    @classmethod
+    def _get_model_description(cls, model_name: str) -> Dict[str, Any]:
+        """
+        Gets the model description from the model_name.
+
+        Args:
+            model_name (str): The name of the model.
+
+        raises:
+            ValueError: If the model_name is not supported.
+
+        Returns:
+            Dict[str, Any]: The model description.
+        """
+        for model in cls.list_supported_models():
+            if model_name.lower() == model["model"].lower():
+                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,
+        **kwargs,
+    ) -> 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,
+            local_files_only=kwargs.get("local_files_only", False),
+        )
+
+    @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, **kwargs) -> 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,
+                        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))
+
+        raise ValueError(f"Could not download model {model['model']} from any source.")

fastembed/common/models.py

@@ -0,0 +1,54 @@
+import json
+from pathlib import Path
+
+import numpy as np
+from tokenizers import Tokenizer, AddedToken
+
+
+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}")
+
+    tokenizer_path = model_dir / "tokenizer.json"
+    if not tokenizer_path.exists():
+        raise ValueError(f"Could not find tokenizer.json in {model_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}")
+
+    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}")
+
+    with open(str(config_path)) as config_file:
+        config = json.load(config_file)
+
+    with open(str(tokenizer_config_path)) as tokenizer_config_file:
+        tokenizer_config = json.load(tokenizer_config_file)
+
+    with open(str(tokens_map_path)) as tokens_map_file:
+        tokens_map = json.load(tokens_map_file)
+
+    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"]
+    )
+
+    for token in tokens_map.values():
+        if isinstance(token, str):
+            tokenizer.add_special_tokens([token])
+        elif isinstance(token, dict):
+            tokenizer.add_special_tokens([AddedToken(**token)])
+
+    return tokenizer
+
+
+def normalize(input_array, p=2, dim=1, eps=1e-12) -> np.ndarray:
+    # Calculate the Lp norm along the specified dimension
+    norm = np.linalg.norm(input_array, ord=p, axis=dim, keepdims=True)
+    norm = np.maximum(norm, eps)  # Avoid division by zero
+    normalized_array = input_array / norm
+    return normalized_array

fastembed/common/onnx_model.py

@@ -0,0 +1,165 @@
+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,
+)
+
+import numpy as np
+import onnxruntime as ort
+
+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")
+
+OnnxProvider = Union[str, Tuple[str, Dict[Any, Any]]]
+
+
+class OnnxModel(Generic[T]):
+    @classmethod
+    def _get_worker_class(cls) -> Type["EmbeddingWorker"]:
+        raise NotImplementedError("Subclasses must implement this method")
+
+    @classmethod
+    def _post_process_onnx_output(cls, output: Tuple[np.ndarray, np.ndarray]) -> Iterable[T]:
+        raise NotImplementedError("Subclasses must implement this method")
+
+    def __init__(self) -> None:
+        self.model = None
+        self.tokenizer = None
+
+    def _preprocess_onnx_input(self, onnx_input: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
+        """
+        Preprocess the onnx input.
+        """
+        return onnx_input
+
+    def load_onnx_model(
+        self,
+        model_dir: Path,
+        model_file: str,
+        threads: Optional[int],
+        providers: Optional[Sequence[OnnxProvider]] = None,
+    ) -> None:
+        model_path = model_dir / model_file
+
+        # 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}"
+                )
+
+        so = ort.SessionOptions()
+        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.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)
+        input_ids = np.array([e.ids for e in encoded])
+        attention_mask = np.array([e.attention_mask for e in encoded])
+
+        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
+            ),
+        }
+
+        onnx_input = self._preprocess_onnx_input(onnx_input)
+
+        model_output = self.model.run(None, onnx_input)
+        embeddings = model_output[0]
+        return embeddings, attention_mask
+
+    def _embed_documents(
+        self,
+        model_name: str,
+        cache_dir: str,
+        documents: Union[str, Iterable[str]],
+        batch_size: int = 256,
+        parallel: Optional[int] = None,
+    ) -> Iterable[T]:
+        is_small = False
+
+        if isinstance(documents, str):
+            documents = [documents]
+            is_small = True
+
+        if isinstance(documents, list):
+            if len(documents) < batch_size:
+                is_small = True
+
+        if parallel == 0:
+            parallel = os.cpu_count()
+
+        if parallel is None or is_small:
+            for batch in iter_batch(documents, batch_size):
+                yield from self._post_process_onnx_output(self.onnx_embed(batch))
+        else:
+            start_method = "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
+            params = {
+                "model_name": model_name,
+                "cache_dir": cache_dir,
+            }
+            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)
+
+
+class EmbeddingWorker(Worker):
+    def init_embedding(
+        self,
+        model_name: str,
+        cache_dir: str,
+    ) -> OnnxModel:
+        raise NotImplementedError()
+
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str,
+    ):
+        self.model = self.init_embedding(model_name, cache_dir)
+
+    @classmethod
+    def start(cls, model_name: str, cache_dir: str, **kwargs: Any) -> "EmbeddingWorker":
+        return cls(
+            model_name=model_name,
+            cache_dir=cache_dir,
+        )
+
+    def process(self, items: Iterable[Tuple[int, Any]]) -> Iterable[Tuple[int, Any]]:
+        for idx, batch in items:
+            embeddings, attn_mask = self.model.onnx_embed(batch)
+            yield idx, (embeddings, attn_mask)

fastembed/common/utils.py

@@ -0,0 +1,33 @@
+import os
+import tempfile
+from itertools import islice
+from pathlib import Path
+from typing import Union, Iterable, Generator, Optional
+
+
+def iter_batch(iterable: Union[Iterable, Generator], size: int) -> Iterable:
+    """
+    >>> list(iter_batch([1,2,3,4,5], 3))
+    [[1, 2, 3], [4, 5]]
+    """
+    source_iter = iter(iterable)
+    while source_iter:
+        b = list(islice(source_iter, size))
+        if len(b) == 0:
+            break
+        yield b
+
+
+def define_cache_dir(cache_dir: Optional[str] = None) -> Path:
+    """
+    Define the cache directory for fastembed
+    """
+    if cache_dir is None:
+        default_cache_dir = os.path.join(tempfile.gettempdir(), "fastembed_cache")
+        cache_path = Path(os.getenv("FASTEMBED_CACHE_PATH", default_cache_dir))
+    else:
+        cache_path = Path(cache_dir)
+
+    cache_path.mkdir(parents=True, exist_ok=True)
+
+    return cache_path

fastembed/embedding.py

@@ -0,0 +1,24 @@
+from typing import Optional
+
+from loguru import logger
+
+from fastembed import TextEmbedding
+
+logger.warning(
+    "DefaultEmbedding, FlagEmbedding, JinaEmbedding are deprecated."
+    "Use from fastembed import TextEmbedding instead."
+)
+
+DefaultEmbedding = TextEmbedding
+FlagEmbedding = TextEmbedding
+
+
+class JinaEmbedding(TextEmbedding):
+    def __init__(
+        self,
+        model_name: str = "jinaai/jina-embeddings-v2-base-en",
+        cache_dir: Optional[str] = None,
+        threads: Optional[int] = None,
+        **kwargs,
+    ):
+        super().__init__(model_name, cache_dir, threads, **kwargs)