Update app.py

app.py

@@ -1,26 +1,43 @@
 import streamlit as st
-from transformers import pipeline
+from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
+import torch
 
-def main():
-    # Define pipelines
-    summarizer_ntg = pipeline(model="mrm8488/t5-base-finetuned-summarize-news")
-    model = pipeline(model="Lauraayu/News_Classi_Model")
+# Define the summarization pipeline
+summarizer_ntg = pipeline("summarization", model="mrm8488/t5-base-finetuned-summarize-news")
+
+# Load the tokenizer and model for classification
+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:")
+
+# 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']
     
-    # Streamlit application title
-    st.title("News Article Classifier")
-    st.write("Enter a news article text to get its category:")
+    # Tokenize the summarized text
+    inputs = tokenizer_bb(summary, return_tensors="pt", truncation=True, padding=True, max_length=512)
     
-    # Text input for user to enter the news article text
-    text = st.text_area("Enter the news article text here:")
+    # Move inputs and model to the same device (GPU or CPU)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    inputs = {k: v.to(device) for k, v in inputs.items()}
+    model_bb.to(device)
     
-    # Perform summarization and classification when the user clicks the "Classify" button
-    if st.button("Classify"):
-        # Perform text summarization
-        summary = summarizer_ntg(text)[0]
-        
-        # Perform classification
-        output = model(summary)
-        category = output[0]["label"]
-
-        # Display the summary and classification result
-        st.write("Category:", category)
+    # 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)