Update src/inference.py

src/inference.py

@@ -1,43 +1,30 @@
+import joblib
 import torch
 import torch.nn.functional as F
-import joblib
-
 from src.data_processing import clean_text
-from src.model_def       import EmotionTransformer
+from src.model_def import EmotionTransformer
 
-# 1. Load your saved artifacts:
-le    = joblib.load("label_encoder.pkl")
-vocab = joblib.load("vocab.pkl")  # saved vocab dict
+# Load artifacts
+vocab = joblib.load('vocab.pkl')
+le    = joblib.load('label_encoder.pkl')
+DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-# 2. Recreate & load your model
-model = EmotionTransformer(
-    vocab_size=len(vocab), embed_dim=64, num_heads=4,
-    num_classes=len(le.classes_)
-)
-model.load_state_dict(
-    torch.load("emotion_transformer_model.pth", map_location=device)
-)
+# Recreate model
+model = EmotionTransformer(len(vocab), num_classes=len(le.classes_)).to(DEVICE)
+model.load_state_dict(torch.load('emotion_transformer_model.pth', map_location=DEVICE))
 model.eval()
-MAX_LEN = 32
-
-# 3. Preprocess + predict
-
-def preprocess_input(text):
-    # note: your pad_sequence logic from the notebook
-    tokens  = clean_text(text).split()
-    encoded = [vocab.get(tok, vocab['<UNK>']) for tok in tokens]
-    padded  = encoded[:MAX_LEN] + [vocab['<PAD>']] * max(0, MAX_LEN - len(encoded))
-    return torch.tensor([padded], dtype=torch.long).to(device)
 
+MAX_LEN = 32
 
 def predict(text):
-    x = preprocess_input(text)
-    # Bayesian-style dropout inference
+    toks = clean_text(text).split()
+    idxs = [vocab.get(tok,1) for tok in toks]
+    pad  = (idxs + [0]*MAX_LEN)[:MAX_LEN]
+    x = torch.tensor([pad], dtype=torch.long).to(DEVICE)
+
+    # MC-dropout inference
     model.train()
     with torch.no_grad():
         probs = torch.stack([F.softmax(model(x), dim=1) for _ in range(5)])
-        avg   = probs.mean(dim=0)
-    idx = avg.argmax(dim=1).item()
-    return le.inverse_transform([idx])[0]
+    avg = probs.mean(dim=0)
+    return le.inverse_transform([avg.argmax().item()])[0]