import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# 1. 指定你的模型 ID
model_id = "Jobfromearth/MobileLLM-80M-Finetuned"
print(f"🚀 正在初始化... 加载模型: {model_id}")

# 自动检测设备：如果有显卡驱动则用 cuda，否则用 cpu
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"检测到运行设备: {device}")

# 2. 加载模型
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id, 
    torch_dtype=torch.float32,
    device_map=device,       # 动态指定设备
    trust_remote_code=True, 
    low_cpu_mem_usage=True
)

# 2. 定义聊天逻辑
def chat_response(message, history):
    # 1. 定义分隔符
    separator = "#"
    
    # 2. 自动判断用户输入
    if separator in message:
        instruction_part, input_part = message.split(separator, 1)
        instruction = instruction_part.strip()
        input_context = input_part.strip()
    else:
        instruction = message
        input_context = ""

    # 3. 严格构造 Prompt
    prompt = f"""Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.### Instruction:{instruction}### Input:{input_context}### Response:"""
    
    # --- 修复部分 ---
    # 将 .to("cuda") 改为 .to(device)
    inputs = tokenizer([prompt], return_tensors="pt").to(device)
    
    outputs = model.generate(
        **inputs,
        max_new_tokens=64,
        repetition_penalty=1.2,
        do_sample=True,
        temperature=0.1, 
        top_p=0.9,
        use_cache=True,
        eos_token_id=tokenizer.eos_token_id,
        pad_token_id=tokenizer.eos_token_id,
    )
    
    response = tokenizer.batch_decode(outputs)[0]
    
    if "### Response:" in response:
        response = response.split("### Response:")[-1]
    if "### Instruction:" in response:
        response = response.split("### Instruction:")[0]
        
    return response.strip().replace(tokenizer.eos_token, "")

# 3. 启动界面
demo = gr.ChatInterface(
    fn=chat_response,
    title="MobileLLM-60M ChatBot (Pro)",
    description="Tip: For translation or rewriting tasks, it is recommended to use '#' to separate instructions and content for better results.",
    examples=[
        "The capital of France is?",
        "Translate to Swedish # The weather is very good today.",
        "Classify sentiment # I hate waiting in line, it is so annoying.",
        "Extract the name and age to JSON # My friend Alice is 25 years old."
    ],
)

if __name__ == "__main__":
    demo.launch()