Langchain-Chatchat
111 строк · 3.9 Кб
1from fastchat.conversation import Conversation2from server.model_workers.base import *3from fastchat import conversation as conv4import sys5from typing import List, Literal, Dict6from configs import logger, log_verbose7
8
9class FangZhouWorker(ApiModelWorker):10"""11火山方舟
12"""
13
14def __init__(15self,16*,17model_names: List[str] = ["fangzhou-api"],18controller_addr: str = None,19worker_addr: str = None,20version: Literal["chatglm-6b-model"] = "chatglm-6b-model",21**kwargs,22):23kwargs.update(model_names=model_names, controller_addr=controller_addr, worker_addr=worker_addr)24kwargs.setdefault("context_len", 16384)25super().__init__(**kwargs)26self.version = version27def do_chat(self, params: ApiChatParams) -> Dict:28from volcengine.maas.v2 import MaasService29
30params.load_config(self.model_names[0])31maas = MaasService('maas-api.ml-platform-cn-beijing.volces.com', 'cn-beijing')32maas.set_ak(params.api_key)33maas.set_sk(params.secret_key)34
35# document: "https://www.volcengine.com/docs/82379/1099475"36req = {37"parameters": {38# 这里的参数仅为示例,具体可用的参数请参考具体模型的 API 说明39"max_new_tokens": params.max_tokens,40"temperature": params.temperature,41},42"messages": params.messages,43}44
45text = ""46if log_verbose:47self.logger.info(f'{self.__class__.__name__}:maas: {maas}')48for resp in maas.stream_chat(params.version, unicode_escape_data(req)):49error = resp.error50if error and error.code_n > 0:51data = {52"error_code": error.code_n,53"text": error.message,54"error": {55"message": error.message,56"type": "invalid_request_error",57"param": None,58"code": None,59}60}61self.logger.error(f"请求方舟 API 时发生错误:{data}")62yield data63elif chunk := resp.choices and resp.choices[0].message.content:64text += chunk65yield {"error_code": 0, "text": text}66else:67data = {68"error_code": 500,69"text": f"请求方舟 API 时发生未知的错误: {resp}"70}71self.logger.error(data)72yield data73break74
75def get_embeddings(self, params):76print("embedding")77print(params)78
79def make_conv_template(self, conv_template: str = None, model_path: str = None) -> Conversation:80return conv.Conversation(81name=self.model_names[0],82system_message="你是一个聪明、对人类有帮助的人工智能,你可以对人类提出的问题给出有用、详细、礼貌的回答。",83messages=[],84roles=["user", "assistant", "system"],85sep="\n### ",86stop_str="###",87)88
89
90def unicode_escape_data(data):91if isinstance(data, str):92return data.encode('unicode_escape').decode('ascii')93elif isinstance(data, dict):94return {key: unicode_escape_data(value) for key, value in data.items()}95elif isinstance(data, list):96return [unicode_escape_data(item) for item in data]97else:98return data99
100if __name__ == "__main__":101import uvicorn102from server.utils import MakeFastAPIOffline103from fastchat.serve.model_worker import app104
105worker = FangZhouWorker(106controller_addr="http://127.0.0.1:20001",107worker_addr="http://127.0.0.1:21005",108)109sys.modules["fastchat.serve.model_worker"].worker = worker110MakeFastAPIOffline(app)111uvicorn.run(app, port=21005)