Fixed the connection to the forward component

This commit is contained in:
Alexandre Henrique 2023-08-29 16:34:40 -03:00
commit 9530230c2c

View file

@ -1,23 +1,42 @@
from typing import Optional from typing import Optional
from langflow import CustomComponent from langflow import CustomComponent
from langchain.llms import HuggingFaceEndpoint from langchain.llms import HuggingFaceEndpoint
from langchain.llms.base import BaseLLM
class HuggingFaceEndpointsComponent(CustomComponent): class HuggingFaceEndpointsComponent(CustomComponent):
display_name: str = "Hugging Face Endpoints" display_name: str = "Hugging Face Inference API"
description: str = "LLM model from Inference Hugging Face Endpoints." description: str = "LLM model from Hugging Face Inference API."
def build_config(self): def build_config(self):
return { return {
"endpoint_url": {"display_name": "Endpoint URL", "password": True}, "endpoint_url": {"display_name": "Endpoint URL", "password": True},
"task": {"display_name": "Task", "type": "select", "options": ["text2text-generation", "text-generation", "summarization"]}, "task": {
"display_name": "Task",
"type": "select",
"options": ["text2text-generation", "text-generation", "summarization"],
},
"huggingfacehub_api_token": {"display_name": "API token", "password": True},
"model_kwargs": {
"display_name": "Model Keyword Arguments",
"field_type": "code",
},
"code": {"show": False}, "code": {"show": False},
} }
def build( def build(
self, endpoint_url: str, task="text2text-generation", self,
) -> HuggingFaceEndpoint: endpoint_url: str,
task="text2text-generation",
huggingfacehub_api_token: Optional[str] = None,
model_kwargs: Optional[dict] = None,
) -> BaseLLM:
try: try:
output = HuggingFaceEndpoint(endpoint_url=endpoint_url, task=task) output = HuggingFaceEndpoint(
endpoint_url=endpoint_url,
task=task,
huggingfacehub_api_token=huggingfacehub_api_token,
)
except Exception as e: except Exception as e:
raise ValueError("Could not connect to HuggingFace Endpoints API.") from e raise ValueError("Could not connect to HuggingFace Endpoints API.") from e
return output return output