自定义 LLM：LangChain与文心一言擦出火花

import jsonimport timefrom typing import Any, List, Mapping, Optional, Dict, Union, Tupleimport loggingimport requestsfrom langchain.callbacks.manager import CallbackManagerForLLMRunfrom langchain.llms.base import LLMfrom langchain.utils import get_from_dict_or_envfrom pydantic import Field, root_validatorlogger = logging.getLogger(__name__)def get_access_token(api_key: str, secret_key: str):    """    使用 API Key，Secret Key 获取access_token    """    url = f"https://aip.baidubce.com/oauth/2.0/token?grant_type=client_credentials&client_id={api_key}&client_secret={secret_key}"    payload = json.dumps("")    headers = {        'Content-Type': 'application/json',        'Accept': 'application/json'    }    resp = requests.request("POST", url, headers=headers, data=payload)    return resp.json().get("access_token")class ErnieLLm(LLM):    url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/eb-instant"    model_name: str = Field(default="ERNIE-Bot-turbo", alias="model")    request_timeout: Optional[Union[float, Tuple[float, float]]] = None    temperature: float = 0.95    """temperature 说明：    （1）较高的数值会使输出更加随机，而较低的数值会使其更加集中和确定    （2）默认0.95，范围 (0, 1.0]，不能为0    （3）建议该参数和top_p只设置1个    （4）建议top_p和temperature不要同时更改    """    top_p: float = 0.8    """top_p 说明：    （1）影响输出文本的多样性，取值越大，生成文本的多样性越强    （2）默认0.8，取值范围 [0, 1.0]    （3）建议该参数和temperature只设置1个    （4）建议top_p和temperature不要同时更改    """    penalty_score: float = 1.0    """通过对已生成的token增加惩罚，减少重复生成的现象。说明：    （1）值越大表示惩罚越大    （2）默认1.0，取值范围：[1.0, 2.0]    """    ernie_api_key: Optional[str] = None    """文心一言大模型 apiKey"""    ernie_secret_key: Optional[str] = None    """文心一言大模型 secretKey"""    user_id: Optional[str] = None    """表示最终用户的唯一标识符，可以监视和检测滥用行为，防止接口恶意调用"""    streaming: bool = False    """是否以流式接口的形式返回数据，默认false"""    cache: bool = False    """是否开启缓存，默认为false"""    model_kwargs: Dict[str, Any] = Field(default_factory=dict)    """Holds any model parameters valid for `create` call not explicitly specified."""    @root_validator()    def validate_environment(cls, values: Dict) -> Dict:        """Validate that api key and python package exists in environment."""        values["ernie_api_key"] = get_from_dict_or_env(            values, "ernie_api_key", "ERNIE_API_KEY"        )        values["ernie_secret_key"] = get_from_dict_or_env(            values,            "ernie_secret_key",            "ERNIE_SECRET_KEY"        )        return values    @property    def _default_params(self) -> Dict[str, Any]:        """获取调用Ennie API的默认参数。"""        normal_params = {            "temperature": self.temperature,            "top_p": self.top_p,            "penalty_score": self.penalty_score,            "request_timeout": self.request_timeout,        }        return {**normal_params, **self.model_kwargs}    def _construct_query(self, prompt: str) -> Dict:        """构造请求体"""        query = {            "messages": [                {                    "role": "user",                    "content": prompt                }            ],            "stream": self.streaming,            "temperature": self.temperature,            "top_p": self.top_p,            "penalty_score": self.penalty_score,            "user_id": self.user_id,        }        return query    @property    def _identifying_params(self) -> Mapping[str, Any]:        """Get the identifying parameters."""        return {**{"model_name": self.model_name}, **self._default_params}    @property    def _llm_type(self) -> str:        return "ernie"    def _call(            self,            prompt: str,            stop: Optional[List[str]] = None,            run_manager: Optional[CallbackManagerForLLMRun] = None,            **kwargs: Any,    ) -> str:        """_call 实现对模型的调用"""        # construct query        query = self._construct_query(prompt=prompt)        # post        _headers = {"Content-Type": "application/json"}        with requests.session() as session:            resp = session.post(                self.url + "?access_token=" + get_access_token(self.ernie_api_key, self.ernie_secret_key),                json=query,                headers=_headers,                timeout=60)            if resp.status_code == 200:                resp_json = resp.json()                predictions = resp_json["result"]                return predictions        return "请求失败"

from langchain.cache import InMemoryCache# 启动llm的缓存langchain.llm_cache = InMemoryCache()

llm = ErnieLLm(temperature=0.95, top_p=1.0, cache=True)s = time.perf_counter()# 第一次调用print(llm("给我讲一个笑话"))elapsed = time.perf_counter() - sprint("\033[1m" + f"第一次调用耗时 {elapsed:0.2f} 秒." + "\033[0m")s = time.perf_counter()# 第一次调用print(llm("给我讲一个笑话"))elapsed = time.perf_counter() - sprint("\033[1m" + f"第二次调用耗时 {elapsed:0.2f} 秒." + "\033[0m")

当然可以，这是一个关于两只鸟的笑话：有两只小鸟，一只小鸟问：“哥哥，人们都说你长得好帅，你觉得自己帅吗？”哥哥小鸟羞涩地回答：“不，我不觉得自己帅，我只是很可爱。”而第二只小鸟打断了他：“哥呀，人家说的是你旁边的蝴蝶卷毛哈。”第一次调用耗时 2.37 秒.当然可以，这是一个关于两只鸟的笑话：有两只小鸟，一只小鸟问：“哥哥，人们都说你长得好帅，你觉得自己帅吗？”哥哥小鸟羞涩地回答：“不，我不觉得自己帅，我只是很可爱。”而第二只小鸟打断了他：“哥呀，人家说的是你旁边的蝴蝶卷毛哈。”第二次调用耗时 0.00 秒.