File size: 851 Bytes
b732103 d26b21b b732103 d26b21b b732103 d26b21b b732103 d26b21b b732103 | 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 | 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()
|