Update app.py

app.py

@@ -3,11 +3,17 @@ from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassifica
 import torch
 
 # Define the summarization pipeline
-summarizer_ntg = pipeline("text2text-generation", model="mrm8488/t5-base-finetuned-summarize-news")
+try:
+    summarizer_ntg = pipeline("text2text-generation", model="mrm8488/t5-base-finetuned-summarize-news")
+except Exception as e:
+    st.error(f"Error loading summarization model: {e}")
 
 # Load the tokenizer and model for classification
-tokenizer_bb = AutoTokenizer.from_pretrained("Lauraayu/News_Classi_Model")
-model_bb = AutoModelForSequenceClassification.from_pretrained("Lauraayu/News_Classi_Model")
+try:
+    tokenizer_bb = AutoTokenizer.from_pretrained("Lauraayu/News_Classi_Model")
+    model_bb = AutoModelForSequenceClassification.from_pretrained("Lauraayu/News_Classi_Model")
+except Exception as e:
+    st.error(f"Error loading classification model or tokenizer: {e}")
 
 # Streamlit application title
 st.title("News Article Summarizer and Classifier")
@@ -18,26 +24,32 @@ 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']
-    
-    # Tokenize the summarized text
-    inputs = tokenizer_bb(summary, return_tensors="pt", truncation=True, padding=True, max_length=512)
-    
-    # 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 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)
+    if not text:
+        st.error("Please enter some text to classify.")
+    else:
+        try:
+            # Perform text summarization
+            summary = summarizer_ntg(text)[0]['summary_text']
+            
+            # Tokenize the summarized text
+            inputs = tokenizer_bb(summary, return_tensors="pt", truncation=True, padding=True, max_length=512)
+            
+            # 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 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)
+        except Exception as e:
+            st.error(f"Error during summarization or classification: {e}")