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

fix: hf hosted inference check (#1128)

Browse source
This commit is contained in:
takatost 2023-09-09 00:29:48 +08:00 • committed by GitHub
commit c4d8bdc3db
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 69 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

8
api/core/model_providers/models/llm/huggingface_hub_model.py
View file

@ -1,6 +1,5 @@
from typing import List, Optional, Any
from langchain import HuggingFaceHub
from langchain.callbacks.manager import Callbacks
from langchain.schema import LLMResult
@ -9,6 +8,7 @@ from core.model_providers.models.llm.base import BaseLLM
from core.model_providers.models.entity.message import PromptMessage
from core.model_providers.models.entity.model_params import ModelMode, ModelKwargs
from core.third_party.langchain.llms.huggingface_endpoint_llm import HuggingFaceEndpointLLM
from core.third_party.langchain.llms.huggingface_hub_llm import HuggingFaceHubLLM
class HuggingfaceHubModel(BaseLLM):
@ -31,7 +31,7 @@ class HuggingfaceHubModel(BaseLLM):
streaming=streaming
)
else:
client = HuggingFaceHub(
client = HuggingFaceHubLLM(
repo_id=self.name,
task=self.credentials['task_type'],
model_kwargs=provider_model_kwargs,
@ -88,4 +88,6 @@ class HuggingfaceHubModel(BaseLLM):
if 'baichuan' in self.name.lower():
return False
return True
return True
else:
return False