import gradio as gr
from transformers import pipeline

# 加载模型
model_name = "Dominic0406/convai2_gpt2"
pipe = pipeline("text-generation", model=model_name)

def generate_response(dialog):
    # 对话是一个字符串，使用分隔符分割多轮对话
    dialog = dialog.split("\n")
    input_text = " ".join(dialog)
    response = pipe(input_text, max_length=50, num_return_sequences=1)[0]['generated_text']
    return response

# 创建Gradio接口
iface = gr.Interface(
    fn=generate_response,
    inputs=gr.Textbox(lines=7, placeholder="Enter your dialog here..."),
    outputs="text",
    title="GPT-2 Conversational AI",
    description="Enter your dialog, one message per line. The model will generate a response based on the entire conversation."
)

# 启动接口
if __name__ == "__main__":
    iface.launch()