File size: 6,283 Bytes
e60c070 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 |
#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File : zhipuai_llm.py
@Time : 2023/10/16 22:06:26
@Author : 0-yy-0
@Version : 1.0
@Contact : 310484121@qq.com
@License : (C)Copyright 2017-2018, Liugroup-NLPR-CASIA
@Desc : 基于智谱 AI 大模型自定义 LLM 类
'''
from __future__ import annotations
import logging
from typing import (
Any,
AsyncIterator,
Dict,
Iterator,
List,
Optional,
)
from langchain.callbacks.manager import (
AsyncCallbackManagerForLLMRun,
CallbackManagerForLLMRun,
)
from langchain.llms.base import LLM
from langchain.pydantic_v1 import Field, root_validator
from langchain.schema.output import GenerationChunk
from langchain.utils import get_from_dict_or_env
from llm.self_llm import Self_LLM
logger = logging.getLogger(__name__)
class ZhipuAILLM(Self_LLM):
"""Zhipuai hosted open source or customized models.
To use, you should have the ``zhipuai`` python package installed, and
the environment variable ``zhipuai_api_key`` set with
your API key and Secret Key.
zhipuai_api_key are required parameters which you could get from
https://open.bigmodel.cn/usercenter/apikeys
Example:
.. code-block:: python
from langchain.llms import ZhipuAILLM
zhipuai_model = ZhipuAILLM(model="chatglm_std", temperature=temperature)
"""
model_kwargs: Dict[str, Any] = Field(default_factory=dict)
client: Any
model: str = "chatglm_std"
"""Model name in chatglm_pro, chatglm_std, chatglm_lite. """
zhipuai_api_key: Optional[str] = None
incremental: Optional[bool] = True
"""Whether to incremental the results or not."""
streaming: Optional[bool] = False
"""Whether to streaming the results or not."""
# streaming = -incremental
request_timeout: Optional[int] = 60
"""request timeout for chat http requests"""
top_p: Optional[float] = 0.8
temperature: Optional[float] = 0.95
request_id: Optional[float] = None
@root_validator()
def validate_enviroment(cls, values: Dict) -> Dict:
values["zhipuai_api_key"] = get_from_dict_or_env(
values,
"zhipuai_api_key",
"ZHIPUAI_API_KEY",
)
params = {
"zhipuai_api_key": values["zhipuai_api_key"],
"model": values["model"],
}
try:
import zhipuai
zhipuai.api_key = values["zhipuai_api_key"]
values["client"] = zhipuai.model_api
except ImportError:
raise ValueError(
"zhipuai package not found, please install it with "
"`pip install zhipuai`"
)
return values
@property
def _identifying_params(self) -> Dict[str, Any]:
return {
**{"model": self.model},
**super()._identifying_params,
}
@property
def _llm_type(self) -> str:
"""Return type of llm."""
return "zhipuai"
@property
def _default_params(self) -> Dict[str, Any]:
"""Get the default parameters for calling OpenAI API."""
normal_params = {
"streaming": self.streaming,
"top_p": self.top_p,
"temperature": self.temperature,
"request_id": self.request_id,
}
return {**normal_params, **self.model_kwargs}
def _convert_prompt_msg_params(
self,
prompt: str,
**kwargs: Any,
) -> dict:
return {
**{"prompt": prompt, "model": self.model},
**self._default_params,
**kwargs,
}
def _call(
self,
prompt: str,
stop: Optional[List[str]] = None,
run_manager: Optional[CallbackManagerForLLMRun] = None,
**kwargs: Any,
) -> str:
"""Call out to an zhipuai models endpoint for each generation with a prompt.
Args:
prompt: The prompt to pass into the model.
Returns:
The string generated by the model.
Example:
.. code-block:: python
response = zhipuai_model("Tell me a joke.")
"""
if self.streaming:
completion = ""
for chunk in self._stream(prompt, stop, run_manager, **kwargs):
completion += chunk.text
return completion
params = self._convert_prompt_msg_params(prompt, **kwargs)
response_payload = self.client.invoke(**params)
return response_payload["data"]["choices"][-1]["content"].strip('"').strip(" ")
async def _acall(
self,
prompt: str,
stop: Optional[List[str]] = None,
run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
**kwargs: Any,
) -> str:
if self.streaming:
completion = ""
async for chunk in self._astream(prompt, stop, run_manager, **kwargs):
completion += chunk.text
return completion
params = self._convert_prompt_msg_params(prompt, **kwargs)
response = await self.client.async_invoke(**params)
return response_payload
def _stream(
self,
prompt: str,
stop: Optional[List[str]] = None,
run_manager: Optional[CallbackManagerForLLMRun] = None,
**kwargs: Any,
) -> Iterator[GenerationChunk]:
params = self._convert_prompt_msg_params(prompt, **kwargs)
for res in self.client.invoke(**params):
if res:
chunk = GenerationChunk(text=res)
yield chunk
if run_manager:
run_manager.on_llm_new_token(chunk.text)
async def _astream(
self,
prompt: str,
stop: Optional[List[str]] = None,
run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
**kwargs: Any,
) -> AsyncIterator[GenerationChunk]:
params = self._convert_prompt_msg_params(prompt, **kwargs)
async for res in await self.client.ado(**params):
if res:
chunk = GenerationChunk(text=res["data"]["choices"]["content"])
yield chunk
if run_manager:
await run_manager.on_llm_new_token(chunk.text)
|