import argparse
import os
import platform
import warnings

import re
pattern = re.compile("[\n]+")

import torch
from accelerate import init_empty_weights, load_checkpoint_and_dispatch
from huggingface_hub import snapshot_download

from transformers.generation.utils import logger
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer

parser = argparse.ArgumentParser()
parser.add_argument("--model_name", default="DAMO-NLP-MT/polylm-multialpaca-13b", 
                    choices=["DAMO-NLP-MT/polylm-multialpaca-13b"], type=str)
parser.add_argument("--multi_round", action="store_true",
                    help="Turn multiple rounds interaction on.")
parser.add_argument("--gpu", default="0", type=str)
args = parser.parse_args()

os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
num_gpus = len(args.gpu.split(","))

if args.model_name in ["DAMO-NLP-MT/polylm-multialpaca-13b-int8", "DAMO-NLP-MT/polylm-multialpaca-13b-int4"] and num_gpus > 1:
    raise ValueError("Quantized models do not support model parallel. Please run on a single GPU (e.g., --gpu 0).")

logger.setLevel("ERROR")
warnings.filterwarnings("ignore")

model_path = args.model_name
if not os.path.exists(args.model_name):
    model_path = snapshot_download(args.model_name)

config = AutoConfig.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)

if num_gpus > 1:  
    print("Waiting for all devices to be ready, it may take a few minutes...")
    with init_empty_weights():
        raw_model = AutoModelForCausalLM.from_config(config)
    raw_model.tie_weights()
    model = load_checkpoint_and_dispatch(
        raw_model, model_path, device_map="auto", no_split_module_classes=["GPT2Block"]
    )
else:
    print("Loading model files, it may take a few minutes...")
    model = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto").cuda()

def clear():
    os.system('cls' if platform.system() == 'Windows' else 'clear')
    
def main():
    print("欢迎使用 PolyLM 多语言人工智能助手！输入内容即可进行对话。输入 clear 以清空对话历史，输入 stop 以终止对话。")
    prompt = ""
    while True:
        query = input()
        if query.strip() == "stop":
            break
        if query.strip() == "clear":
            if args.multi_round:
                prompt = ""
            clear()
            continue
        
        text = query.strip()
        text = re.sub(pattern, "\n", text)
        if args.multi_round:
            prompt += f"{text}\n\n"
        else:
            prompt = f"{text}\n\n"
        
        inputs = tokenizer(prompt, return_tensors="pt")
        with torch.no_grad():
            outputs = model.generate(
                inputs.input_ids.cuda(), 
                attention_mask=inputs.attention_mask.cuda(), 
                max_length=1024, 
                do_sample=True, 
                top_p=0.8, 
                temperature=0.7,
                repetition_penalty=1.02,
                num_return_sequences=1, 
                eos_token_id=2,
                early_stopping=True)
            response = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
            if args.multi_round:
                prompt += f"{response}\n"
            print(f">>> {response}")
    
if __name__ == "__main__":
    main()