import os
import torch
import requests
import json
from transformers import AutoModelForCausalLM, AutoTokenizer
from configs import ip, api_port, model_path
os.environ['CUDA_LAUNCH_BLOCKING'] = '1'

class Linly:
    def __init__(self, mode='api', model_path="Linly-AI/Chinese-LLaMA-2-7B-hf", prefix_prompt = '''请用少于25个字回答以下问题\n\n'''):
        # mode = api need
        # 定义设置的api的服务器,首先记得运行Linly-api-fast.py 填入ip地址和端口号
        self.url = f"http://{ip}:{api_port}" # local server: http://ip:port
        self.headers = {
            "Content-Type": "application/json"
        }
        self.data = {
            "question": "北京有什么好玩的地方？"
        }
        # 全局设定的prompt
        self.prefix_prompt = prefix_prompt
        self.mode = mode
        if mode != 'api':
            self.model, self.tokenizer = self.init_model(model_path)
        self.history = []
    
    def init_model(self, path = "Linly-AI/Chinese-LLaMA-2-7B-hf"):
        model = AutoModelForCausalLM.from_pretrained(path, device_map="cuda:0",
                                                    torch_dtype=torch.bfloat16, trust_remote_code=True)
        tokenizer = AutoTokenizer.from_pretrained(path, use_fast=False, trust_remote_code=True)
        return model, tokenizer   
    
    def generate(self, question, system_prompt=""):
        if self.mode != 'api':
            self.data["question"] = self.message_to_prompt(question, system_prompt)
            inputs = self.tokenizer(self.data["question"], return_tensors="pt").to("cuda:0")
            try:
                generate_ids = self.model.generate(inputs.input_ids, 
                                                   max_new_tokens=2048, 
                                                   do_sample=True, 
                                                   top_k=20, 
                                                   top_p=0.84,
                                                   temperature=1, 
                                                   repetition_penalty=1.15, 
                                                   eos_token_id=2, 
                                                   bos_token_id=1,
                                                   pad_token_id=0)
                response = self.tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
                response = response.split("### Response:")[-1]
                return response
            except:
                return "对不起，你的请求出错了，请再次尝试。\nSorry, your request has encountered an error. Please try again.\n"
        elif self.mode == 'api':
            return self.predict_api(question)
    
    def message_to_prompt(self, message, system_prompt=""):
        system_prompt = self.prefix_prompt + system_prompt
        for interaction in self.history:
            user_prompt, bot_prompt = str(interaction[0]).strip(' '), str(interaction[1]).strip(' ')
            system_prompt = f"{system_prompt} User: {user_prompt} Bot: {bot_prompt}"
        prompt = f"{system_prompt} ### Instruction:{message.strip()}  ### Response:"
        return prompt
        
    def predict_api(self, question):
        # FastAPI Predict 调用API来进行预测
        self.data["question"] = question
        headers = {'Content-Type': 'application/json'}
        data = {"prompt": question}
        response = requests.post(url=self.url, headers=headers, data=json.dumps(data))
        return response.json()['response']
            
    def chat(self, system_prompt, message, history):
        self.history = history
        prompt = self.message_to_prompt(message, system_prompt)
        response = self.generate(prompt)
        self.history.append([message, response])
        return response, self.history
    
    def clear_history(self):
        # 清空历史记录
        self.history = []

def test():
    llm = Linly(mode='offline',model_path='../Linly-AI/Chinese-LLaMA-2-7B-hf')
    answer = llm.generate("如何应对压力？")
    print(answer)

if __name__ == '__main__':
    test()