Update app.py

app.py

@@ -1,22 +1,29 @@
 import streamlit as st
-from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
-import torch
+from transformers import AutoTokenizer, AutoModelWithLMHead
 
-# Define the summarization pipeline
-summarizer_ntg = pipeline("text2text-generation", model="mrm8488/t5-base-finetuned-summarize-news")
+# 加载模型和分词器
+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 application title
-st.title("News Article Summarizer and Classifier")
-st.write("Enter a news article text to get its summary and category.")
+# Streamlit 应用程序界面
+st.title("News Summarization App")
+st.write("Enter the news article text below to generate a summary.")
 
-# Text input for user to enter the news article text
-text = st.text_area("Enter the news article text here:")
+article = st.text_area("News Article", height=300)
+max_len = st.slider("Max Length of Summary", min_value=50, max_value=300, value=150)
 
-# Perform summarization and classification when the user clicks the "Classify" button
-if st.button("Classify"):
-    # Perform text summarization
-    summary = summarizer_ntg(text)[0]['summary_text']
-        
-    # Display the summary and classification result
-    st.write("Summary:", summary)
+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.")