Update app.py

app.py

@@ -1,13 +1,3 @@
-# Upgrade pip to avoid dependency issues
-pip install --upgrade pip  
-
-# Install PyTorch and related libraries
-pip install torch torchvision torchaudio  
-
-# Verify installation
-python -c "import torch; print(torch.__version__)"  
-
-
 import streamlit as st
 import pandas as pd
 import torch
@@ -15,19 +5,24 @@ from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
 # Load the fine-tuned model
 MODEL_NAME = "dinusha11/finetuned-distilbert-news"
+
+# Ensure model runs on CPU if GPU is unavailable
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
+model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME).to(device)
 
-# Define label mapping
+# Define label mapping (Ensure this matches your model's label order)
 LABEL_MAPPING = {0: "Business", 1: "Opinion", 2: "Sports", 3: "Political_gossip", 4: "World_news"}
 
-# Function to classify text
-def classify_text(text):
-    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=128)
+# Function to classify a batch of text
+def classify_texts(texts):
+    inputs = tokenizer(texts, return_tensors="pt", padding=True, truncation=True, max_length=128)
+    inputs = {key: val.to(device) for key, val in inputs.items()}  # Move inputs to device
     with torch.no_grad():
         outputs = model(**inputs)
-    predicted_class = torch.argmax(outputs.logits, dim=1).item()
-    return LABEL_MAPPING[predicted_class]
+    predicted_classes = torch.argmax(outputs.logits, dim=1).cpu().numpy()
+    return [LABEL_MAPPING[pred] for pred in predicted_classes]
 
 # Streamlit UI
 st.title("News Classification App")
@@ -38,18 +33,19 @@ uploaded_file = st.file_uploader("Upload CSV", type=["csv"])
 
 if uploaded_file:
     df = pd.read_csv(uploaded_file)
-    
+
     if "text" not in df.columns:
         st.error("CSV must contain a 'text' column.")
     else:
-        # Preprocess text
+        # Preprocess text (handle missing values, strip spaces, and convert to lowercase)
         df["text"] = df["text"].fillna("").str.strip().str.lower()
-        
-        # Apply classification
-        df["class"] = df["text"].apply(classify_text)
-        
+
+        # Apply batch classification
+        df["class"] = classify_texts(df["text"].tolist())
+
         # Download output CSV
         output_csv = df.to_csv(index=False).encode("utf-8")
         st.download_button("Download Results", data=output_csv, file_name="output.csv", mime="text/csv")
 
         st.write("Classification Complete! Download your file above.")
+