import time
import torch
import transformers


model_dir = './models/Qwen2.5-0.5B-Instruct'
device_map = 'cuda'
template = '你是一个逻辑严密、极具常识、幽默风趣的 AI 助手，请仔细思考用户提问，并简洁准确地回答问题。'


tokenizer = transformers.AutoTokenizer.from_pretrained(
    pretrained_model_name_or_path = model_dir,
    local_files_only = True,
    trust_remote_code = True
)


model = transformers.AutoModelForCausalLM.from_pretrained(
    pretrained_model_name_or_path = model_dir,
    local_files_only = True,
    trust_remote_code = True,
    dtype = torch.bfloat16,
    device_map = device_map,
    low_cpu_mem_usage = True
)
model.eval()


streamer = transformers.TextStreamer(
    tokenizer = tokenizer,
    skip_prompt = True,
    skip_special_tokens = True
)


while True:
    content = input('💁 : ')
    if content.strip().lower() in ['quit', 'exit']:
        print('\n👋👋👋👋👋👋👋👋👋👋👋👋👋👋👋👋')
        break
    if not content.strip():
        continue

    messages = [
        {'role': 'system', 'content': template},
        {'role': 'user', 'content': content},
        # {'role': 'assistant', 'content': ''}
    ]


    prompt = tokenizer.apply_chat_template(
        messages,
        tokenize = True,
        add_generation_prompt = True,
        return_tensors = 'pt',
        return_dict = True
    ).to(model.device)

    print('🤖 : ', end = '', flush = True)

    with torch.no_grad():
        start_time = time.time()

        outputs = model.generate(
            **prompt,
            max_length = None,
            max_new_tokens = 4096,
            temperature = 0.6,
            top_p = 0.9,
            repetition_penalty = 1.1,
            do_sample = True,
            streamer = streamer,
            pad_token_id = tokenizer.eos_token_id
        )

        cost_time = time.time() - start_time

        tokens = outputs[0].shape[0] - prompt['input_ids'].shape[1]
        tps = tokens / cost_time if cost_time > 0 else 0

        response: str = tokenizer.decode(
            outputs[0][prompt['input_ids'].shape[1]:],
            skip_special_tokens = True
        )
    print('')
    print(f'✅ cost time = {cost_time:.3f} s | generate token = {tokens} | speed = {tps:.3f} tokens/s\n')