Source code for langchain_community.llms.huggingface_hub
import json
from typing import Any, Dict, List, Mapping, Optional
from langchain_core._api.deprecation import deprecated
from langchain_core.callbacks import CallbackManagerForLLMRun
from langchain_core.language_models.llms import LLM
from langchain_core.pydantic_v1 import Extra, root_validator
from langchain_core.utils import get_from_dict_or_env
from langchain_community.llms.utils import enforce_stop_tokens
# key: task
# value: key in the output dictionary
VALID_TASKS_DICT = {
"translation": "translation_text",
"summarization": "summary_text",
"conversational": "generated_text",
"text-generation": "generated_text",
"text2text-generation": "generated_text",
}
[docs]@deprecated("0.0.21", removal="0.3.0", alternative="HuggingFaceEndpoint")
class HuggingFaceHub(LLM):
"""HuggingFaceHub 模型。
! 此类已被弃用,您应该使用 HuggingFaceEndpoint 替代。
要使用,您应该安装 ``huggingface_hub`` python 包,并设置环境变量 ``HUGGINGFACEHUB_API_TOKEN`` 为您的 API 令牌,或将其作为构造函数的命名参数传递。
支持 `text-generation`, `text2text-generation`, `conversational`, `translation`, 和 `summarization`。
示例:
.. code-block:: python
from langchain_community.llms import HuggingFaceHub
hf = HuggingFaceHub(repo_id="gpt2", huggingfacehub_api_token="my-api-key")
"""
client: Any #: :meta private:
repo_id: Optional[str] = None
"""要使用的模型名称。
如果未提供,则将使用所选任务的默认模型。"""
task: Optional[str] = None
"""调用模型的任务。
应该是一个返回`generated_text`、`summary_text`或`translation_text`的任务。"""
model_kwargs: Optional[dict] = None
"""传递给模型的关键字参数。"""
huggingfacehub_api_token: Optional[str] = None
class Config:
"""此pydantic对象的配置。"""
extra = Extra.forbid
@root_validator()
def validate_environment(cls, values: Dict) -> Dict:
"""验证环境中是否存在API密钥和Python包。"""
huggingfacehub_api_token = get_from_dict_or_env(
values, "huggingfacehub_api_token", "HUGGINGFACEHUB_API_TOKEN"
)
try:
from huggingface_hub import HfApi, InferenceClient
repo_id = values["repo_id"]
client = InferenceClient(
model=repo_id,
token=huggingfacehub_api_token,
)
if not values["task"]:
if not repo_id:
raise ValueError(
"Must specify either `repo_id` or `task`, or both."
)
# Use the recommended task for the chosen model
model_info = HfApi(token=huggingfacehub_api_token).model_info(
repo_id=repo_id
)
values["task"] = model_info.pipeline_tag
if values["task"] not in VALID_TASKS_DICT:
raise ValueError(
f"Got invalid task {values['task']}, "
f"currently only {VALID_TASKS_DICT.keys()} are supported"
)
values["client"] = client
except ImportError:
raise ImportError(
"Could not import huggingface_hub python package. "
"Please install it with `pip install huggingface_hub`."
)
return values
@property
def _identifying_params(self) -> Mapping[str, Any]:
"""获取识别参数。"""
_model_kwargs = self.model_kwargs or {}
return {
**{"repo_id": self.repo_id, "task": self.task},
**{"model_kwargs": _model_kwargs},
}
@property
def _llm_type(self) -> str:
"""llm的返回类型。"""
return "huggingface_hub"
def _call(
self,
prompt: str,
stop: Optional[List[str]] = None,
run_manager: Optional[CallbackManagerForLLMRun] = None,
**kwargs: Any,
) -> str:
"""调用HuggingFace Hub的推理端点。
参数:
prompt: 传递给模型的提示。
stop: 生成时可选的停止词列表。
返回:
模型生成的字符串。
示例:
.. code-block:: python
response = hf("Tell me a joke.")
"""
_model_kwargs = self.model_kwargs or {}
parameters = {**_model_kwargs, **kwargs}
response = self.client.post(
json={"inputs": prompt, "parameters": parameters}, task=self.task
)
response = json.loads(response.decode())
if "error" in response:
raise ValueError(f"Error raised by inference API: {response['error']}")
response_key = VALID_TASKS_DICT[self.task] # type: ignore
if isinstance(response, list):
text = response[0][response_key]
else:
text = response[response_key]
if stop is not None:
# This is a bit hacky, but I can't figure out a better way to enforce
# stop tokens when making calls to huggingface_hub.
text = enforce_stop_tokens(text, stop)
return text