Update app.py

app.py

@@ -1,29 +1,46 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelWithLMHead
+from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
+import torch
 
-# 加载模型和分词器
+# Load the tokenizer and model for classification
 tokenizer = AutoTokenizer.from_pretrained("mrm8488/t5-base-finetuned-summarize-news")
 model = AutoModelWithLMHead.from_pretrained("mrm8488/t5-base-finetuned-summarize-news")
 
-# 定义摘要函数
+tokenizer_bb = AutoTokenizer.from_pretrained("your-username/your-model-name")
+model_bb = AutoModelForSequenceClassification.from_pretrained("your-username/your-model-name")
+
+# Streamlit application title
+st.title("News Article Summarizer and Classifier")
+st.write("Enter a news article text to get its summary and category.")
+
+# Text input for user to enter the news article text
+text = st.text_area("Enter the news article text here:")
+
 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.")
+    
+# Perform summarization and classification when the user clicks the "Classify" button
+if st.button("Classify"):
+    # Perform text summarization
+    with st.spinner("Generating summary..."):
+            summary = summarize(article)
+    
+    # Tokenize the summarized text
+    inputs = tokenizer_bb(summary, return_tensors="pt", truncation=True, padding=True, max_length=512)
+    
+    # Perform text classification
+    with torch.no_grad():
+        outputs = model_bb(**inputs)
+    
+    # Get the predicted label
+    predicted_label_id = torch.argmax(outputs.logits, dim=-1).item()
+    label_mapping = model_bb.config.id2label
+    predicted_label = label_mapping[predicted_label_id]
+    
+    # Display the summary and classification result
+    st.write("Summary:", summary)
+    st.write("Category:", predicted_label)