File size: 2,387 Bytes
a1ca7f2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 | # inference.py - Simple inference script
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import gradio as gr
class IncidentClassifier:
def __init__(self, model_name="Rajeshwartiwari/incident-classification-model"):
self.model_name = model_name
self.tokenizer = None
self.model = None
self.id2label = None
def load(self):
"""Load the model and tokenizer"""
print(f"Loading model: {self.model_name}")
self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
self.model = AutoModelForSequenceClassification.from_pretrained(self.model_name)
self.id2label = self.model.config.id2label
print(f"Model loaded with {len(self.id2label)} categories")
return self
def predict(self, text):
"""Make prediction on input text"""
if not text.strip():
return "Please enter text", {}
# Tokenize
inputs = self.tokenizer(
text,
return_tensors="pt",
truncation=True,
padding=True,
max_length=128
)
# Predict
with torch.no_grad():
outputs = self.model(**inputs)
probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
# Get top prediction
predicted_id = outputs.logits.argmax().item()
predicted_label = self.id2label.get(predicted_id, "Unknown")
# Get all confidences
confidences = {}
for idx, prob in enumerate(probabilities[0]):
label = self.id2label.get(idx, f"Class_{idx}")
confidences[label] = float(prob) * 100
return predicted_label, confidences
# Quick test
if __name__ == "__main__":
classifier = IncidentClassifier().load()
test_cases = [
"Oracle database connection error ORA-12154",
"Email not syncing in Outlook",
"VPN keeps disconnecting every few minutes",
"Monitor screen flickering issues"
]
for test in test_cases:
label, confidences = classifier.predict(test)
top_3 = dict(sorted(confidences.items(), key=lambda x: x[1], reverse=True)[:3])
print(f"\n📝 Input: {test}")
print(f" 🎯 Predicted: {label}")
print(f" 📊 Top 3: {top_3}") |