jyapayne/langflow
1
0
Fork 0
Code Issues Pull requests Projects Releases 1 Packages Wiki Activity Actions

Refactor HuggingFaceEmbeddingsComponent to use langchain_community embeddings

Browse source
This commit is contained in:
anovazzi1 2024-01-10 17:22:41 -03:00
commit 8fc03d06ed
1 changed files with 6 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

16
src/backend/langflow/components/embeddings/HuggingFaceEmbeddings.py
View file

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