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


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=[]):
        '''单轮问答'''
        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)
        except Exception as e:
            logger.error(f"chat inference failed, {e}")
        return output_text


    def chat_stream(self, prompt: str, history=[]):
        '''流式服务'''
        if self.use_vllm:
            from fastllm_pytools import llm
            # Fastllm
            for response in self.model.stream_response(prompt, history=[]):
                yield response
        else:
            # HuggingFace
            current_length = 0
            past_key_values = None
            for response, _, past_key_values in self.model.stream_chat(self.tokenizer, prompt, history=history,
                                                                past_key_values=past_key_values,
                                                                return_past_key_values=True):
                output_text = response[current_length:]

                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()