inferencer.py 10.5 KB
Newer Older
Rayyyyy's avatar
Rayyyyy committed
1
2
3
4
5
6
7
8
9
10
11
12
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


Rayyyyy's avatar
Rayyyyy committed
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
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",
}

Rayyyyy's avatar
Rayyyyy committed
33
34
35
36
37
38
39
40
41
42
43
44
45
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:

Rayyyyy's avatar
update  
Rayyyyy committed
46
    def __init__(self, model_path: str, use_vllm: bool, stream_chat: bool, tensor_parallel_size: int):
Rayyyyy's avatar
Rayyyyy committed
47
        self.use_vllm = use_vllm
Rayyyyy's avatar
Rayyyyy committed
48
49
50
        self.stream_chat = stream_chat
        # huggingface
        self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
Rayyyyy's avatar
update  
Rayyyyy committed
51
52
53
54
55
56
        # 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()

Rayyyyy's avatar
Rayyyyy committed
57
        if self.use_vllm:
Rayyyyy's avatar
update  
Rayyyyy committed
58
59
60
61
62
63
64
65
66
67
            ## 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
Rayyyyy's avatar
Rayyyyy committed
68
69
            try:
                if self.stream_chat:
Rayyyyy's avatar
Rayyyyy committed
70
                    # fastllm的流式初始化
Rayyyyy's avatar
Rayyyyy committed
71
72
                    self.model = llm.model(model_path)
                else:
Rayyyyy's avatar
update  
Rayyyyy committed
73
                    self.model = llm.from_hf(self.model, self.tokenizer, dtype="float16")
Rayyyyy's avatar
Rayyyyy committed
74
75

            except Exception as e:
Rayyyyy's avatar
Rayyyyy committed
76
                logger.error(f"fastllm initial failed, {e}")
Rayyyyy's avatar
Rayyyyy committed
77
78
79


    def chat(self, prompt: str, history=[]):
Rayyyyy's avatar
update  
Rayyyyy committed
80
        '''单轮问答'''
Rayyyyy's avatar
Rayyyyy committed
81
        import re
Rayyyyy's avatar
Rayyyyy committed
82
83
        output_text = ''
        try:
Rayyyyy's avatar
Rayyyyy committed
84
            if self.use_vllm:
Rayyyyy's avatar
update  
Rayyyyy committed
85

Rayyyyy's avatar
Rayyyyy committed
86
87
88
89
90
91
                output_text = self.model.response(prompt)
            else:
                output_text, _ = self.model.chat(self.tokenizer,
                                                    prompt,
                                                    history,
                                                    do_sample=False)
Rayyyyy's avatar
Rayyyyy committed
92
93
94
95
96
97
98
99
            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}")

Rayyyyy's avatar
Rayyyyy committed
100
        except Exception as e:
Rayyyyy's avatar
Rayyyyy committed
101
            logger.error(f"chat inference failed, {e}")
Rayyyyy's avatar
Rayyyyy committed
102
103
104
105
106
        return output_text


    def chat_stream(self, prompt: str, history=[]):
        '''流式服务'''
Rayyyyy's avatar
Rayyyyy committed
107
        import re
Rayyyyy's avatar
Rayyyyy committed
108
        if self.use_vllm:
Rayyyyy's avatar
Rayyyyy committed
109
110
111
            from fastllm_pytools import llm
            # Fastllm
            for response in self.model.stream_response(prompt, history=[]):
Rayyyyy's avatar
Rayyyyy committed
112
113
114
115
116
117
118
119
                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}")

Rayyyyy's avatar
Rayyyyy committed
120
121
122
123
124
                yield response
        else:
            # HuggingFace
            current_length = 0
            for response, _, past_key_values in self.model.stream_chat(self.tokenizer, prompt, history=history,
Rayyyyy's avatar
Rayyyyy committed
125
                                                                past_key_values=None,
Rayyyyy's avatar
Rayyyyy committed
126
127
                                                                return_past_key_values=True):
                output_text = response[current_length:]
Rayyyyy's avatar
Rayyyyy committed
128
129
130
131
132
133
134
                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}")
Rayyyyy's avatar
Rayyyyy committed
135
136
137
138
139
140
141
142
143
144

                yield output_text
                current_length = len(response)


class LLMInference:
    def __init__(self,
                 model_path: str,
                 tensor_parallel_size: int,
                 device: str = 'cuda',
Rayyyyy's avatar
update  
Rayyyyy committed
145
146
                 use_vllm: bool = False,
                 stream_chat: bool = False
Rayyyyy's avatar
Rayyyyy committed
147
148
149
150
                 ) -> None:

        self.device = device

Rayyyyy's avatar
update  
Rayyyyy committed
151
        self.inference = InferenceWrapper(model_path=model_path,
Rayyyyy's avatar
Rayyyyy committed
152
                                          use_vllm=use_vllm,
Rayyyyy's avatar
update  
Rayyyyy committed
153
                                          stream_chat=stream_chat,
Rayyyyy's avatar
Rayyyyy committed
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
                                          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):
Rayyyyy's avatar
Rayyyyy committed
177
    '''启动 Web 服务器,接收 HTTP 请求,并通过调用本地的 LLM 推理服务生成响应. '''
Rayyyyy's avatar
Rayyyyy committed
178
179
180
181
182
    config = configparser.ConfigParser()
    config.read(args.config_path)

    bind_port = int(config['default']['bind_port'])
    model_path = config['llm']['local_llm_path']
Rayyyyy's avatar
Rayyyyy committed
183
    use_vllm = config.getboolean('llm', 'use_vllm')
Rayyyyy's avatar
Rayyyyy committed
184
    inference_wrapper = InferenceWrapper(model_path,
Rayyyyy's avatar
Rayyyyy committed
185
                                         use_vllm=use_vllm,
Rayyyyy's avatar
update  
Rayyyyy committed
186
                                         tensor_parallel_size=1,
Rayyyyy's avatar
Rayyyyy committed
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
                                         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)


Rayyyyy's avatar
Rayyyyy committed
207
208
209
210
211
212
213
214
215
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}")


Rayyyyy's avatar
Rayyyyy committed
216
217
218
219
220
221
def parse_args():
    '''参数'''
    parser = argparse.ArgumentParser(
        description='Feature store for processing directories.')
    parser.add_argument(
        '--config_path',
Rayyyyy's avatar
update  
Rayyyyy committed
222
        default='../config.ini',
Rayyyyy's avatar
Rayyyyy committed
223
224
225
226
227
228
229
        help='config目录')
    parser.add_argument(
        '--query',
        default=['请问下产品的服务器保修或保修政策?'],
        help='提问的问题.')
    parser.add_argument(
        '--DCU_ID',
Rayyyyy's avatar
Rayyyyy committed
230
        type=str,
Rayyyyy's avatar
update  
Rayyyyy committed
231
        default='1',
Rayyyyy's avatar
Rayyyyy committed
232
        help='设置DCU卡号,卡号之间用英文逗号隔开,输入样例:"0,1,2"')
Rayyyyy's avatar
Rayyyyy committed
233
234
235
236
237
238
239
240
241
242
    parser.add_argument(
        '--stream_chat',
        action='store_true',
        help='启用流式对话方式')
    args = parser.parse_args()
    return args


def main():
    args = parse_args()
Rayyyyy's avatar
Rayyyyy committed
243
    set_envs(args.DCU_ID)
Rayyyyy's avatar
Rayyyyy committed
244
245
246
247
248
    llm_inference(args)


if __name__ == '__main__':
    main()