Refactor HuggingFaceEmbeddingsComponent to use langchain_community embeddings

This commit is contained in:
anovazzi1 2024-01-10 17:22:41 -03:00
commit 8fc03d06ed

View file

@ -1,7 +1,6 @@
from langflow import CustomComponent
from typing import Optional, Any, Dict
from langflow.field_typing import Embeddings
from typing import Optional, Dict
from langchain_community.embeddings.huggingface import HuggingFaceEmbeddings
class HuggingFaceEmbeddingsComponent(CustomComponent):
display_name = "HuggingFaceEmbeddings"
@ -13,9 +12,8 @@ class HuggingFaceEmbeddingsComponent(CustomComponent):
def build_config(self):
return {
"cache_folder": {"display_name": "Cache Folder", "advanced": True},
"client": {"display_name": "Client", "advanced": True},
"encode_kwargs": {"display_name": "Encode Kwargs", "advanced": True},
"model_kwargs": {"display_name": "Model Kwargs", "advanced": True},
"encode_kwargs": {"display_name": "Encode Kwargs", "advanced": True,"field_type":"dict"},
"model_kwargs": {"display_name": "Model Kwargs","field_type":"dict", "advanced": True},
"model_name": {"display_name": "Model Name"},
"multi_process": {"display_name": "Multi Process", "advanced": True},
}
@ -23,15 +21,13 @@ class HuggingFaceEmbeddingsComponent(CustomComponent):
def build(
self,
cache_folder: Optional[str] = None,
client: Optional[Any] = None,
encode_kwargs: Optional[Dict] = None,
model_kwargs: Optional[Dict] = None,
model_name: str = "sentence-transformers/all-mpnet-base-v2",
multi_process: bool = False,
) -> Embeddings:
return Embeddings(
) -> HuggingFaceEmbeddings:
return HuggingFaceEmbeddings(
cache_folder=cache_folder,
client=client,
encode_kwargs=encode_kwargs,
model_kwargs=model_kwargs,
model_name=model_name,