import gradio as gr
from transformers import pipeline

# 加载模型
pipe = pipeline(
    "text-generation",
    model="Name-isname/olmo3-190m-zh-full",
    device_map="auto"
)

def respond(message, history):
    # ChatInterface 的 history 格式通常是 [[user_msg1, bot_msg1], [user_msg2, bot_msg2]]
    # 我们直接在这里构建 prompt
    full_prompt = ""
    for user_input, bot_response in history:
        full_prompt += f"用户: {user_input}\n助手: {bot_response}\n"
    
    full_prompt += f"用户: {message}\n助手:"

    # 生成回复（注意：不要在 pipe 中传入 history 参数，只传 prompt）
    output = pipe(
        full_prompt,
        max_new_tokens=256,
        do_sample=True,
        temperature=0.7,
        top_k=50,
        top_p=0.9,
        repetition_penalty=1.2
    )
    
    # 提取生成的文本
    generated_text = output[0]["generated_text"]
    # 只返回新生成的回复内容
    response = generated_text.replace(full_prompt, "").strip()
    return response

# 使用最稳妥的 ChatInterface 写法
demo = gr.ChatInterface(fn=respond)

if __name__ == "__main__":
    demo.launch(theme='soft')