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

# 加载模型和分词器
model_name = "XiaomiMiMo/MiMo-7B-RL"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map="auto",
    trust_remote_code=True
)
@spaces.GPU(duration=120)
def predict(message, history):
    # 构建输入
    history_text = ""
    for human, assistant in history:
        history_text += f"Human: {human}\nAssistant: {assistant}\n"
    prompt = f"{history_text}Human: {message}\nAssistant:"
    
    # 生成回复
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    
    # 使用流式生成
    streamer = tokenizer.decode
    response = ""
    
    for outputs in model.generate(
        **inputs,
        max_new_tokens=10000,
        do_sample=True,
        temperature=0.7,
        top_p=0.9,
        repetition_penalty=1.1,
        pad_token_id=tokenizer.eos_token_id,
        stream_output=True
    ):
        next_token = outputs[0][inputs.input_ids.shape[1]:]
        next_token_text = streamer(next_token, skip_special_tokens=True)
        response += next_token_text
        yield response.strip()

# 创建Gradio界面
demo = gr.ChatInterface(
    predict,
    title="MiMo-7B-RL 聊天机器人",
    description="这是一个基于小米 MiMo-7B-RL 模型的聊天机器人。",
    examples=["你好！", "请介绍一下你自己", "你能做什么？"],
    theme=gr.themes.Soft()
)

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