Update app.py

app.py

@@ -1,67 +1,125 @@
 import gradio as gr
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
 import torch
+import pandas as pd
 import os
 
-model_hub_path = "AJC1/ag_news_distilbert_finetuned" 
-target_names = ["World", "Sports", "Business", "Sci/Tech"]
+# --- Configuration ---
+# Use your verified public model path
+MODEL_HUB_PATH = "AJC1/ag_news_distilbert_finetuned" 
+TARGET_NAMES = ["World", "Sports", "Business", "Sci/Tech"]
 
-
-# Load Model and Tokenizer
+# --- 1. Load Model (Cached for performance) ---
 try:
-    # Load model and tokenizer directly from the Hugging Face Model Hub
-    tokenizer = AutoTokenizer.from_pretrained(model_hub_path)
-    model = AutoModelForSequenceClassification.from_pretrained(model_hub_path)
-    model.eval() 
-    print(f"Model loaded successfully from Hugging Face Hub: {model_hub_path}")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_HUB_PATH)
+    model = AutoModelForSequenceClassification.from_pretrained(MODEL_HUB_PATH)
+    model.eval()
+    print("Model loaded successfully.")
 except Exception as e:
-    print(f"Failed to load model from Hub. Check the repository name. Error: {e}")
+    print(f"Error loading model: {e}")
 
+# --- 2. Prediction Functions ---
 
-# Define the Prediction Function
-def classify_news(text):
-    """Takes input text, tokenizes it, and returns the predicted class name and confidence scores."""
-    if not text or len(text.strip()) < 5:
-        return "Please enter a longer news snippet or headline.", {}
-    
+def predict_single_text(text):
+    """The core logic: text -> dict of scores"""
+    if not text: return {}
     inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
-    
     with torch.no_grad():
-        outputs = model(**inputs)
-        logits = outputs.logits
-        probabilities = torch.softmax(logits, dim=1)[0].tolist()
+        logits = model(**inputs).logits
+    probs = torch.softmax(logits, dim=1)[0].tolist()
+    return {TARGET_NAMES[i]: v for i, v in enumerate(probs)}
+
+def process_csv_file(file_obj):
+    """The enterprise logic: CSV -> Classified CSV"""
+    try:
+        # Load the uploaded CSV
+        df = pd.read_csv(file_obj.name)
         
-    confidence = {
-        target_names[i]: prob for i, prob in enumerate(probabilities)
-    }
+        # Validation: Check if it has text
+        if df.empty:
+            return None, "Error: Uploaded file is empty."
+            
+        # Smart column detection: Look for 'text', 'headline', or use the first column
+        target_col = None
+        for col in ['text', 'headline', 'title', 'content']:
+            if col in df.columns:
+                target_col = col
+                break
+        if not target_col:
+            target_col = df.columns[0] # Fallback to first column
+
+        # Run predictions (Iterating for safety, batching could be faster but more complex)
+        predicted_labels = []
+        confidence_scores = []
+        
+        for text in df[target_col].astype(str):
+            scores = predict_single_text(text)
+            # Get the top label
+            top_label = max(scores, key=scores.get)
+            predicted_labels.append(top_label)
+            confidence_scores.append(f"{scores[top_label]:.2f}")
+
+        # Add results to dataframe
+        df['Predicted_Category'] = predicted_labels
+        df['Confidence'] = confidence_scores
+        
+        # Save to a temporary output file
+        output_path = "classified_results.csv"
+        df.to_csv(output_path, index=False)
+        
+        return output_path, f"Success! Processed {len(df)} rows. Download your results below."
+
+    except Exception as e:
+        return None, f"Error processing file: {str(e)}"
+
+# --- 3. The Professional Tabbed Interface ---
+with gr.Blocks(title="AG News Enterprise Classifier") as demo:
     
-    return confidence
+    gr.Markdown("# 📰 Automated News Routing System")
+    gr.Markdown("Select a workflow below: Single-item checking or Bulk file processing.")
+
+    with gr.Tabs():
+        
+        # === TAB 1: Single Input (For Demo/Editors) ===
+        with gr.TabItem("Live Check"):
+            with gr.Row():
+                with gr.Column():
+                    text_input = gr.Textbox(lines=4, label="Input News Headline", placeholder="Paste text here...")
+                    submit_btn = gr.Button("Classify Content", variant="primary")
+                    
+                with gr.Column():
+                    label_output = gr.Label(num_top_classes=4, label="Category Prediction")
+            
+            # Link functionality
+            submit_btn.click(fn=predict_single_text, inputs=text_input, outputs=label_output)
+            
+            # Examples
+            gr.Examples(
+                examples=[
+                    ["Wall Street tumbles as tech stocks sell off."],
+                    ["Manchester United signs new striker for record fee."],
+                    ["NASA discovers water on Mars surface."]
+                ],
+                inputs=text_input
+            )
 
-# Create the Gradio Interface 
-iface = gr.Interface(
-    fn=classify_news,
-    inputs=gr.Textbox(
-        lines=5, 
-        label="Input News Headline or Snippet",
-        placeholder="Example: Russia and Canada discuss gas pipeline auction..."
-    ),
-    outputs=gr.Label(
-        num_top_classes=4, 
-        label="Predicted Category and Confidence"
-    ), 
-    title="AG News Classifier: Fine-Tuned DistilBERT",
-    description=(
-        "Enter a news headline or short article. The model will predict its category "
-        "(World, Sports, Business, or Sci/Tech) and show the confidence scores for all categories."
-    ),
-    examples=[
-        ["The global stock market rallied today after the central bank cut interest rates."], 
-        ["New research shows quantum entanglement may enable faster computing."], 
-        ["Manchester United defeats Liverpool in a stunning Premier League match."],
-        ["The President's cabinet held an emergency summit on trade negotiations."],
-        ["AT&T Wireless ships mobile IM gadget US mobile network operator AT&T Wireless today launched Ogo, its first non-voice messaging device, pitche..."]
-    ]
-)
+        # === TAB 2: Batch Processing (For Operations) ===
+        with gr.TabItem("Bulk Analysis (CSV)"):
+            gr.Markdown("Upload a CSV file containing news headlines. The system will append a 'Category' column and return the file.")
+            
+            with gr.Row():
+                file_input = gr.File(label="Upload CSV File", file_types=[".csv"])
+                file_output = gr.File(label="Download Classified Results")
+            
+            status_text = gr.Textbox(label="Status", interactive=False)
+            process_btn = gr.Button("Process Batch", variant="primary")
+            
+            # Link functionality
+            process_btn.click(
+                fn=process_csv_file, 
+                inputs=file_input, 
+                outputs=[file_output, status_text]
+            )
 
 if __name__ == "__main__":
-    iface.launch()
+    demo.launch()