app.py

import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import pandas as pd
import json

# Configuration
MODEL_NAME = "distilbert-base-uncased-finetuned-sst-2-english"  # Using a default model for testing
# Replace with your model: "Rajeshwartiwari/incident-classification-model"

# Default categories (replace with your actual categories)
DEFAULT_CATEGORIES = [
    "Software", "Hardware", "Network", "Database", 
    "Security", "Access", "Email", "VPN", "Application"
]

class IncidentClassifier:
    def __init__(self):
        self.tokenizer = None
        self.model = None
        self.categories = DEFAULT_CATEGORIES
        self.id2label = {i: label for i, label in enumerate(self.categories)}
        self.loaded = False
        
    def load_model(self):
        """Load the model - simplified for Hugging Face Spaces"""
        try:
            print("Loading tokenizer and model...")
            self.tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
            self.model = AutoModelForSequenceClassification.from_pretrained(
                MODEL_NAME,
                num_labels=len(self.categories)
            )
            self.loaded = True
            return "✅ Model loaded successfully!"
        except Exception as e:
            return f"❌ Error loading model: {str(e)}"
    
    def predict(self, short_desc, detailed_desc=""):
        """Make prediction"""
        if not self.loaded:
            return "Please load the model first", {}
            
        if not short_desc.strip():
            return "Please enter incident description", {}
        
        # Combine text
        full_text = f"{short_desc} {detailed_desc}".strip()
        
        # Tokenize
        inputs = self.tokenizer(
            full_text, 
            return_tensors="pt", 
            truncation=True, 
            max_length=128
        )
        
        # Predict
        with torch.no_grad():
            outputs = self.model(**inputs)
            probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
        
        # Get predictions
        predicted_id = torch.argmax(probabilities).item()
        predicted_label = self.id2label.get(predicted_id % len(self.categories), "Unknown")
        
        # Get all confidence scores
        confidences = {}
        for idx, prob in enumerate(probabilities[0]):
            if idx < len(self.categories):
                label = self.categories[idx]
                confidences[label] = float(prob) * 100
            else:
                break
        
        # Sort by confidence
        sorted_confidences = dict(sorted(confidences.items(), key=lambda x: x[1], reverse=True))
        
        return predicted_label, sorted_confidences
    
    def batch_predict(self, csv_file):
        """Process batch of incidents from CSV"""
        if not self.loaded:
            return pd.DataFrame({"Error": ["Model not loaded"]})
        
        try:
            # Read CSV
            df = pd.read_csv(csv_file.name)
            
            results = []
            for idx, row in df.iterrows():
                short_desc = str(row.get('short_description', row.get('description', '')))
                if len(short_desc) > 100:
                    short_display = short_desc[:100] + "..."
                else:
                    short_display = short_desc
                
                predicted, confidences = self.predict(short_desc)
                top_conf = max(confidences.values()) if confidences else 0
                
                results.append({
                    "Incident": short_display,
                    "Predicted Category": predicted,
                    "Confidence": f"{top_conf:.1f}%",
                    "Top 3 Predictions": ", ".join([
                        f"{k}: {v:.1f}%" 
                        for k, v in list(confidences.items())[:3]
                    ])
                })
            
            return pd.DataFrame(results)
            
        except Exception as e:
            return pd.DataFrame({"Error": [f"Failed to process CSV: {str(e)}"]})

# Initialize classifier
classifier = IncidentClassifier()

# Create Gradio interface
def create_interface():
    with gr.Blocks(title="Incident Classifier", theme=gr.themes.Soft()) as demo:
        gr.Markdown("""
        # 🔧 Incident Classification System
        Automatically classify and route IT incidents to the correct team
        """)
        
        # Model Status Section
        with gr.Row():
            status_box = gr.Textbox(
                label="Model Status", 
                value="⚠️ Model not loaded", 
                interactive=False
            )
            load_btn = gr.Button("🔄 Load Model", variant="primary")
        
        # Single Incident Classification
        with gr.Row():
            with gr.Column():
                gr.Markdown("### 📝 Single Incident")
                short_desc = gr.Textbox(
                    label="Short Description*",
                    placeholder="Brief description of the issue...",
                    lines=2
                )
                detailed_desc = gr.Textbox(
                    label="Detailed Description (Optional)",
                    placeholder="Additional details, error messages...",
                    lines=3
                )
                predict_btn = gr.Button("🔍 Classify", variant="primary")
                
                # Results
                with gr.Row():
                    prediction = gr.Textbox(label="Predicted Category", interactive=False)
                    confidence = gr.Textbox(label="Top Confidence", interactive=False)
                
                confidences_chart = gr.Label(
                    label="Confidence Scores",
                    num_top_classes=5
                )
        
        # Batch Processing
        with gr.Row():
            with gr.Column():
                gr.Markdown("### 📁 Batch Processing")
                gr.Markdown("Upload CSV with 'short_description' column")
                
                file_input = gr.File(
                    label="Upload CSV",
                    file_types=[".csv"],
                    type="filepath"
                )
                batch_btn = gr.Button("📊 Process Batch", variant="secondary")
                batch_output = gr.Dataframe(label="Results")
        
        # Examples
        gr.Markdown("### 💡 Example Incidents")
        examples = gr.Examples(
            examples=[
                ["Oracle database connection error ORA-12154", "Users cannot connect to production Oracle DB"],
                ["Outlook not syncing emails", "Email client stopped receiving new messages"],
                ["VPN keeps disconnecting", "Cisco AnyConnect drops connection every 5 minutes"],
                ["Monitor screen flickering", "Display flickers with dark backgrounds"]
            ],
            inputs=[short_desc, detailed_desc],
            label="Try these examples:"
        )
        
        # Event Handlers
        def on_load_model():
            message = classifier.load_model()
            if classifier.loaded:
                return f"✅ Model loaded! {len(classifier.categories)} categories available"
            return message
        
        def on_predict(short, detailed):
            if not classifier.loaded:
                return "Please load model first", {}, "0%"
            
            predicted, confidences = classifier.predict(short, detailed)
            top_conf = max(confidences.values()) if confidences else 0
            
            return predicted, confidences, f"{top_conf:.1f}%"
        
        # Connect events
        load_btn.click(
            fn=on_load_model,
            outputs=status_box
        )
        
        predict_btn.click(
            fn=on_predict,
            inputs=[short_desc, detailed_desc],
            outputs=[prediction, confidences_chart, confidence]
        )
        
        batch_btn.click(
            fn=classifier.batch_predict,
            inputs=file_input,
            outputs=batch_output
        )
    
    return demo

# Launch the app
if __name__ == "__main__":
    # Try to load model on startup
    print("Initializing Incident Classifier...")
    
    # Create and launch interface
    demo = create_interface()
    
    # For Hugging Face Spaces, use share=False
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        debug=True
    )