Update src/inference.py

src/inference.py

@@ -1,74 +1,30 @@
 import joblib
 import torch
-import numpy as np
-from collections import Counter
-from sklearn.preprocessing import LabelEncoder
-from sklearn.model_selection import train_test_split
-from torch.utils.data import Dataset, DataLoader
-from src.data_processing import load_and_clean_data
+import torch.nn.functional as F
+from src.data_processing import clean_text
 from src.model_def import EmotionTransformer
 
-# Hyperparameters
-MAX_LEN = 32
-BATCH_SIZE = 16
-EPOCHS = 5
-LR = 1e-3
+# Load artifacts
+vocab = joblib.load('vocab.pkl')
+le    = joblib.load('label_encoder.pkl')
 DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-# Dataset wrapper
-dclass EmotionDataset(Dataset):
-    def __init__(self, X, y):
-        self.X = torch.tensor(X, dtype=torch.long)
-        self.y = torch.tensor(y, dtype=torch.long)
-    def __len__(self): return len(self.X)
-    def __getitem__(self, idx): return self.X[idx], self.y[idx]
-
-
-def train():
-    df = load_and_clean_data()
-    toks = df['clean'].str.split()
-
-    # Build vocab
-    ctr = Counter(tok for sent in toks for tok in sent)
-    vocab = {w:i+2 for i,(w,_) in enumerate(ctr.most_common())}
-    vocab['<PAD>'], vocab['<UNK>'] = 0, 1
-    joblib.dump(vocab, 'vocab.pkl')
-
-    # Encode + pad
-    X = [
-        ([vocab.get(tok,1) for tok in sent] + [0]*(MAX_LEN-len(sent)))[:MAX_LEN]
-        for sent in toks
-    ]
+# 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()
 
-    # Encode labels
-    le = LabelEncoder()
-    y = le.fit_transform(df['label'])
-    joblib.dump(le, 'label_encoder.pkl')
-
-    # Split & loader
-    X_tr, X_va, y_tr, y_va = train_test_split(X, y, test_size=0.2,
-                                             stratify=y, random_state=42)
-    tr_loader = DataLoader(EmotionDataset(X_tr, y_tr), batch_size=BATCH_SIZE, shuffle=True)
-    va_loader = DataLoader(EmotionDataset(X_va, y_va), batch_size=BATCH_SIZE)
-
-    # Model, optimizer, loss
-    model = EmotionTransformer(len(vocab), num_classes=len(le.classes_)).to(DEVICE)
-    opt = torch.optim.Adam(model.parameters(), lr=LR)
-    crit = torch.nn.CrossEntropyLoss()
-
-    # Training loop
-    for epoch in range(EPOCHS):
-        model.train(); total_loss = 0
-        for xb, yb in tr_loader:
-            xb, yb = xb.to(DEVICE), yb.to(DEVICE)
-            opt.zero_grad()
-            loss = crit(model(xb), yb)
-            loss.backward(); opt.step()
-            total_loss += loss.item()
-        print(f"Epoch {epoch+1}/{EPOCHS} Loss: {total_loss/len(tr_loader):.4f}")
-
-    # Save weights
-    torch.save(model.state_dict(), 'emotion_transformer_model.pth')
+MAX_LEN = 32
 
-if __name__=='__main__':
-    train()
+def predict(text):
+    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)
+    return le.inverse_transform([avg.argmax().item()])[0]