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

refactor: Update HuggingFaceEmbeddingsComponent and HuggingFaceInferenceAPIEmbeddingsComponent to use new Inputs/Outputs format

Browse source
This commit is contained in:
Cezar Vasconcelos 2024-06-19 21:12:07 +00:00
commit f93bdf8cd8
2 changed files with 50 additions and 51 deletions
Download patch file Download diff file Expand all files Collapse all files

44
src/backend/base/langflow/components/embeddings/HuggingFaceEmbeddings.py
View file

@ -2,10 +2,12 @@ from typing import Dict, Optional
from langchain_community.embeddings.huggingface import HuggingFaceEmbeddings from langchain_community.embeddings.huggingface import HuggingFaceEmbeddings
from langflow.custom import CustomComponent from langflow.base.models.model import LCModelComponent
from langflow.field_typing import Embeddings
from langflow.io import BoolInput, DictInput, TextInput, Output
class HuggingFaceEmbeddingsComponent(CustomComponent): class HuggingFaceEmbeddingsComponent(LCModelComponent):
display_name = "Hugging Face Embeddings" display_name = "Hugging Face Embeddings"
description = "Generate embeddings using HuggingFace models." description = "Generate embeddings using HuggingFace models."
documentation = ( documentation = (
@ -13,27 +15,23 @@ class HuggingFaceEmbeddingsComponent(CustomComponent):
) )
icon = "HuggingFace" icon = "HuggingFace"
def build_config(self): inputs = [
return { TextInput(name="cache_folder", display_name="Cache Folder", advanced=True),
"cache_folder": {"display_name": "Cache Folder", "advanced": True}, DictInput(name="encode_kwargs", display_name="Encode Kwargs", advanced=True),
"encode_kwargs": {"display_name": "Encode Kwargs", "advanced": True, "field_type": "dict"}, DictInput(name="model_kwargs", display_name="Model Kwargs", advanced=True),
"model_kwargs": {"display_name": "Model Kwargs", "field_type": "dict", "advanced": True}, TextInput(name="model_name", display_name="Model Name", value="sentence-transformers/all-mpnet-base-v2"),
"model_name": {"display_name": "Model Name"}, BoolInput(name="multi_process", display_name="Multi Process", advanced=True),
"multi_process": {"display_name": "Multi Process", "advanced": True}, ]
}
def build( outputs = [
self, Output(display_name="Embeddings", name="embeddings", method="build_embeddings"),
cache_folder: Optional[str] = None, ]
encode_kwargs: Optional[Dict] = {},
model_kwargs: Optional[Dict] = {}, def build_embeddings(self) -> Embeddings:
model_name: str = "sentence-transformers/all-mpnet-base-v2",
multi_process: bool = False,
) -> HuggingFaceEmbeddings:
return HuggingFaceEmbeddings( return HuggingFaceEmbeddings(
cache_folder=cache_folder, cache_folder=self.cache_folder,
encode_kwargs=encode_kwargs, encode_kwargs=self.encode_kwargs,
model_kwargs=model_kwargs, model_kwargs=self.model_kwargs,
model_name=model_name, model_name=self.model_name,
multi_process=multi_process, multi_process=self.multi_process,
) )

57
src/backend/base/langflow/components/embeddings/HuggingFaceInferenceAPIEmbeddings.py
View file

@ -3,42 +3,43 @@ from typing import Dict, Optional
from langchain_community.embeddings.huggingface import HuggingFaceInferenceAPIEmbeddings from langchain_community.embeddings.huggingface import HuggingFaceInferenceAPIEmbeddings
from pydantic.v1.types import SecretStr from pydantic.v1.types import SecretStr
from langflow.custom import CustomComponent from langflow.base.models.model import LCModelComponent
from langflow.field_typing import Embeddings
from langflow.io import BoolInput, DictInput, FloatInput, Output, SecretStrInput, TextInput
class HuggingFaceInferenceAPIEmbeddingsComponent(CustomComponent): class HuggingFaceInferenceAPIEmbeddingsComponent(LCModelComponent):
display_name = "Hugging Face API Embeddings" display_name = "Hugging Face API Embeddings"
description = "Generate embeddings using Hugging Face Inference API models." description = "Generate embeddings using Hugging Face Inference API models."
documentation = "https://github.com/huggingface/text-embeddings-inference" documentation = "https://github.com/huggingface/text-embeddings-inference"
icon = "HuggingFace" icon = "HuggingFace"
def build_config(self): inputs = [
return { SecretStrInput(name="api_key", display_name="API Key", advanced=True),
"api_key": {"display_name": "API Key", "password": True, "advanced": True}, TextInput(name="api_url", display_name="API URL", advanced=True, value="http://localhost:8080"),
"api_url": {"display_name": "API URL", "advanced": True}, TextInput(name="model_name", display_name="Model Name", value="BAAI/bge-large-en-v1.5"),
"model_name": {"display_name": "Model Name"}, TextInput(name="cache_folder", display_name="Cache Folder", advanced=True),
"cache_folder": {"display_name": "Cache Folder", "advanced": True}, DictInput(name="encode_kwargs", display_name="Encode Kwargs", advanced=True),
"encode_kwargs": {"display_name": "Encode Kwargs", "advanced": True, "field_type": "dict"}, DictInput(name="model_kwargs", display_name="Model Kwargs", advanced=True),
"model_kwargs": {"display_name": "Model Kwargs", "field_type": "dict", "advanced": True}, BoolInput(name="multi_process", display_name="Multi Process", advanced=True),
"multi_process": {"display_name": "Multi Process", "advanced": True}, ]
}
def build( outputs = [
self, Output(display_name="Embeddings", name="embeddings", method="build_embeddings"),
api_key: Optional[str] = "", ]
api_url: str = "http://localhost:8080",
model_name: str = "BAAI/bge-large-en-v1.5", def build_embeddings(self) -> Embeddings:
cache_folder: Optional[str] = None, if not self.api_key:
encode_kwargs: Optional[Dict] = {},
model_kwargs: Optional[Dict] = {},
multi_process: bool = False,
) -> HuggingFaceInferenceAPIEmbeddings:
if api_key:
secret_api_key = SecretStr(api_key)
else:
raise ValueError("API Key is required") raise ValueError("API Key is required")
api_key = SecretStr(self.api_key)
return HuggingFaceInferenceAPIEmbeddings( return HuggingFaceInferenceAPIEmbeddings(
api_key=secret_api_key, api_key=api_key,
api_url=api_url, api_url=self.api_url,
model_name=model_name, model_name=self.model_name,
cache_folder=self.cache_folder,
encode_kwargs=self.encode_kwargs,
model_kwargs=self.model_kwargs,
multi_process=self.multi_process,
) )