Update language/intent.py

language/intent.py

@@ -2,10 +2,11 @@ import torch
 import torch.nn as nn
 from core.device import DEVICE
 
+
 class IntentClassifier(nn.Module):
     def __init__(self, input_dim: int, intent_labels: list):
         """
-        input_dim: dimension of sentence embeddings (HIDDEN_DIM)
+        input_dim: dimension of sentence embeddings
         intent_labels: list of intent names
         """
         super().__init__()
@@ -13,23 +14,25 @@ class IntentClassifier(nn.Module):
         self.intent_labels = intent_labels
         self.num_intents = len(intent_labels)
 
-        self.classifier = nn.Linear(input_dim, self.num_intents)
+        self.classifier = nn.Sequential(
+            nn.Linear(input_dim, 512),
+            nn.ReLU(),
+            nn.Dropout(0.3),
+            nn.Linear(512, 256),
+            nn.ReLU(),
+            nn.Dropout(0.3),
+            nn.Linear(256, self.num_intents)
+        )
+
         self.softmax = nn.Softmax(dim=-1)
 
         self.to(DEVICE)
 
     def forward(self, sentence_embedding: torch.Tensor) -> torch.Tensor:
-        """
-        sentence_embedding: (batch_size, input_dim)
-        returns logits: (batch_size, num_intents)
-        """
         sentence_embedding = sentence_embedding.to(DEVICE)
         return self.classifier(sentence_embedding)
 
-    def predict(self, sentence_embedding: torch.Tensor) -> dict:
-        """
-        Returns intent label + confidence
-        """
+    def predict(self, sentence_embedding: torch.Tensor, threshold=0.4) -> dict:
         with torch.no_grad():
             logits = self.forward(sentence_embedding)
             probs = self.softmax(logits)
@@ -37,7 +40,26 @@ class IntentClassifier(nn.Module):
             confidence, idx = torch.max(probs, dim=-1)
             label = self.intent_labels[idx.item()]
 
+            if confidence.item() < threshold:
+                label = "question_general"
+
             return {
                 "intent": label,
                 "confidence": confidence.item()
             }
+
+
+# ------------------------
+# RULE-BASED OVERRIDE
+# ------------------------
+
+def override_intent(text: str, predicted: str) -> str:
+    text = text.lower().strip()
+
+    if "who are you" in text:
+        return "self_identity"
+
+    if text.startswith(("hi", "hello", "hey")):
+        return "greeting"
+
+    return predicted