|
from abc import ABC |
|
import requests |
|
from typing import Optional, List |
|
from langchain.llms.base import LLM |
|
|
|
from models.loader import LoaderCheckPoint |
|
from models.base import (RemoteRpcModel, |
|
AnswerResult) |
|
from typing import ( |
|
Collection, |
|
Dict |
|
) |
|
|
|
|
|
def _build_message_template() -> Dict[str, str]: |
|
""" |
|
:return: 结构 |
|
""" |
|
return { |
|
"role": "", |
|
"content": "", |
|
} |
|
|
|
|
|
class FastChatOpenAILLM(RemoteRpcModel, LLM, ABC): |
|
api_base_url: str = "http://localhost:8000/v1" |
|
model_name: str = "chatglm-6b" |
|
max_token: int = 10000 |
|
temperature: float = 0.01 |
|
top_p = 0.9 |
|
checkPoint: LoaderCheckPoint = None |
|
history = [] |
|
history_len: int = 10 |
|
|
|
def __init__(self, checkPoint: LoaderCheckPoint = None): |
|
super().__init__() |
|
self.checkPoint = checkPoint |
|
|
|
@property |
|
def _llm_type(self) -> str: |
|
return "FastChat" |
|
|
|
@property |
|
def _check_point(self) -> LoaderCheckPoint: |
|
return self.checkPoint |
|
|
|
@property |
|
def _history_len(self) -> int: |
|
return self.history_len |
|
|
|
def set_history_len(self, history_len: int = 10) -> None: |
|
self.history_len = history_len |
|
|
|
@property |
|
def _api_key(self) -> str: |
|
pass |
|
|
|
@property |
|
def _api_base_url(self) -> str: |
|
return self.api_base_url |
|
|
|
def set_api_key(self, api_key: str): |
|
pass |
|
|
|
def set_api_base_url(self, api_base_url: str): |
|
self.api_base_url = api_base_url |
|
|
|
def call_model_name(self, model_name): |
|
self.model_name = model_name |
|
|
|
def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str: |
|
print(f"__call:{prompt}") |
|
try: |
|
import openai |
|
|
|
openai.api_key = "EMPTY" |
|
openai.api_base = self.api_base_url |
|
except ImportError: |
|
raise ValueError( |
|
"Could not import openai python package. " |
|
"Please install it with `pip install openai`." |
|
) |
|
|
|
completion = openai.ChatCompletion.create( |
|
model=self.model_name, |
|
messages=self.build_message_list(prompt) |
|
) |
|
print(f"response:{completion.choices[0].message.content}") |
|
print(f"+++++++++++++++++++++++++++++++++++") |
|
return completion.choices[0].message.content |
|
|
|
|
|
def build_message_list(self, query) -> Collection[Dict[str, str]]: |
|
build_message_list: Collection[Dict[str, str]] = [] |
|
history = self.history[-self.history_len:] if self.history_len > 0 else [] |
|
for i, (old_query, response) in enumerate(history): |
|
user_build_message = _build_message_template() |
|
user_build_message['role'] = 'user' |
|
user_build_message['content'] = old_query |
|
system_build_message = _build_message_template() |
|
system_build_message['role'] = 'system' |
|
system_build_message['content'] = response |
|
build_message_list.append(user_build_message) |
|
build_message_list.append(system_build_message) |
|
|
|
user_build_message = _build_message_template() |
|
user_build_message['role'] = 'user' |
|
user_build_message['content'] = query |
|
build_message_list.append(user_build_message) |
|
return build_message_list |
|
|
|
def generatorAnswer(self, prompt: str, |
|
history: List[List[str]] = [], |
|
streaming: bool = False): |
|
|
|
try: |
|
import openai |
|
|
|
openai.api_key = "EMPTY" |
|
openai.api_base = self.api_base_url |
|
except ImportError: |
|
raise ValueError( |
|
"Could not import openai python package. " |
|
"Please install it with `pip install openai`." |
|
) |
|
|
|
completion = openai.ChatCompletion.create( |
|
model=self.model_name, |
|
messages=self.build_message_list(prompt) |
|
) |
|
|
|
history += [[prompt, completion.choices[0].message.content]] |
|
answer_result = AnswerResult() |
|
answer_result.history = history |
|
answer_result.llm_output = {"answer": completion.choices[0].message.content} |
|
|
|
yield answer_result |
|
|
