From 8be42c4e69160b020b92ec3174f363107ea26fde Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Wed, 21 Feb 2024 22:50:11 -0300 Subject: [PATCH] Refactor HuggingFaceEndpointsComponent imports and handle model_kwargs parameter --- .../langflow/components/llms/HuggingFaceEndpoints.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/backend/langflow/components/llms/HuggingFaceEndpoints.py b/src/backend/langflow/components/llms/HuggingFaceEndpoints.py index 0d28d5b9b..579e3c0b3 100644 --- a/src/backend/langflow/components/llms/HuggingFaceEndpoints.py +++ b/src/backend/langflow/components/llms/HuggingFaceEndpoints.py @@ -1,7 +1,8 @@ from typing import Optional -from langflow import CustomComponent -from langchain.llms.huggingface_endpoint import HuggingFaceEndpoint + from langchain.llms.base import BaseLLM +from langchain.llms.huggingface_endpoint import HuggingFaceEndpoint +from langflow import CustomComponent class HuggingFaceEndpointsComponent(CustomComponent): @@ -31,11 +32,11 @@ class HuggingFaceEndpointsComponent(CustomComponent): model_kwargs: Optional[dict] = None, ) -> BaseLLM: try: - output = HuggingFaceEndpoint( + output = HuggingFaceEndpoint( # type: ignore endpoint_url=endpoint_url, task=task, huggingfacehub_api_token=huggingfacehub_api_token, - model_kwargs=model_kwargs, + model_kwargs=model_kwargs or {}, ) except Exception as e: raise ValueError("Could not connect to HuggingFace Endpoints API.") from e