File size: 861 Bytes
78ffd72 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 |
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
def load_model():
model_path = "model"
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForSequenceClassification.from_pretrained(model_path)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
model.eval()
return tokenizer, model, device
def classify_email(text, tokenizer, model, device):
inputs = tokenizer(text, return_tensors="pt", max_length=256, padding="max_length", truncation=True)
inputs = {k: v.to(device) for k, v in inputs.items()}
with torch.no_grad():
logits = model(**inputs).logits
label_map = {0: "Incident", 1: "Request", 2: "Change", 3: "Problem"}
pred = torch.argmax(logits, dim=1).item()
return label_map[pred]
|