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

Refactor HuggingFaceEndpointsComponent imports and handle model_kwargs parameter

Browse source
This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-02-21 22:50:11 -03:00
commit 8be42c4e69
1 changed files with 5 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

9
src/backend/langflow/components/llms/HuggingFaceEndpoints.py
View file

@ -1,7 +1,8 @@
from typing import Optional from typing import Optional
from langflow import CustomComponent
from langchain.llms.huggingface_endpoint import HuggingFaceEndpoint
from langchain.llms.base import BaseLLM from langchain.llms.base import BaseLLM
from langchain.llms.huggingface_endpoint import HuggingFaceEndpoint
from langflow import CustomComponent
class HuggingFaceEndpointsComponent(CustomComponent): class HuggingFaceEndpointsComponent(CustomComponent):
@ -31,11 +32,11 @@ class HuggingFaceEndpointsComponent(CustomComponent):
model_kwargs: Optional[dict] = None, model_kwargs: Optional[dict] = None,
) -> BaseLLM: ) -> BaseLLM:
try: try:
output = HuggingFaceEndpoint( output = HuggingFaceEndpoint( # type: ignore
endpoint_url=endpoint_url, endpoint_url=endpoint_url,
task=task, task=task,
huggingfacehub_api_token=huggingfacehub_api_token, huggingfacehub_api_token=huggingfacehub_api_token,
model_kwargs=model_kwargs, model_kwargs=model_kwargs or {},
) )
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