\n",
" \n",
" | 0 | \n",
- " BAAI/bge-small-en | \n",
- " 384 | \n",
- " Fast English model | \n",
- " 0.20 | \n",
- "
\n",
- " \n",
- " | 1 | \n",
- " BAAI/bge-small-en-v1.5 | \n",
- " 384 | \n",
- " Fast and Default English model | \n",
- " 0.13 | \n",
- "
\n",
- " \n",
- " | 2 | \n",
- " BAAI/bge-small-zh-v1.5 | \n",
- " 512 | \n",
- " Fast and recommended Chinese model | \n",
- " 0.10 | \n",
- "
\n",
- " \n",
- " | 3 | \n",
" BAAI/bge-base-en | \n",
" 768 | \n",
" Base English model | \n",
" 0.50 | \n",
"
\n",
" \n",
- " | 4 | \n",
+ " 1 | \n",
" BAAI/bge-base-en-v1.5 | \n",
" 768 | \n",
" Base English model, v1.5 | \n",
" 0.44 | \n",
"
\n",
" \n",
- " | 5 | \n",
- " sentence-transformers/all-MiniLM-L6-v2 | \n",
+ " 2 | \n",
+ " BAAI/bge-large-en-v1.5 | \n",
+ " 1024 | \n",
+ " Large English model, v1.5 | \n",
+ " 1.34 | \n",
+ "
\n",
+ " \n",
+ " | 3 | \n",
+ " BAAI/bge-small-en | \n",
" 384 | \n",
- " Sentence Transformer model, MiniLM-L6-v2 | \n",
- " 0.09 | \n",
+ " Fast English model | \n",
+ " 0.20 | \n",
+ "
\n",
+ " \n",
+ " | 4 | \n",
+ " BAAI/bge-small-en-v1.5 | \n",
+ " 384 | \n",
+ " Fast and Default English model | \n",
+ " 0.13 | \n",
+ "
\n",
+ " \n",
+ " | 5 | \n",
+ " BAAI/bge-small-zh-v1.5 | \n",
+ " 512 | \n",
+ " Fast and recommended Chinese model | \n",
+ " 0.10 | \n",
"
\n",
" \n",
" | 6 | \n",
@@ -106,46 +106,76 @@
" English embedding model supporting 8192 sequence length | \n",
" 0.13 | \n",
"
\n",
+ " \n",
+ " | 9 | \n",
+ " sentence-transformers/all-MiniLM-L6-v2 | \n",
+ " 384 | \n",
+ " Sentence Transformer model, MiniLM-L6-v2 | \n",
+ " 0.09 | \n",
+ "
\n",
+ " \n",
+ " | 10 | \n",
+ " xenova/multilingual-e5-large | \n",
+ " 1024 | \n",
+ " Multilingual model. Recommended for non-English languages | \n",
+ " 2.24 | \n",
+ "
\n",
+ " \n",
+ " | 11 | \n",
+ " xenova/paraphrase-multilingual-mpnet-base-v2 | \n",
+ " 768 | \n",
+ " Sentence-transformers model for tasks like clustering or semantic search | \n",
+ " 1.11 | \n",
+ "
\n",
" \n",
"\n",
""
],
"text/plain": [
- " model dim \\\n",
- "0 BAAI/bge-small-en 384 \n",
- "1 BAAI/bge-small-en-v1.5 384 \n",
- "2 BAAI/bge-small-zh-v1.5 512 \n",
- "3 BAAI/bge-base-en 768 \n",
- "4 BAAI/bge-base-en-v1.5 768 \n",
- "5 sentence-transformers/all-MiniLM-L6-v2 384 \n",
- "6 intfloat/multilingual-e5-large 1024 \n",
- "7 jinaai/jina-embeddings-v2-base-en 768 \n",
- "8 jinaai/jina-embeddings-v2-small-en 512 \n",
+ " model dim \\\n",
+ "0 BAAI/bge-base-en 768 \n",
+ "1 BAAI/bge-base-en-v1.5 768 \n",
+ "2 BAAI/bge-large-en-v1.5 1024 \n",
+ "3 BAAI/bge-small-en 384 \n",
+ "4 BAAI/bge-small-en-v1.5 384 \n",
+ "5 BAAI/bge-small-zh-v1.5 512 \n",
+ "6 intfloat/multilingual-e5-large 1024 \n",
+ "7 jinaai/jina-embeddings-v2-base-en 768 \n",
+ "8 jinaai/jina-embeddings-v2-small-en 512 \n",
+ "9 sentence-transformers/all-MiniLM-L6-v2 384 \n",
+ "10 xenova/multilingual-e5-large 1024 \n",
+ "11 xenova/paraphrase-multilingual-mpnet-base-v2 768 \n",
"\n",
- " description \\\n",
- "0 Fast English model \n",
- "1 Fast and Default English model \n",
- "2 Fast and recommended Chinese model \n",
- "3 Base English model \n",
- "4 Base English model, v1.5 \n",
- "5 Sentence Transformer model, MiniLM-L6-v2 \n",
- "6 Multilingual model, e5-large. Recommend using this model for non-English languages \n",
- "7 English embedding model supporting 8192 sequence length \n",
- "8 English embedding model supporting 8192 sequence length \n",
+ " description \\\n",
+ "0 Base English model \n",
+ "1 Base English model, v1.5 \n",
+ "2 Large English model, v1.5 \n",
+ "3 Fast English model \n",
+ "4 Fast and Default English model \n",
+ "5 Fast and recommended Chinese model \n",
+ "6 Multilingual model, e5-large. Recommend using this model for non-English languages \n",
+ "7 English embedding model supporting 8192 sequence length \n",
+ "8 English embedding model supporting 8192 sequence length \n",
+ "9 Sentence Transformer model, MiniLM-L6-v2 \n",
+ "10 Multilingual model. Recommended for non-English languages \n",
+ "11 Sentence-transformers model for tasks like clustering or semantic search \n",
"\n",
- " size_in_GB \n",
- "0 0.20 \n",
- "1 0.13 \n",
- "2 0.10 \n",
- "3 0.50 \n",
- "4 0.44 \n",
- "5 0.09 \n",
- "6 2.24 \n",
- "7 0.55 \n",
- "8 0.13 "
+ " size_in_GB \n",
+ "0 0.50 \n",
+ "1 0.44 \n",
+ "2 1.34 \n",
+ "3 0.20 \n",
+ "4 0.13 \n",
+ "5 0.10 \n",
+ "6 2.24 \n",
+ "7 0.55 \n",
+ "8 0.13 \n",
+ "9 0.09 \n",
+ "10 2.24 \n",
+ "11 1.11 "
]
},
- "execution_count": 1,
+ "execution_count": 6,
"metadata": {},
"output_type": "execute_result"
}
@@ -175,7 +205,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
- "version": "3.9.17"
+ "version": "3.11.7"
},
"orig_nbformat": 4
},
diff --git a/fooling_around/Throughput_Across_Models.ipynb b/experiments/Throughput_Across_Models.ipynb
similarity index 100%
rename from fooling_around/Throughput_Across_Models.ipynb
rename to experiments/Throughput_Across_Models.ipynb
diff --git a/fastembed/__init__.py b/fastembed/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/fastembed/common/__init__.py b/fastembed/common/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/fastembed/common/model_management.py b/fastembed/common/model_management.py
new file mode 100644
index 0000000..732d5e7
--- /dev/null
+++ b/fastembed/common/model_management.py
@@ -0,0 +1,196 @@
+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
+
+
+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 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) -> 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.
+ Returns:
+ Path: The path to the model directory.
+ """
+
+ return snapshot_download(
+ repo_id=hf_source_repo,
+ ignore_patterns=["model.safetensors", "pytorch_model.bin"],
+ cache_dir=cache_dir,
+ )
+
+ @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
+
+ if model_dir.exists():
+ 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"
+
+ 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) -> 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:
+ try:
+ 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.")
+
+ 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.")
diff --git a/fastembed/common/models.py b/fastembed/common/models.py
new file mode 100644
index 0000000..74cbdab
--- /dev/null
+++ b/fastembed/common/models.py
@@ -0,0 +1,52 @@
+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["pad_token_id"], 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
diff --git a/fastembed/common/utils.py b/fastembed/common/utils.py
new file mode 100644
index 0000000..89f881b
--- /dev/null
+++ b/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_dir = Path(os.getenv("FASTEMBED_CACHE_PATH", default_cache_dir))
+ else:
+ cache_dir = Path(cache_dir)
+
+ cache_dir.mkdir(parents=True, exist_ok=True)
+
+ return cache_dir
diff --git a/fastembed/embedding.py b/fastembed/embedding.py
index 620b5d5..e7c207e 100644
--- a/fastembed/embedding.py
+++ b/fastembed/embedding.py
@@ -1,694 +1,21 @@
-import functools
-import json
-import os
-import shutil
-import tarfile
-import tempfile
-from abc import ABC, abstractmethod
-from itertools import islice
-from multiprocessing import get_all_start_methods
-from pathlib import Path
-from typing import Any, Callable, Dict, Generator, Iterable, List, Optional, Tuple, Union
+from typing import Optional
-import numpy as np
-import onnxruntime as ort
-import requests
-from tokenizers import AddedToken, Tokenizer
-from tqdm import tqdm
-from huggingface_hub import snapshot_download
-from huggingface_hub.utils import RepositoryNotFoundError
from loguru import logger
-from fastembed.parallel_processor import ParallelWorkerPool, Worker
+from fastembed.text.text_embedding import TextEmbedding
+logger.warning("DefaultEmbedding, FlagEmbedding, JinaEmbedding are deprecated." " Use TextEmbedding instead.")
-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
+DefaultEmbedding = TextEmbedding
+FlagEmbedding = TextEmbedding
-def locate_model_file(model_dir: Path, file_names: List[str]):
- """
- 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 path in model_dir.rglob("*.onnx"):
- for file_name in file_names:
- if path.is_file() and path.name == file_name:
- return path
-
- raise ValueError(f"Could not find either of {', '.join(file_names)} in {model_dir}")
-
-
-def normalize(input_array, p=2, dim=1, eps=1e-12):
- # 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
-
-
-class EmbeddingModel:
- @classmethod
- def load_tokenizer(cls, 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["pad_token_id"], 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 __init__(
- self,
- path: Path,
- model_name: str,
- max_length: int = 512,
- max_threads: int = None,
- ):
- self.path = path
- self.model_name = model_name
- model_path = locate_model_file(self.path, ["model.onnx", "model_optimized.onnx"])
-
- # List of Execution Providers: https://onnxruntime.ai/docs/execution-providers
- onnx_providers = ["CPUExecutionProvider"]
-
- # Hacky support for multilingual model
- self.exclude_token_type_ids = False
- if model_name == "intfloat/multilingual-e5-large":
- self.exclude_token_type_ids = True
-
- so = ort.SessionOptions()
- so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
-
- if max_threads is not None:
- so.intra_op_num_threads = max_threads
- so.inter_op_num_threads = max_threads
-
- self.tokenizer = self.load_tokenizer(self.path, max_length=max_length)
- 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),
- }
-
- if not self.exclude_token_type_ids:
- onnx_input["token_type_ids"] = np.array(
- [np.zeros(len(e), dtype=np.int64) for e in input_ids], dtype=np.int64
- )
-
- model_output = self.model.run(None, onnx_input)
- embeddings = model_output[0]
- return embeddings, attention_mask
-
-
-class EmbeddingWorker(Worker):
- def __init__(
- self,
- path: Path,
- model_name: str,
- max_length: int = 512,
- ):
- self.model = EmbeddingModel(path=path, model_name=model_name, max_length=max_length, max_threads=1)
-
- @classmethod
- def start(cls, path: Path, model_name: str, max_length: int = 512, **kwargs: Any) -> "EmbeddingWorker":
- return cls(
- path=path,
- model_name=model_name,
- max_length=max_length,
- )
-
- 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)
-
-
-class Embedding(ABC):
- """
- Abstract class for embeddings.
-
- Inherits:
- ABC: Abstract base class
-
- Raises:
- NotImplementedError: Raised when you call an abstract method that has not been implemented.
- PermissionError: _description_
- ValueError: Several possible reasons: 1) targz_path does not exist or is not a file, 2) targz_path is not a .tar.gz file, 3) An error occurred while decompressing targz_path, 4) Could not find model_dir in cache_dir, 5) Could not find tokenizer.json in model_dir, 6) Could not find model.onnx in model_dir.
- NotImplementedError: _description_
-
- Returns:
- _type_: _description_
-
- Yields:
- _type_: _description_
- """
-
- # Internal helper decorator to maintain backward compatibility
- # by supporting a fallback to download from Google Cloud Storage (GCS)
- # if the model couldn't be downloaded from HuggingFace.
- def gcs_fallback(hf_download_method: Callable) -> Callable:
- @functools.wraps(hf_download_method)
- def wrapper(self, *args, **kwargs):
- try:
- return hf_download_method(self, *args, **kwargs)
- except (EnvironmentError, RepositoryNotFoundError, ValueError) as e:
- logger.exception(
- f"Could not download model from HuggingFace: {e}"
- "Falling back to download from Google Cloud Storage"
- )
- return self.retrieve_model_gcs(*args, **kwargs)
-
- return wrapper
-
- @abstractmethod
- def embed(self, texts: Iterable[str], batch_size: int = 256, parallel: int = None) -> List[np.ndarray]:
- raise NotImplementedError
-
- @classmethod
- def list_supported_models(cls, exclude: List[str] = []) -> List[Dict[str, Any]]:
- """Lists the supported models.
-
- Args:
- exclude (List[str], optional): Keys to exclude from the result. Defaults to [].
-
- Returns:
- List[Dict[str, Any]]: A list of dictionaries containing the model information.
- """
- models_file_path = Path(__file__).with_name("models.json")
- with open(models_file_path, "r") as file:
- models = json.load(file)
-
- models = [{k: v for k, v in model.items() if k not in exclude} for model in models]
-
- return models
-
- @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, model_name: str, cache_dir: Optional[str] = None) -> str:
- """
- Downloads a model from HuggingFace Hub.
- Args:
- model_name (str): Name of the model to download.
- cache_dir (Optional[str]): The path to the cache directory.
- Raises:
- ValueError: If the model_name is not in the format