import gradio as gr, os
from transformers import BartForConditionalGeneration

# 加载 BART 模型
model = BartForConditionalGeneration.from_pretrained("facebook/bart-large-cnn")

def generate_summary(file):
    # 重置文件指针位置
    file.seek(0)
    
    # 读取上传的文本文件内容
    text_content = file.read()

    # 使用模型进行处理（摘要生成）
    summary_ids = model.generate(text_content, max_length=150, min_length=50, length_penalty=2.0, num_beams=4, early_stopping=True)
    summary = model.decode(summary_ids[0], skip_special_tokens=True)

    return summary

demo = gr.Interface(
    fn=generate_summary,
    inputs=gr.File(),
    outputs="text",
    live=False
)

# 启动应用
demo.launch(share=True)

# 加载 BART 模型
model = BartForConditionalGeneration.from_pretrained("models/fnlp/bart-base-chinese")

def generate_summary(file):
    # 重置文件指针位置
    file.seek(0)
    
    # 读取上传的文本文件内容
    text_content = file.read()

    # 使用模型进行处理（摘要生成）
    summary_ids = model.generate(text_content, max_length=150, min_length=50, length_penalty=2.0, num_beams=4, early_stopping=True)
    summary = model.decode(summary_ids[0], skip_special_tokens=True)

    return summary

demo = gr.Interface(
    fn=generate_summary,
    inputs=gr.File(),
    outputs="text",
    live=False
)

# 启动应用
demo.launch(share=True)