mirror of https://github.com/hpcaitech/ColossalAI
151 lines
5.2 KiB
Python
151 lines
5.2 KiB
Python
"""
|
||
LLM wrapper for Pangu
|
||
|
||
Usage:
|
||
|
||
# URL: “盘古大模型套件管理”->点击“服务管理”->“模型列表”->点击想要使用的模型的“复制路径”
|
||
# USERNAME: 华为云控制台：“我的凭证”->“API凭证”下的“IAM用户名”，也就是你登录IAM账户的名字
|
||
# PASSWORD: IAM用户的密码
|
||
# DOMAIN_NAME: 华为云控制台：“我的凭证”->“API凭证”下的“用户名”，也就是公司管理IAM账户的总账户名
|
||
|
||
os.environ["URL"] = ""
|
||
os.environ["URLNAME"] = ""
|
||
os.environ["PASSWORD"] = ""
|
||
os.environ["DOMAIN_NAME"] = ""
|
||
|
||
pg = Pangu(id=1)
|
||
pg.set_auth_config()
|
||
|
||
res = pg('你是谁') # 您好,我是华为盘古大模型。我能够通过和您对话互动为您提供帮助。请问您有什么想问我的吗?
|
||
"""
|
||
|
||
import http.client
|
||
import json
|
||
from typing import Any, List, Mapping, Optional
|
||
|
||
import requests
|
||
from langchain.llms.base import LLM
|
||
from langchain.utils import get_from_dict_or_env
|
||
|
||
|
||
class Pangu(LLM):
|
||
"""
|
||
A custom LLM class that integrates pangu models
|
||
|
||
"""
|
||
|
||
n: int
|
||
gen_config: dict = None
|
||
auth_config: dict = None
|
||
|
||
def __init__(self, gen_config=None, **kwargs):
|
||
super(Pangu, self).__init__(**kwargs)
|
||
if gen_config is None:
|
||
self.gen_config = {"user": "User", "max_tokens": 50, "temperature": 0.95, "n": 1}
|
||
else:
|
||
self.gen_config = gen_config
|
||
|
||
@property
|
||
def _identifying_params(self) -> Mapping[str, Any]:
|
||
"""Get the identifying parameters."""
|
||
return {"n": self.n}
|
||
|
||
@property
|
||
def _llm_type(self) -> str:
|
||
return "pangu"
|
||
|
||
def _call(self, prompt: str, stop: Optional[List[str]] = None, **kwargs) -> str:
|
||
"""
|
||
Args:
|
||
prompt: The prompt to pass into the model.
|
||
stop: A list of strings to stop generation when encountered
|
||
|
||
Returns:
|
||
The string generated by the model
|
||
"""
|
||
# Update the generation arguments
|
||
for key, value in kwargs.items():
|
||
if key in self.gen_config:
|
||
self.gen_config[key] = value
|
||
|
||
response = self.text_completion(prompt, self.gen_config, self.auth_config)
|
||
text = response["choices"][0]["text"]
|
||
if stop is not None:
|
||
for stopping_words in stop:
|
||
if stopping_words in text:
|
||
text = text.split(stopping_words)[0]
|
||
return text
|
||
|
||
def set_auth_config(self, **kwargs):
|
||
url = get_from_dict_or_env(kwargs, "url", "URL")
|
||
username = get_from_dict_or_env(kwargs, "username", "USERNAME")
|
||
password = get_from_dict_or_env(kwargs, "password", "PASSWORD")
|
||
domain_name = get_from_dict_or_env(kwargs, "domain_name", "DOMAIN_NAME")
|
||
|
||
region = url.split(".")[1]
|
||
auth_config = {}
|
||
auth_config["endpoint"] = url[url.find("https://") + 8 : url.find(".com") + 4]
|
||
auth_config["resource_path"] = url[url.find(".com") + 4 :]
|
||
auth_config["auth_token"] = self.get_latest_auth_token(region, username, password, domain_name)
|
||
self.auth_config = auth_config
|
||
|
||
def get_latest_auth_token(self, region, username, password, domain_name):
|
||
url = f"https://iam.{region}.myhuaweicloud.com/v3/auth/tokens"
|
||
payload = json.dumps(
|
||
{
|
||
"auth": {
|
||
"identity": {
|
||
"methods": ["password"],
|
||
"password": {"user": {"name": username, "password": password, "domain": {"name": domain_name}}},
|
||
},
|
||
"scope": {"project": {"name": region}},
|
||
}
|
||
}
|
||
)
|
||
headers = {"Content-Type": "application/json"}
|
||
|
||
response = requests.request("POST", url, headers=headers, data=payload)
|
||
return response.headers["X-Subject-Token"]
|
||
|
||
def text_completion(self, text, gen_config, auth_config):
|
||
conn = http.client.HTTPSConnection(auth_config["endpoint"])
|
||
payload = json.dumps(
|
||
{
|
||
"prompt": text,
|
||
"user": gen_config["user"],
|
||
"max_tokens": gen_config["max_tokens"],
|
||
"temperature": gen_config["temperature"],
|
||
"n": gen_config["n"],
|
||
}
|
||
)
|
||
headers = {
|
||
"X-Auth-Token": auth_config["auth_token"],
|
||
"Content-Type": "application/json",
|
||
}
|
||
conn.request("POST", auth_config["resource_path"], payload, headers)
|
||
res = conn.getresponse()
|
||
data = res.read()
|
||
data = json.loads(data.decode("utf-8"))
|
||
return data
|
||
|
||
def chat_model(self, messages, gen_config, auth_config):
|
||
conn = http.client.HTTPSConnection(auth_config["endpoint"])
|
||
payload = json.dumps(
|
||
{
|
||
"messages": messages,
|
||
"user": gen_config["user"],
|
||
"max_tokens": gen_config["max_tokens"],
|
||
"temperature": gen_config["temperature"],
|
||
"n": gen_config["n"],
|
||
}
|
||
)
|
||
headers = {
|
||
"X-Auth-Token": auth_config["auth_token"],
|
||
"Content-Type": "application/json",
|
||
}
|
||
conn.request("POST", auth_config["resource_path"], payload, headers)
|
||
res = conn.getresponse()
|
||
data = res.read()
|
||
data = json.loads(data.decode("utf-8"))
|
||
return data
|