warleagle/Medical_Chat_Autocompletion

1class Conversation:2    def __init__(3        self,4    ):5        6        self.message_template = "<s>{role}\n{content}</s>"7        self.response_template = "<s>bot\n"8        self.messages = [{9            "role": "system",10            "content": "Предложи ответ оператора технической поддержки на вопрос пользователя из чата."11        }]12 13    def add_user_message(self, message):14        self.messages.append({15            "role": "user",16            "content": message17        })18 19    def add_bot_message(self, message):20        self.messages.append({21            "role": "bot",22            "content": message23        })24 25    def get_prompt(self):26        final_text = ""27        for message in self.messages:28            message_text = self.message_template.format(**message)29            final_text += message_text30        final_text += "<s>bot\n"31        return final_text.strip()32 33 34def generate(model, tokenizer, prompt, generation_config):35    data = tokenizer(prompt, return_tensors="pt", add_special_tokens=False)36    data = {k: v.to(model.device) for k, v in data.items()}37    output_ids = model.generate(38        **data,39        generation_config=generation_config40    )[0]41    output_ids = output_ids[len(data["input_ids"][0]):]42    output = tokenizer.decode(output_ids, skip_special_tokens=True)43    return output.strip()