import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# Load model & tokenizer
model = AutoModelForSequenceClassification.from_pretrained("final_model")
tokenizer = AutoTokenizer.from_pretrained("final_model")
id2label = {0: 'Incident', 1: 'Request', 2: 'Problem', 3: 'Change'}

def predict(text):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=256)
    with torch.no_grad():
        outputs = model(**inputs)
        probs = torch.nn.functional.softmax(outputs.logits, dim=1)[0]
        predicted_class_id = torch.argmax(probs).item()
        return {
            id2label[i]: float(probs[i]) for i in range(len(probs))
        }

# UI
demo = gr.Interface(fn=predict, inputs="text", outputs="label", title="Email Classifier")

# Launch
demo.launch()