Class BaiChuanModel

src/models/baichuan_model.py:16–91

class BaiChuanModel(ToolModel):
    def __init__(self, model_path: str, peft_path: str = None, template: str = "default", trust_remote_code=True, tensor_parallel_size=1, gpu_memory_utilization=0.25):
        self.model_path = model_path
        self.peft_path = peft_path
        self.template = template
        self.trust_remote_code = trust_remote_code
        self.tensor_parallel_size = tensor_parallel_size
        self.gpu_memory_utilization = gpu_memory_utilization
        self.generation_config = GenerationConfig.from_pretrained(model_path)
        self.load_model(self.model_path, self.peft_path, self.trust_remote_code, self.tensor_parallel_size, self.gpu_memory_utilization)

    def generate(
        self, prompts: str, 
        template: str = None,
        generate_configs: GenerateConfigs =None,
        history: list = None,
    ) -> list:
        '''产出对应结果'''
        template = self.template if template is None else template

        params = self.generate_params(generate_configs)

        if template == "default":
            inputs = self.tokenizer(prompts, return_tensors="pt")
            inputs["input_ids"] = inputs["input_ids"].cuda()

            inputs.update(params)
            output = self.model.generate(**inputs)
            predict = self.tokenizer.decode(output[0].tolist())[len(prompts):]
            predict = predict.replace("<|endoftext|>", "").replace("</s>",  "")
            return predict
        elif template != "default":
            messages = [{"role": "user" if idx==0 else "assistant", "content": ii} for i in history for idx, ii in enumerate(i)]
            messages.append({"role": "user", "content": prompts})
            output = self.model.chat(self.tokenizer, messages=messages, generation_config=self.generation_config)
            return output

    def generate_params(
        self, generate_configs: GenerateConfigs,
    ):
        '''generate param'''
        kargs = generate_configs.dict()
        params = {
                    "max_new_tokens": kargs.get("max_new_tokens", 128),
                    "top_k": kargs.get("top_k", 50),
                    "top_p": kargs.get("top_p", 0.95),
                    "temperature": kargs.get("temperature", 1.0),
        }
        self.generation_config.max_new_tokens = kargs.get("max_new_tokens", 128)
        self.generation_config.top_k = kargs.get("top_k", 50)
        self.generation_config.top_p = kargs.get("top_p", 0.95)
        self.generation_config.temperature = kargs.get("temperature", 1.0)

        # params = {
        #             "n": 1,
        #             "max_tokens": kargs.get("max_new_tokens", 128),
        #             "best_of": kargs.get("beam_bums", 1), 
        #             "top_k": kargs.get("top_k", 50),