inferencer.py

import time
import os
import configparser
import argparse
from multiprocessing import Value
from aiohttp import web
import torch
from loguru import logger

from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModel


COMMON = {
    "<光合组织登记网址>": "https://www.hieco.com.cn/partner?from=timeline",
    "<官网>": "https://www.sugon.com/after_sale/policy?sh=1",
    "<平台联系方式>": "1、访问官网,根据您所在地地址联系平台人员,网址地址：https://www.sugon.com/about/contact;\n2、点击人工客服进行咨询;\n3、请您拨打中科曙光服务热线400-810-0466联系人工进行咨询。",
    "<购买与维修的咨询方法>": "1、确定付费处理,可以微信搜索'sugon中科曙光服务'小程序,选择'在线报修'业务\n2、先了解价格,可以微信搜索'sugon中科曙光服务'小程序,选择'其他咨询'业务\n3、请您拨打中科曙光服务热线400-810-0466",
    "<服务器续保流程>": "1、微信搜索'sugon中科曙光服务'小程序,选择'延保与登记'业务\n2、点击人工客服进行登记\n3、请您拨打中科曙光服务热线400-810-0466根据语音提示选择维保与购买",
    "<XC内外网OS网盘链接>": "【腾讯文档】XC内外网OS网盘链接：https://docs.qq.com/sheet/DTWtXbU1BZHJvWkJm",
    "<W360-G30机器,安装Win7使用的镜像链接>": "W360-G30机器,安装Win7使用的镜像链接:https://pan.baidu.com/s/1SjHqCP6kJ9KzdJEBZDEynw;提取码:x6m4",
    "<麒麟系统搜狗输入法下载链接>": "软件下载链接（百度云盘）:链接:https://pan.baidu.com/s/18Iluvs4BOAfFET0yFMBeLQ,提取码:bhkf",
    "<X660 G45 GPU服务器拆解视频网盘链接>": "链接: https://pan.baidu.com/s/1RkRGh4XY1T2oYftGnjLp4w;提取码: v2qi",
    "<DS800,SANTRICITY存储IBM版本模拟器网盘链接>": "链接:https://pan.baidu.com/s/1euG9HGbPfrVbThEB8BX76g;提取码:o2ya",
    "<E80-D312（X680-G55）风冷整机组装说明下载链接>": "链接:https://pan.baidu.com/s/17KDpm-Z9lp01WGp9sQaQ4w;提取码:0802",
    "<X680 G55 风冷相关资料下载链接>": "链接:https://pan.baidu.com/s/1KQ-hxUIbTWNkc0xzrEQLjg;提取码:0802",
    "<R620 G51刷写EEPROM下载>": "下载链接如下:http://10.2.68.104/tools/bytedance/eeprom/",
    "<X7450A0服务器售后培训文件网盘链接>": "网盘下载:https://pan.baidu.com/s/1tZJIf_IeQLOWsvuOawhslQ?pwd=kgf1;提取码:kgf1",
    "<福昕阅读器补丁链接>": "补丁链接: https://pan.baidu.com/s/1QJQ1kHRplhhFly-vxJquFQ,提取码: aupx1",
    "<W330-H35A_22DB4/W3335HA安装win7网盘链接>": "硬盘链接: https://pan.baidu.com/s/1fDdGPH15mXiw0J-fMmLt6Q提取码: k97i",
    "<X680 G55服务器售后培训资料网盘链接>": "云盘连接下载:链接:https://pan.baidu.com/s/1gaok13DvNddtkmk6Q-qLYg?pwd=xyhb提取码:xyhb",
}

def build_history_messages(prompt, history, system: str = None):
    history_messages = []
    if system is not None and len(system) > 0:
        history_messages.append({'role': 'system', 'content': system})
    for item in history:
        history_messages.append({'role': 'user', 'content': item[0]})
        history_messages.append({'role': 'assistant', 'content': item[1]})
    history_messages.append({'role': 'user', 'content': prompt})
    return history_messages


class InferenceWrapper:

    def __init__(self, model_path: str, use_vllm: bool, stream_chat: bool, tensor_parallel_size: int):
        self.use_vllm = use_vllm
        self.stream_chat = stream_chat
        # huggingface
        self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
        # self.model = AutoModelForCausalLM.from_pretrained(model_path,
        #                                                   trust_remote_code=True,
        #                                                   torch_dtype=torch.float16).cuda().eval()
        model = AutoModel.from_pretrained(model_path, trust_remote_code=True).half().cuda()
        self.model = model.eval()

        if self.use_vllm:
            ## vllm
            # from vllm import LLM, SamplingParams
            #
            # self.sampling_params = SamplingParams(temperature=1, top_p=0.95)
            # self.llm = LLM(model=model_path,
            #                trust_remote_code=True,
            #                enforce_eager=True,
            #                tensor_parallel_size=tensor_parallel_size)
            ## fastllm
            from fastllm_pytools import llm
            try:
                if self.stream_chat:
                    # fastllm的流式初始化
                    self.model = llm.model(model_path)
                else:
                    self.model = llm.from_hf(self.model, self.tokenizer, dtype="float16")

            except Exception as e:
                logger.error(f"fastllm initial failed, {e}")


    def chat(self, prompt: str, history=[]):
        '''单轮问答'''
        import re
        output_text = ''
        try:
            if self.use_vllm:

                output_text = self.model.response(prompt)
            else:
                output_text, _ = self.model.chat(self.tokenizer,
                                                    prompt,
                                                    history,
                                                    do_sample=False)
            matchObj = re.match('.*(<.*>).*', output_text)
            if matchObj:
                obj = matchObj.group(1)
                replace_str = COMMON.get(obj)

                output_text = output_text.replace(obj, replace_str)
                logger.info(f"{obj} be replaced {replace_str}, after {output_text}")

        except Exception as e:
            logger.error(f"chat inference failed, {e}")
        return output_text


    def chat_stream(self, prompt: str, history=[]):
        '''流式服务'''
        import re
        if self.use_vllm:
            from fastllm_pytools import llm
            # Fastllm
            for response in self.model.stream_response(prompt, history=[]):
                matchObj = re.match('.*(<.*>).*', response)
                if matchObj:
                    obj = matchObj.group(1)
                    replace_str = COMMON.get(obj)

                    response = response.replace(obj, replace_str)
                    logger.info(f"{obj} be replaced {replace_str}, after {response}")

                yield response
        else:
            # HuggingFace
            current_length = 0
            for response, _, past_key_values in self.model.stream_chat(self.tokenizer, prompt, history=history,
                                                                past_key_values=None,
                                                                return_past_key_values=True):
                output_text = response[current_length:]
                matchObj = re.match('.*(<.*>).*', output_text)
                if matchObj:
                    obj = matchObj.group(1)
                    replace_str = COMMON.get(obj)

                    output_text = output_text.replace(obj, replace_str)
                    logger.info(f"{obj} be replaced {replace_str}, after {output_text}")

                yield output_text
                current_length = len(response)


class LLMInference:
    def __init__(self,
                 model_path: str,
                 tensor_parallel_size: int,
                 device: str = 'cuda',
                 use_vllm: bool = False,
                 stream_chat: bool = False
                 ) -> None:

        self.device = device

        self.inference = InferenceWrapper(model_path=model_path,
                                          use_vllm=use_vllm,
                                          stream_chat=stream_chat,
                                          tensor_parallel_size=tensor_parallel_size)

    def generate_response(self, prompt, history=[]):
        output_text = ''
        error = ''
        time_tokenizer = time.time()

        try:
            output_text = self.inference.chat(prompt, history)

        except Exception as e:
            error = str(e)
            logger.error(error)

        time_finish = time.time()

        logger.debug('output_text：{} \ntimecost {} '.format(output_text,
            time_finish - time_tokenizer))

        return output_text, error


def llm_inference(args):
    '''启动 Web 服务器，接收 HTTP 请求，并通过调用本地的 LLM 推理服务生成响应. '''
    config = configparser.ConfigParser()
    config.read(args.config_path)

    bind_port = int(config['default']['bind_port'])
    model_path = config['llm']['local_llm_path']
    use_vllm = config.getboolean('llm', 'use_vllm')
    inference_wrapper = InferenceWrapper(model_path,
                                         use_vllm=use_vllm,
                                         tensor_parallel_size=1,
                                         stream_chat=args.stream_chat)
    async def inference(request):
        start = time.time()
        input_json = await request.json()

        prompt = input_json['prompt']
        history = input_json['history']
        if args.stream_chat:
            text = inference_wrapper.stream_chat(prompt=prompt, history=history)
        else:
            text = inference_wrapper.chat(prompt=prompt, history=history)
        end = time.time()
        logger.debug('问题:{} 回答:{} \ntimecost {} '.format(prompt, text, end - start))
        return web.json_response({'text': text})

    app = web.Application()
    app.add_routes([web.post('/inference', inference)])
    web.run_app(app, host='0.0.0.0', port=bind_port)


def set_envs(dcu_ids):
    try:
        os.environ["CUDA_VISIBLE_DEVICES"] = dcu_ids
        logger.info(f"Set environment variable CUDA_VISIBLE_DEVICES to {dcu_ids}")
    except Exception as e:
        logger.error(f"{e}, but got {dcu_ids}")
        raise ValueError(f"{e}")


def parse_args():
    '''参数'''
    parser = argparse.ArgumentParser(
        description='Feature store for processing directories.')
    parser.add_argument(
        '--config_path',
        default='../config.ini',
        help='config目录')
    parser.add_argument(
        '--query',
        default=['请问下产品的服务器保修或保修政策？'],
        help='提问的问题.')
    parser.add_argument(
        '--DCU_ID',
        type=str,
        default='1',
        help='设置DCU卡号,卡号之间用英文逗号隔开,输入样例:"0,1,2"')
    parser.add_argument(
        '--stream_chat',
        action='store_true',
        help='启用流式对话方式')
    args = parser.parse_args()
    return args


def main():
    args = parse_args()
    set_envs(args.DCU_ID)
    llm_inference(args)


if __name__ == '__main__':
    main()