from transformers import AutoModelForCausalLM, AutoTokenizer
import argparse
import torch
import types
import modeling_dfm
import numpy as np
import random
import time
parser = argparse.ArgumentParser()
parser.add_argument("--model_name", type=str, default="/lustre/fs11/portfolios/nvr/projects/nvr_elm_llm/users/chengyuew/workspace/LMFlow/output_models/finetune_qwen_1.5B_nemotron_shizhe_no_reasoning_pad_bd_size_128_8nodes_seq_length_2k/checkpoint-15000")
args = parser.parse_args()

model_name = args.model_name

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype="auto",
    device_map="cuda:0",
    trust_remote_code=True
)
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model.mdm_sample = types.MethodType(modeling_dfm.DFM_QwenForCausalLM.mdm_sample_cache_small_block, model)
# model.sample_with_sampling_controls = types.MethodType(modeling_dfm.DFM_QwenForCausalLM.sample_with_sampling_controls, model)

# 初始化对话历史和模式状态
# detailed_thinking_mode = False  # True表示"on"，False表示"off"
messages = [
    {
        "role": "system",
        "content": "You are a helpful assistant."
    }
]

def fix_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)

# def update_system_message():
#     """更新系统消息以反映当前模式"""
#     global messages
#     mode_text = "detailed thinking on" if detailed_thinking_mode else "detailed thinking off"
#     messages[0]["content"] = mode_text

fix_seed(42)
print("聊天机器人已启动！输入 'exit' 退出对话。")
print("输入 'shift' 切换detailed thinking模式。")
print("输入 'clear' 清空对话记录。")
# print(f"当前模式: {'on' if detailed_thinking_mode else 'off'}")
print("-" * 50)

while True:
    # 获取用户输入
    user_input = input("用户: ").strip()

    # 检查是否退出
    if user_input.lower() == "exit":
        print("再见！")
        break

    # 检查是否清空对话记录
    if user_input.lower() == "clear":
        messages = messages[:1] if messages[0]["role"] == "system" else messages[:0]
        print("对话记录已清空！")
        continue

    # 检查是否切换模式
    # if user_input.lower() == "shift":
    #     detailed_thinking_mode = not detailed_thinking_mode
    #     update_system_message()
    #     print(f"模式已切换为: {'on' if detailed_thinking_mode else 'off'}")
    #     continue

    # 如果用户输入为空，继续下一轮

    #
    if not user_input:
        continue

    # messages.append({"role": "user", "content": user_input + "\nPlease reason step by step, and put your final answer within \\boxed{{}}."})
    messages.append({"role": "user", "content": user_input})

    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        # enable_thinking=detailed_thinking_mode
    )
    print(text)
    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
    start_time = time.time()
    generated_ids = model.mdm_sample(
        model_inputs["input_ids"],
        tokenizer=tokenizer,
        block_size=32,
        max_tokens=2048,
        small_block_size=8,
        temperature=0.0,
        top_p=1,
        block_cache_refresh_interval=16,
        threshold=1,
        use_block_cache=False,
        mask_id=tokenizer.encode("|<MASK>|")[0]
    )
    end_time = time.time()
    print(f"生成时间: {end_time - start_time} 秒")
    print(f"生成速度: {(generated_ids.shape[1]-model_inputs['input_ids'].shape[1]) / (end_time - start_time)} tokens/s")
    response = tokenizer.decode(generated_ids[0][model_inputs["input_ids"].shape[1]:], skip_special_tokens=True)
    # response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
    print(f"AI: {response}")

    # 添加AI回复到对话历史
    messages.append({"role": "assistant", "content": response})

    print("-" * 50)
