import gradio as gr
from transformers import pipeline

# 加载预训练的摘要模型 "shleifer/distilbart-cnn-12-6"
summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6")

# 定义摘要函数
def summarize(input_text):
    # 使用模型生成摘要
    summary = summarizer(input_text, max_length=130, min_length=30, do_sample=False)
    return summary[0]['summary_text']
    
def summarize_with_error_handling(input_text):
    try:
        return summarize(input_text)
    except Exception as e:
        return f"An error occurred: {str(e)}"
        
# 使用 Gradio 构建前端
gr.close_all()  # 关闭其他可能在运行的 Gradio 应用
demo = gr.Interface(fn=summarize_with_error_handling, inputs="text", outputs="text", title="Text Summarization with distilbart-cnn-12-6")


# 启动 Gradio 应用
demo.launch()