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

tokenizer = AutoTokenizer.from_pretrained("openbmb/MiniCPM3-4B")

model = AutoModelForCausalLM.from_pretrained("openbmb/MiniCPM3-4B", trust_remote_code=True)

# 定义生成回复的函数
def generate_response(input_text):
    inputs = tokenizer.encode(input_text, return_tensors="pt")
    # 为生成模型配置生成参数
    outputs = model.generate(inputs, max_length=200, num_return_sequences=1, do_sample=True, top_k=50, top_p=0.95)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

# 使用 Gradio 创建聊天界面
iface = gr.Interface(
    fn=generate_response,
    inputs="text",
    outputs="text",
    title="MiniCPM-3 中文聊天机器人",
    description="这是一个基于 MiniCPM-3 的简单聊天机器人，可以进行中文对话"
)

# 启动 Gradio 应用
iface.launch()