model / app.py
Nikpatil's picture
Create app.py
b732103 verified
Raw History Blame Contribute Delete
851 Bytes
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()