diff --git a/src/backend/langflow/components/embeddings/HuggingFaceEmbeddings.py b/src/backend/langflow/components/embeddings/HuggingFaceEmbeddings.py index 4d8d73c4a..c3dc04ec8 100644 --- a/src/backend/langflow/components/embeddings/HuggingFaceEmbeddings.py +++ b/src/backend/langflow/components/embeddings/HuggingFaceEmbeddings.py @@ -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,