models.py

from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import os

def load_model():
    model_path = "sathish2352/email-classifier-model"
    
    # Set HF_HOME to use a writable cache dir
    os.environ["HF_HOME"] = "/tmp/huggingface"
    os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface/transformers"
    os.makedirs("/tmp/huggingface/transformers", exist_ok=True)

    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]