diff --git a/llama-index-integrations/embeddings/llama-index-embeddings-huggingface/llama_index/embeddings/huggingface/base.py b/llama-index-integrations/embeddings/llama-index-embeddings-huggingface/llama_index/embeddings/huggingface/base.py index cca7da667cb78ec1eba77c261a9e59bea4bafb81..4c1a81f2a2a6c73c7b941e4eb6017287972f5258 100644 --- a/llama-index-integrations/embeddings/llama-index-embeddings-huggingface/llama_index/embeddings/huggingface/base.py +++ b/llama-index-integrations/embeddings/llama-index-embeddings-huggingface/llama_index/embeddings/huggingface/base.py @@ -68,6 +68,7 @@ class HuggingFaceEmbedding(BaseEmbedding): trust_remote_code: bool = False, device: Optional[str] = None, callback_manager: Optional[CallbackManager] = None, + safe_serialization: Optional[bool] = None, ): self._device = device or infer_torch_device() @@ -80,7 +81,10 @@ class HuggingFaceEmbedding(BaseEmbedding): else DEFAULT_HUGGINGFACE_EMBEDDING_MODEL ) model = AutoModel.from_pretrained( - model_name, cache_dir=cache_folder, trust_remote_code=trust_remote_code + model_name, + cache_dir=cache_folder, + trust_remote_code=trust_remote_code, + safe_serialization=safe_serialization, ) elif model_name is None: # Extract model_name from model model_name = model.name_or_path