import os
os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"

from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
import torch

# 1. 基础模型和分词器路径
base_model_name = "./model"
lora_path = "./lora"

# 【核心修改】自动检测设备：如果是 Apple Silicon 则使用 mps，否则用 cpu
if torch.backends.mps.is_available():
    device = torch.device("mps")
    print("已检测到 Apple Silicon，将使用 MPS 进行硬件加速 🚀")
else:
    device = torch.device("cpu")
    print("未检测到 MPS，将使用 CPU")

# 加载分词器
tokenizer = AutoTokenizer.from_pretrained(lora_path)

# 加载原模型（在 Mac 上建议去掉 device_map="auto"，直接用 .to(device) 更稳定）
base_model = AutoModelForCausalLM.from_pretrained(
    base_model_name,
    torch_dtype=torch.float16
).to(device)

# 调整模型的 Token Embeddings 大小
base_model.resize_token_embeddings(len(tokenizer))

# 加载 LoRA 权重
model = PeftModel.from_pretrained(base_model, lora_path)

# 测试对话
while True:
    prompt = input(">  ")
    prompt = "Human: " + prompt + "\nBot:"
    if prompt == "exit":
        break
    # 【核心修改】将输入移动到 device（即 mps）
    inputs = tokenizer(prompt, return_tensors="pt").to(device)

    outputs = model.generate(
        **inputs,
        max_new_tokens=100,
        temperature=0.7,
        top_p=0.9,
        do_sample=True,
        pad_token_id=tokenizer.eos_token_id,
        repetition_penalty=1.2,
        no_repeat_ngram_size=3
    )

    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    print(response.replace(" ", ""))