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 )