Source code for langchain_community.llms.nlpcloud

from typing import Any, Dict, List, Mapping, Optional

from langchain_core.callbacks import CallbackManagerForLLMRun
from langchain_core.language_models.llms import LLM
from langchain_core.pydantic_v1 import Extra, SecretStr, root_validator
from langchain_core.utils import convert_to_secret_str, get_from_dict_or_env


class NLPCloud(LLM):
    """NLPCloud大型语言模型。

    要使用，您应该已安装``nlpcloud`` python包，并且
    将环境变量``NLPCLOUD_API_KEY``设置为您的API密钥。

    示例:
        .. code-block:: python

            from langchain_community.llms import NLPCloud
            nlpcloud = NLPCloud(model="finetuned-gpt-neox-20b")
    """

    client: Any  #: :meta private:
    model_name: str = "finetuned-gpt-neox-20b"
    """要使用的模型名称。"""
    gpu: bool = True
    """是否使用GPU或不使用"""
    lang: str = "en"
    """要使用的语言（多语言插件）"""
    temperature: float = 0.7
    """使用哪种采样温度。"""
    max_length: int = 256
    """生成完成的最大令牌数。"""
    length_no_input: bool = True
    """min_length和max_length是否应该包括输入的长度。"""
    remove_input: bool = True
    """从API响应中删除输入文本"""
    remove_end_sequence: bool = True
    """是否删除结束序列标记。"""
    bad_words: List[str] = []
    """不允许生成的令牌列表。"""
    top_p: float = 1.0
    """每一步需要考虑的标记的总概率质量。"""
    top_k: int = 50
    """保留用于top-k过滤的最高概率标记数。"""
    repetition_penalty: float = 1.0
    """对重复的标记进行惩罚。1.0表示没有惩罚。"""
    num_beams: int = 1
    """Beam search的beam数量。"""
    num_return_sequences: int = 1
    """每个提示生成多少个完成。"""

    nlpcloud_api_key: Optional[SecretStr] = None

    class Config:
        """此pydantic对象的配置。"""

        extra = Extra.forbid

    @root_validator()
    def validate_environment(cls, values: Dict) -> Dict:
        """验证环境中是否存在API密钥和Python包。"""
        values["nlpcloud_api_key"] = convert_to_secret_str(
            get_from_dict_or_env(values, "nlpcloud_api_key", "NLPCLOUD_API_KEY")
        )
        try:
            import nlpcloud

            values["client"] = nlpcloud.Client(
                values["model_name"],
                values["nlpcloud_api_key"].get_secret_value(),
                gpu=values["gpu"],
                lang=values["lang"],
            )
        except ImportError:
            raise ImportError(
                "Could not import nlpcloud python package. "
                "Please install it with `pip install nlpcloud`."
            )
        return values

    @property
    def _default_params(self) -> Mapping[str, Any]:
        """获取调用NLPCloud API 的默认参数。"""
        return {
            "temperature": self.temperature,
            "max_length": self.max_length,
            "length_no_input": self.length_no_input,
            "remove_input": self.remove_input,
            "remove_end_sequence": self.remove_end_sequence,
            "bad_words": self.bad_words,
            "top_p": self.top_p,
            "top_k": self.top_k,
            "repetition_penalty": self.repetition_penalty,
            "num_beams": self.num_beams,
            "num_return_sequences": self.num_return_sequences,
        }

    @property
    def _identifying_params(self) -> Mapping[str, Any]:
        """获取识别参数。"""
        return {
            **{"model_name": self.model_name},
            **{"gpu": self.gpu},
            **{"lang": self.lang},
            **self._default_params,
        }

    @property
    def _llm_type(self) -> str:
        """llm的返回类型。"""
        return "nlpcloud"

    def _call(
        self,
        prompt: str,
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> str:
        """调用NLPCloud的create端点。

参数：
    prompt: 传递给模型的提示。
    stop: 此接口不支持（在init方法中传入）

返回：
    模型生成的字符串。

示例：
    .. code-block:: python

        response = nlpcloud("Tell me a joke.")
"""
        if stop and len(stop) > 1:
            raise ValueError(
                "NLPCloud only supports a single stop sequence per generation."
                "Pass in a list of length 1."
            )
        elif stop and len(stop) == 1:
            end_sequence = stop[0]
        else:
            end_sequence = None
        params = {**self._default_params, **kwargs}
        response = self.client.generation(prompt, end_sequence=end_sequence, **params)
        return response["generated_text"]