Lauraayu's picture
Update app.py
5888fc0 verified
raw
history blame
1.29 kB
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.")