""" NextActionPredictor model definition. Architecture: 2-hidden-layer MLP (12-dim → 64 → 64 → 6) Input layout (12 dims, indices): [0-5] domain one-hot — ecommerce | telecom | banking | cibil | insurance | general [6-10] entity flags — has_ORG | has_AMOUNT | has_DATE | has_REF_ID | has_ACCOUNT [11] prior_contact — 1.0 if user has already contacted the company, else 0.0 Output: all 6 EscalationActions sorted by confidence descending. Fallback: if no checkpoint exists, DOMAIN_ACTION_PRIORS rule-based mapping is used so the pipeline never crashes. """ from __future__ import annotations import logging import os from dataclasses import dataclass from typing import Optional import torch import torch.nn as nn from src.next_action.priors import ACTION_METADATA, DOMAIN_ACTION_PRIORS logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # Label / feature constants — single source of truth for all modules # --------------------------------------------------------------------------- ACTION_LABELS: list[str] = [ "company_support", "nch", "trai", "rbi_ombudsman", "irdai", "legal" ] ACTION2ID: dict[str, int] = {a: i for i, a in enumerate(ACTION_LABELS)} # Domain ordering must match the one-hot encoding used during training DOMAIN_LABELS_ORDERED: list[str] = [ "ecommerce", "telecom", "banking", "cibil", "insurance", "general" ] # The 5 NER entity types used as routing signals (PERSON excluded — role reference) ENTITY_NAMES: list[str] = ["ORG", "AMOUNT", "DATE", "REF_ID", "ACCOUNT"] FEATURE_DIM: int = 12 # 6 + 5 + 1 NUM_ACTIONS: int = 6 # --------------------------------------------------------------------------- # Feature vector builder — used by train.py and predict.py # --------------------------------------------------------------------------- def build_feature_vector( domain: str, entity_flags: list[float], prior_contact: float, ) -> list[float]: """ Construct the 12-dim feature vector. Args: domain: one of DOMAIN_LABELS_ORDERED entity_flags: 5 floats (0.0/1.0) in ENTITY_NAMES order prior_contact: 1.0 if user already contacted the company, else 0.0 Returns: list[float] of length FEATURE_DIM (12) """ domain_oh = [1.0 if d == domain else 0.0 for d in DOMAIN_LABELS_ORDERED] return domain_oh + list(entity_flags) + [float(prior_contact)] def entities_dict_to_flags(entities: dict) -> list[float]: """ Convert an EvidenceNER entities dict {entity_type: value} → 5-dim flag vector. Example: {"ORG": "Flipkart", "AMOUNT": "₹4,299"} → [1.0, 1.0, 0.0, 0.0, 0.0] """ return [1.0 if name in entities and entities[name] else 0.0 for name in ENTITY_NAMES] # --------------------------------------------------------------------------- # Public output type # --------------------------------------------------------------------------- @dataclass class EscalationAction: """A single recommended escalation step.""" action: str authority: str url: str confidence: float # --------------------------------------------------------------------------- # Raw PyTorch module # --------------------------------------------------------------------------- class GUIDE_MLP(nn.Module): """ 2-hidden-layer MLP: 12 → 64 → ReLU → 64 → ReLU → 6. Returns raw logits; softmax is applied at inference time. """ def __init__( self, input_dim: int = FEATURE_DIM, hidden_dim: int = 64, num_classes: int = NUM_ACTIONS, ) -> None: super().__init__() self.net = nn.Sequential( nn.Linear(input_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, num_classes), ) def forward(self, x: torch.Tensor) -> torch.Tensor: # (B, 12) → (B, 6) return self.net(x) # --------------------------------------------------------------------------- # Rule-based fallback # --------------------------------------------------------------------------- def _rule_based_predict(feature_vector: list[float]) -> list[EscalationAction]: """ Return a ranked EscalationAction list from DOMAIN_ACTION_PRIORS. When prior_contact=1, company_support is deprioritised — the user already tried that path. """ domain_oh = feature_vector[:6] prior_contact = feature_vector[11] domain_idx = int(max(range(6), key=lambda i: domain_oh[i])) domain = DOMAIN_LABELS_ORDERED[domain_idx] ordered: list[str] = list(DOMAIN_ACTION_PRIORS[domain]) # e.g. ["company_support","nch","legal"] if prior_contact > 0.5 and ordered[0] == "company_support": # Already tried company — rotate company_support to last position ordered = ordered[1:] + ordered[:1] # Rank-based exponential score: position 0 → highest decay = 0.55 score_map: dict[str, float] = {} s = 1.0 for action in ordered: score_map[action] = s s *= decay for action in ACTION_LABELS: if action not in score_map: score_map[action] = decay ** len(ordered) # small residual total = sum(score_map.values()) return sorted( [ EscalationAction( action=action, authority=ACTION_METADATA[action]["authority"], url=ACTION_METADATA[action]["url"], confidence=round(score_map[action] / total, 4), ) for action in ACTION_LABELS ], key=lambda e: -e.confidence, ) # --------------------------------------------------------------------------- # NextActionPredictor — public wrapper # --------------------------------------------------------------------------- class NextActionPredictor: """ MLP escalation router with rule-based fallback. Pass model_path=None (or a path that does not exist) to use the rule-based fallback only. The pipeline never crashes: if the checkpoint is missing a warning is logged and DOMAIN_ACTION_PRIORS is used instead. """ def __init__(self, model_path: Optional[str] = None) -> None: self._mlp: Optional[GUIDE_MLP] = None self._device = torch.device( "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu" ) if model_path and os.path.isfile(model_path): try: ckpt = torch.load(model_path, map_location=self._device, weights_only=True) self._mlp = GUIDE_MLP() self._mlp.load_state_dict(ckpt["state_dict"]) self._mlp.to(self._device) self._mlp.eval() logger.info("NextActionPredictor loaded from %s on %s", model_path, self._device) except Exception: logger.warning( "Failed to load NextActionPredictor from %s — using rule-based fallback.", model_path, exc_info=True, ) else: logger.info( "No NextActionPredictor checkpoint at '%s' — using rule-based fallback.", model_path, ) @property def uses_fallback(self) -> bool: return self._mlp is None def predict(self, feature_vector: list[float]) -> list[EscalationAction]: """Return all 6 EscalationActions ranked by confidence (highest first).""" if self._mlp is None: return _rule_based_predict(feature_vector) x = torch.tensor(feature_vector, dtype=torch.float32).unsqueeze(0).to(self._device) with torch.no_grad(): logits = self._mlp(x)[0].cpu() # (6,) probs: list[float] = torch.softmax(logits, dim=-1).tolist() return sorted( [ EscalationAction( action=ACTION_LABELS[i], authority=ACTION_METADATA[ACTION_LABELS[i]]["authority"], url=ACTION_METADATA[ACTION_LABELS[i]]["url"], confidence=round(probs[i], 4), ) for i in range(NUM_ACTIONS) ], key=lambda e: -e.confidence, )