import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
import coremltools as ct
import torch


# 1. 设定模型路径和运行设备（确保该路径下是转换好的标准 transformers 格式）
model_path = "./minimind-3-tinymodel"  # 换成你的 SFT 模型本地路径或 HF 仓库名
device = "cuda" if torch.cuda.is_available() else "cpu"

# 2. 加载 Tokenizer 和模型
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True).half().to(device)
model.eval()

# 3. 设置流式输出器
streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)

print("🤖 SFT 模型加载完成！请输入你的问题（输入 exit 退出）：\n" + "=" * 40)

# 4. 简单的单轮/多轮交互循环
while True:
    prompt = input("💬 用户: ").strip()
    if not prompt or prompt.lower() == 'exit':
        break

    # 构造符合 SFT 模型的 Chat 结构
    messages = [{"role": "user", "content": prompt}]

    # 核心：使用 apply_chat_template 自动加上模型的特殊 Chat 标记
    input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(input_text, return_tensors="pt", truncation=True).to(device)

    # 5. 模型流式生成
    print("🧠 助手: ", end="")
    with torch.no_grad():
        model.generate(
            input_ids=inputs["input_ids"],
            attention_mask=inputs.get("attention_mask"),
            max_new_tokens=512,
            do_sample=True,
            top_p=0.95,
            temperature=0.9,
            streamer=streamer,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id
        )
    print("\n" + "-" * 40)