Upload src\ensemble_router.py with huggingface_hub

src//ensemble_router.py

@@ -0,0 +1,482 @@
+# src/ensemble_router.py
+# SupportMind — Ensemble Confidence-Gated Router
+# Combines DistilBERT (MC Dropout) + TF-IDF Logistic Regression
+# for best-in-class accuracy on ticket routing.
+#
+# Strategy: weighted soft-voting on probability distributions
+#   final_probs = w_bert * bert_probs + w_sklearn * sklearn_probs
+#
+# Why this beats either model alone:
+#   - DistilBERT: captures semantic meaning, handles paraphrases
+#   - TF-IDF+LR : captures keyword/n-gram signals, very confident on clear cases
+#   - Ensemble  : DistilBERT corrects LR on ambiguous tickets,
+#                 LR corrects BERT on keyword-heavy ones
+
+import os
+import gc
+import pickle
+import logging
+import numpy as np
+from typing import Dict, Optional
+
+logger = logging.getLogger(__name__)
+
+# ── Category map ────────────────────────────────────────────────────────────
+CATEGORY_MAP = {
+    0: 'billing',
+    1: 'technical_support',
+    2: 'account_management',
+    3: 'feature_request',
+    4: 'compliance_legal',
+    5: 'onboarding',
+    6: 'general_inquiry',
+    7: 'churn_risk',
+}
+CATEGORY_REVERSE = {v: k for k, v in CATEGORY_MAP.items()}
+
+# ── Routing thresholds ───────────────────────────────────────────────────────
+ROUTE_THRESHOLD   = 0.82   # ensemble conf >= this → auto-route
+CLARIFY_THRESHOLD = 0.58   # ensemble conf >= this → ask 1 question
+ENTROPY_MAX       = 0.32   # ensemble entropy <= this → low ambiguity
+MC_PASSES         = int(os.getenv('SUPPORTMIND_MC_PASSES', '3'))  # CPU demo default
+
+# ── Ensemble weights ─────────────────────────────────────────────────────────
+# BERT weight is higher because it generalises better to unseen phrasing.
+# These are tunable — increase SKLEARN_W if LR is more accurate on your data.
+# BERT weight is significantly higher because DeBERTa-v3 is extremely robust.
+BERT_W   = 0.75
+SKLEARN_W = 0.25
+
+
+class EnsembleRouter:
+    """
+    Ensemble Confidence-Gated Router.
+
+    Combines:
+      1. DistilBERT fine-tuned on support tickets (MC Dropout for uncertainty)
+      2. TF-IDF + Calibrated Logistic Regression baseline
+
+    Falls back to sklearn-only if DistilBERT model weights are absent.
+    Drop-in replacement for ConfidenceGatedRouter — same .route() interface.
+    """
+
+    def __init__(self, model_dir: Optional[str] = None, device: str = 'cpu'):
+        base = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+        ultimate_path = os.path.join(base, 'models', 'deberta_ultimate')
+        standard_path = os.path.join(base, 'models', 'ticket_classifier')
+        
+        if model_dir is None:
+            if os.path.exists(os.path.join(ultimate_path, 'config.json')):
+                self.model_dir = ultimate_path
+            else:
+                self.model_dir = standard_path
+        else:
+            self.model_dir = model_dir
+
+        self._bert_router = None
+        self._sklearn_pipe = None
+        self._bert_available = False
+        self._bert_reason = 'not_loaded'
+        self._sklearn_source = 'unknown'
+
+        # IMPORTANT: Load BERT first and do a warmup pass.
+        # On Windows, unpickling sklearn before PyTorch's first forward pass
+        # causes a segfault in torch.distributed/optree DLLs.
+        self._load_bert(device)
+        if self._bert_available:
+            self._warmup_bert()
+        self._load_sklearn()
+
+        try:
+            from historical_memory import HistoricalMemoryLayer
+            self._memory_layer = HistoricalMemoryLayer()
+        except Exception as e:
+            logger.warning(f"[EnsembleRouter] Could not load Historical Memory Layer: {e}")
+            self._memory_layer = None
+
+        self.model_status = {
+            'mode': 'ensemble_transformer_lr' if self._bert_available else 'sklearn_fallback',
+            'bert_available': self._bert_available,
+            'bert_reason': self._bert_reason,
+            'sklearn_source': self._sklearn_source,
+            'model_dir': os.path.relpath(self.model_dir, base),
+            'memory_available': bool(
+                getattr(getattr(self, '_memory_layer', None), 'is_ready', False)
+            ),
+        }
+
+        logger.info(
+            f"[EnsembleRouter] BERT={'ON' if self._bert_available else 'OFF (fallback)'} | "
+            f"sklearn=ON | weights=({BERT_W}/{SKLEARN_W}) | memory={'ON' if getattr(self, '_memory_layer', None) and self._memory_layer.is_ready else 'OFF'}"
+        )
+
+    def _warmup_bert(self):
+        """Perform a warmup forward pass to initialize PyTorch/CUDA state."""
+        try:
+            self._bert_router.mc_predict("warmup", n_passes=1)
+            logger.info("[EnsembleRouter] BERT warmup complete.")
+        except Exception as e:
+            logger.warning(f"[EnsembleRouter] BERT warmup failed: {e}")
+
+    # ── Model loaders ────────────────────────────────────────────────────────
+
+    def _load_sklearn(self):
+        # Check model_dir first, then fall back to ticket_classifier
+        pkl = os.path.join(self.model_dir, 'sklearn_router.pkl')
+        if not os.path.exists(pkl):
+            base = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+            pkl = os.path.join(base, 'models', 'ticket_classifier', 'sklearn_router.pkl')
+        if not os.path.exists(pkl):
+            logger.warning(
+                "[EnsembleRouter] sklearn_router.pkl not found. "
+                "Using embedded synthetic fallback model."
+            )
+            self._sklearn_pipe = self._build_embedded_sklearn()
+            self._sklearn_source = 'embedded_synthetic'
+            return
+        with open(pkl, 'rb') as f:
+            self._sklearn_pipe = pickle.load(f)
+        base = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+        self._sklearn_source = os.path.relpath(pkl, base)
+        logger.info(f"[EnsembleRouter] sklearn pipeline loaded from {pkl}.")
+
+    def _build_embedded_sklearn(self):
+        """Build a tiny in-memory classifier so clean clones and CI still run."""
+        from sklearn.feature_extraction.text import TfidfVectorizer
+        from sklearn.linear_model import LogisticRegression
+        from sklearn.pipeline import Pipeline
+
+        examples = {
+            'billing': [
+                'invoice is wrong', 'refund request', 'payment failed',
+                'billing charge incorrect', 'subscription price changed',
+            ],
+            'technical_support': [
+                'api returns 500 error', 'export is broken', 'dashboard crash',
+                'integration timeout', 'feature not working',
+            ],
+            'account_management': [
+                'reset password', 'add user account', 'sso login issue',
+                'change admin permission', 'locked out of account',
+            ],
+            'feature_request': [
+                'please add dark mode', 'new feature request',
+                'need custom dashboard', 'enhancement idea',
+            ],
+            'compliance_legal': [
+                'gdpr data request', 'soc 2 audit report',
+                'data processing agreement', 'privacy compliance',
+            ],
+            'onboarding': [
+                'help with setup', 'new user onboarding',
+                'configure integration', 'getting started guide',
+            ],
+            'general_inquiry': [
+                'how do i use this', 'pricing question', 'where is documentation',
+                'do you offer a demo',
+            ],
+            'churn_risk': [
+                'cancel my account', 'switching to competitor',
+                'very frustrated', 'not renewing contract',
+            ],
+        }
+
+        texts, labels = [], []
+        for category, samples in examples.items():
+            for sample in samples:
+                texts.append(sample)
+                labels.append(CATEGORY_REVERSE[category])
+
+        pipeline = Pipeline([
+            ('tfidf', TfidfVectorizer(stop_words='english', ngram_range=(1, 2))),
+            ('clf', LogisticRegression(class_weight='balanced', max_iter=1000)),
+        ])
+        pipeline.fit(texts, labels)
+        return pipeline
+
+    def _load_bert(self, device: str):
+        """Load transformer router when the runtime is configured for it."""
+        disable_transformer = os.getenv('SUPPORTMIND_DISABLE_TRANSFORMER', '0') == '1'
+        force_transformer = os.getenv('SUPPORTMIND_FORCE_TRANSFORMER', '0') == '1'
+
+        if disable_transformer:
+            self._bert_reason = 'disabled_by_SUPPORTMIND_DISABLE_TRANSFORMER'
+            logger.warning("[EnsembleRouter] Transformer loading disabled by environment.")
+            return
+
+        if os.name == 'nt' and not force_transformer:
+            self._bert_reason = 'disabled_on_windows_set_SUPPORTMIND_FORCE_TRANSFORMER_to_enable'
+            logger.warning(
+                "[EnsembleRouter] Transformer loading disabled on Windows by default "
+                "to avoid native PyTorch/safetensors crashes. Set "
+                "SUPPORTMIND_FORCE_TRANSFORMER=1 to enable it."
+            )
+            return
+
+        import json, traceback as tb
+        model_bin  = os.path.join(self.model_dir, 'pytorch_model.bin')
+        model_safe = os.path.join(self.model_dir, 'model.safetensors')
+        config     = os.path.join(self.model_dir, 'config.json')
+
+        bert_ready = os.path.exists(config) and (
+            os.path.exists(model_bin) or os.path.exists(model_safe)
+        )
+
+        if not bert_ready:
+            self._bert_reason = 'weights_not_found'
+            logger.warning(
+                "[EnsembleRouter] DistilBERT weights not found — running sklearn-only."
+            )
+            return
+
+        # Check for stale baseline stub (only present before first real training run)
+        try:
+            with open(config) as f:
+                cfg = json.load(f)
+            if cfg.get('model_type') == 'baseline_sklearn':
+                self._bert_reason = 'baseline_stub_config'
+                logger.warning("[EnsembleRouter] config.json is baseline stub — skipping BERT.")
+                return
+        except Exception:
+            pass
+
+        try:
+            from confidence_router import ConfidenceGatedRouter
+            self._bert_router = ConfidenceGatedRouter(self.model_dir, device=device)
+            self._bert_available = not getattr(self._bert_router, '_fallback_mode', False)
+            fallback_reason = getattr(self._bert_router, 'fallback_reason', None)
+            self._bert_reason = (
+                'loaded' if self._bert_available
+                else f'confidence_router_fallback: {fallback_reason or "unknown"}'
+            )
+            gc.collect()
+            if self._bert_available:
+                logger.info(f"[EnsembleRouter] {self._bert_router.model.config.model_type.upper()} loaded successfully.")
+        except (Exception, OSError) as e:
+            logger.error(f"[EnsembleRouter] BERT load failed (likely memory constraint): {e}")
+            # Ensure we don't leave a half-initialized router
+            self._bert_router = None
+            self._bert_available = False
+            self._bert_reason = f'load_failed: {type(e).__name__}'
+            gc.collect()
+
+    # ── Prediction ───────────────────────────────────────────────────────────
+
+    def _sklearn_probs(self, text: str) -> np.ndarray:
+        """Return calibrated probability distribution from sklearn pipeline."""
+        return self._sklearn_pipe.predict_proba([text])[0]   # shape [8]
+
+    def _bert_probs(self, text: str) -> np.ndarray:
+        """Return MC-Dropout probability distribution from DistilBERT."""
+        _, _, _, mean_p, _ = self._bert_router.mc_predict(text, n_passes=MC_PASSES)
+        return mean_p   # shape [8]
+
+    def _blend(self, text: str):
+        """
+        Compute blended probability distribution.
+        Returns: (blended_probs, bert_probs_or_None, sklearn_probs, bert_std_or_None)
+        """
+        sk_probs = self._sklearn_probs(text)
+
+        if self._bert_available:
+            _, _, _, bert_mean, bert_std = self._bert_router.mc_predict(text, MC_PASSES)
+            blended = BERT_W * bert_mean + SKLEARN_W * sk_probs
+            # Re-normalise (floating point can drift slightly)
+            blended = blended / blended.sum()
+            return blended, bert_mean, sk_probs, bert_std
+        else:
+            return sk_probs, None, sk_probs, np.zeros(8)
+
+    # ── Public API ───────────────────────────────────────────────────────────
+
+    def route(self, ticket_text: str, n_passes: int = MC_PASSES) -> Dict:
+        """
+        Route a ticket through the ensemble confidence gate.
+        Returns the same dict schema as ConfidenceGatedRouter.route()
+        so it is a drop-in replacement in api.py.
+        """
+        blended, bert_p, sk_p, bert_std = self._blend(ticket_text)
+
+        confidence  = float(blended.max())
+        entropy     = float(-np.sum(blended * np.log(blended + 1e-9)))
+        
+        # ── Temperature Scaling (T=0.7) ──────────────────────────────────
+        # Sharpen probabilities to reduce noise in unrelated classes.
+        # logits_scaled = logits / T; softmax(logits_scaled)
+        # Since we have probs, we can approximate with power scaling:
+        # p_scaled = p^(1/T) / sum(p^(1/T))
+        T = 0.7
+        blended_sharp = np.power(blended + 1e-9, 1.0 / T)
+        blended_sharp = blended_sharp / blended_sharp.sum()
+        
+        # ── Keyword Reinforcement ────────────────────────────────────────
+        # If text contains specific strong keywords for a category, 
+        # give that category a small 'calibration boost'.
+        reinforce_map = {
+            'billing': ['invoice', 'refund', 'charge', 'payment', 'billing'],
+            'technical_support': ['error', 'bug', 'crash', '500', 'api', 'broken', 'not working'],
+            'account_management': ['login', 'password', 'reset', 'account', 'permission', 'access', 'sso', 'user'],
+            'feature_request': ['feature', 'add', 'request', 'enhancement', 'dark mode', 'new capability', 'could you add'],
+            'compliance_legal': ['gdpr', 'compliance', 'legal', 'audit', 'privacy'],
+            'churn_risk': ['cancel', 'leaving', 'competitor', 'terminate', 'switching'],
+            'onboarding': ['setup', 'configure', 'getting started', 'new user', 'import'],
+        }
+        text_low = ticket_text.lower()
+        for cat, kws in reinforce_map.items():
+            hit_count = sum(1 for kw in kws if kw in text_low)
+            if hit_count:
+                idx = CATEGORY_REVERSE[cat]
+                blended_sharp[idx] *= 1.0 + min(0.45, hit_count * 0.12)
+                blended_sharp[idx] += min(0.12, hit_count * 0.03)
+        
+        # Re-normalise after boost
+        blended_sharp = blended_sharp / blended_sharp.sum()
+        
+        confidence = float(blended_sharp.max())
+        pred_class = int(blended_sharp.argmax())
+        category   = CATEGORY_MAP[pred_class]
+        
+        # ── Visual Confidence Cap (98.5%) ────────────────────────────────
+        # Probabilistic ML should rarely claim 100% certainty.
+        display_confidence = min(confidence, 0.985)
+
+        # Build ranking
+        ranking = sorted(
+            [(CATEGORY_MAP[i], round(float(blended_sharp[i]), 4)) for i in range(8)],
+            key=lambda x: x[1], reverse=True
+        )
+        top_two = [ranking[0][0], ranking[1][0]]
+
+        base = {
+            'confidence':       round(display_confidence, 4),
+            'raw_confidence':   round(confidence, 4),
+            'entropy':          round(entropy,    4),
+            'top_category':     category,
+            'all_probs':        {CATEGORY_MAP[i]: round(float(blended_sharp[i]), 4) for i in range(8)},
+            'std_probs':        {CATEGORY_MAP[i]: round(float(bert_std[i]), 4) for i in range(8)},
+            'category_ranking': ranking,
+            'top_two_classes':  top_two,
+            'mc_passes':        n_passes,
+            # Extra ensemble diagnostics
+            'ensemble': {
+                'bert_available':  self._bert_available,
+                'bert_top':        CATEGORY_MAP[int(bert_p.argmax())] if bert_p is not None else None,
+                'sklearn_top':     CATEGORY_MAP[int(sk_p.argmax())],
+                'bert_weight':     BERT_W if self._bert_available else 0.0,
+                'sklearn_weight':  SKLEARN_W if self._bert_available else 1.0,
+                'agreement':       (
+                    CATEGORY_MAP[int(bert_p.argmax())] == CATEGORY_MAP[int(sk_p.argmax())]
+                    if bert_p is not None else True
+                ),
+            }
+        }
+
+        top1_score = ranking[0][1]
+        top2_score = ranking[1][1]
+        margin = top1_score - top2_score
+        
+        hist_boost = 0.0
+        if getattr(self, '_memory_layer', None) and self._memory_layer.is_ready:
+            hist_boost = self._memory_layer.compute_historical_boost(ticket_text, category)
+            base['historical_boost'] = hist_boost
+
+        base['margin'] = round(margin, 4)
+        base['confidence'] = round(display_confidence, 4)
+
+        critical_labels = ['compliance_legal', 'account_management']
+        
+        effective_conf = confidence + hist_boost
+
+        if category in critical_labels:
+            if effective_conf >= 0.90 and margin >= 0.35 and entropy < 0.60:
+                action = 'route'
+                reason = f'• Safe to auto-route sensitive intent<br>• Confidence: {confidence:.2%}<br>• Margin: {margin:.2f}'
+                if hist_boost > 0: reason += f'<br>• <span style="color:var(--green)">Historical Match Boost: +{hist_boost:.2%}</span>'
+            else:
+                action = 'escalate'
+                reason = f'• Escalated sensitive intent ({category})<br>• Strict confidence/margin threshold not met'
+                if hist_boost > 0: reason += f'<br>• <span style="color:var(--green)">Historical Match Boost: +{hist_boost:.2%}</span> (Insufficient)'
+        elif category == 'technical_support':
+            # Category-specific check for technical support to catch billing misroutes
+            billing_keywords = ['invoice', 'billing', 'charge', 'refund', 'payment', 'subscription', 'plan']
+            has_billing_kw = any(kw in ticket_text.lower() for kw in billing_keywords)
+            
+            if has_billing_kw and 'billing' in [r[0] for r in ranking[:3]]:
+                action = 'clarify'
+                reason = f'• Billing overlap detected<br>• Clarification needed between technical_support and billing'
+            elif effective_conf >= 0.88 and margin >= 0.30 and entropy < 0.65:
+                # Stricter thresholds for technical_support
+                action = 'route'
+                reason = f'• Strong dominant intent<br>• Confidence: {confidence:.2%}<br>• Margin: {margin:.2f}<br>• Safe to auto-route'
+                if hist_boost > 0: reason += f'<br>• <span style="color:var(--green)">Historical Match Boost: +{hist_boost:.2%}</span>'
+            elif effective_conf >= 0.60 and entropy < 1.05:
+                action = 'clarify'
+                reason = f'• Medium ambiguity detected<br>• Clarification needed between {top_two[0]} and {top_two[1]}<br>• Margin: {margin:.2f}'
+                if hist_boost > 0: reason += f'<br>• <span style="color:var(--green)">Historical Match Boost: +{hist_boost:.2%}</span> (Insufficient for auto-route)'
+            else:
+                action = 'escalate'
+                reason = f'• High ambiguity / Low confidence ({confidence:.2%})<br>• Multiple overlapping intents detected<br>• Human triage needed'
+        else:
+            if effective_conf >= 0.85 and margin >= 0.25 and entropy < 0.70:
+                action = 'route'
+                reason = f'• Strong dominant intent<br>• Confidence: {confidence:.2%}<br>• Margin: {margin:.2f}<br>• Safe to auto-route'
+                if hist_boost > 0: reason += f'<br>• <span style="color:var(--green)">Historical Match Boost: +{hist_boost:.2%}</span>'
+            elif effective_conf >= 0.60 and entropy < 1.05:
+                action = 'clarify'
+                reason = f'• Medium ambiguity detected<br>• Clarification needed between {top_two[0]} and {top_two[1]}<br>• Margin: {margin:.2f}'
+                if hist_boost > 0: reason += f'<br>• <span style="color:var(--green)">Historical Match Boost: +{hist_boost:.2%}</span> (Insufficient for auto-route)'
+            else:
+                action = 'escalate'
+                reason = f'• High ambiguity / Low confidence ({confidence:.2%})<br>• Multiple overlapping intents detected<br>• Human triage needed'
+
+        return {**base, 'action': action, 'queue': category if action == 'route' else None, 'reason': reason}
+
+    def batch_route(self, tickets: list, n_passes: int = MC_PASSES) -> list:
+        return [self.route(t, n_passes) for t in tickets]
+
+    # Property to expose model/tokenizer for the SHAP explainer in api.py
+    @property
+    def model(self):
+        if self._bert_available:
+            return self._bert_router.model
+        return None
+
+    @property
+    def tokenizer(self):
+        if self._bert_available:
+            return self._bert_router.tokenizer
+        return None
+
+
+if __name__ == '__main__':
+    logging.basicConfig(level=logging.INFO, format='%(asctime)s [%(levelname)s] %(message)s')
+
+    router = EnsembleRouter()
+
+    tests = [
+        "My invoice from last month is incorrect, please fix the billing.",
+        "The API keeps returning 500 errors since last Tuesday's update.",
+        "I want to cancel — this tool has been broken for weeks.",
+        "How do I add another user to our account?",
+        "We need GDPR data processing agreements for our EU customers.",
+        "Not happy at all, considering switching to a competitor.",
+        "Can you add a dark mode to the dashboard?",
+        "Just signed up — how do I import my existing data?",
+        # Tricky ambiguous cases
+        "Invoice is wrong AND the app keeps crashing.",
+        "Not happy with service",
+    ]
+
+    print(f"\n{'='*90}")
+    print(f"  SupportMind Ensemble Router — BERT={'ON' if router._bert_available else 'OFF (sklearn only)'}")
+    print(f"{'='*90}\n")
+
+    for ticket in tests:
+        r = router.route(ticket)
+        agree = 'AGREE' if r['ensemble']['agreement'] else 'DISAGREE'
+        print(
+            f"[{r['action'].upper():8s}] [{r['confidence']:.2%}] "
+            f"{'H' if r['entropy'] < ENTROPY_MAX else 'L'}-certainty | "
+            f"{r['top_category']:20s} | "
+            f"Models: {agree} | {ticket[:60]}"
+        )