import streamlit as st
from transformers import AutoTokenizer, AutoModelWithLMHead

# 加载模型和分词器
tokenizer = AutoTokenizer.from_pretrained("mrm8488/t5-base-finetuned-summarize-news")
model = AutoModelWithLMHead.from_pretrained("mrm8488/t5-base-finetuned-summarize-news")

# 定义摘要函数
def summarize(text, max_length=150):
    input_ids = tokenizer.encode(text, return_tensors="pt", add_special_tokens=True)
    generated_ids = model.generate(input_ids=input_ids, num_beams=2, max_length=max_length, repetition_penalty=2.5, length_penalty=1.0, early_stopping=True)
    preds = [tokenizer.decode(g, skip_special_tokens=True, clean_up_tokenization_spaces=True) for g in generated_ids]
    return preds[0]

# Streamlit 应用程序界面
st.title("News Summarization App")
st.write("Enter the news article text below to generate a summary.")

article = st.text_area("News Article", height=300)
max_len = st.slider("Max Length of Summary", min_value=50, max_value=300, value=150)

if st.button("Summarize"):
    if article:
        with st.spinner("Generating summary..."):
            summary = summarize(article, max_length=max_len)
            st.write("**Summary:**")
            st.write(summary)
    else:
        st.error("Please enter some text to summarize.")