Allowing HuggingFaceEmbeddings from the cached weight (#3084)

### https://github.com/hwchase17/langchain/issues/3079
Allow initializing HuggingFaceEmbeddings from the cached weight
This commit is contained in:
Azam Iftikhar 2023-04-18 20:00:35 +05:30 committed by GitHub
parent 55f6f80a59
commit 188e9b9beb
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -1,5 +1,5 @@
"""Wrapper around HuggingFace embedding models.""" """Wrapper around HuggingFace embedding models."""
from typing import Any, List from typing import Any, List, Optional
from pydantic import BaseModel, Extra from pydantic import BaseModel, Extra
@ -29,6 +29,9 @@ class HuggingFaceEmbeddings(BaseModel, Embeddings):
client: Any #: :meta private: client: Any #: :meta private:
model_name: str = DEFAULT_MODEL_NAME model_name: str = DEFAULT_MODEL_NAME
"""Model name to use.""" """Model name to use."""
cache_folder: Optional[str] = None
"""Path to store models.
Can be also set by SENTENCE_TRANSFORMERS_HOME enviroment variable."""
def __init__(self, **kwargs: Any): def __init__(self, **kwargs: Any):
"""Initialize the sentence_transformer.""" """Initialize the sentence_transformer."""
@ -36,7 +39,9 @@ class HuggingFaceEmbeddings(BaseModel, Embeddings):
try: try:
import sentence_transformers import sentence_transformers
self.client = sentence_transformers.SentenceTransformer(self.model_name) self.client = sentence_transformers.SentenceTransformer(
self.model_name, self.cache_folder
)
except ImportError: except ImportError:
raise ValueError( raise ValueError(
"Could not import sentence_transformers python package. " "Could not import sentence_transformers python package. "