Spaces:
Sleeping
Sleeping
| """ | |
| HMM Belief Updater for streaming intent classification. | |
| Implements Markov filtering in LOG-SPACE for numerical stability: | |
| - Predict step: log_b_pred = logsumexp(log_T + log_b_prev) | |
| - Update step: log_b = log_b_pred + log_emission - logsumexp(...) | |
| Supports multiple emission transforms for ablation: | |
| - 'power': emission = (p + eps)^alpha (sharpening) | |
| - 'temperature': emission = softmax(logits / T) | |
| - 'isotonic': calibrated probabilities (requires calibrator) | |
| - 'raw': pass-through neural posteriors | |
| """ | |
| import numpy as np | |
| from enum import Enum | |
| from typing import Dict, Optional, Tuple, Callable | |
| from .config import StreamingConfig | |
| class EmissionTransform(Enum): | |
| """Emission transform methods for ablation studies.""" | |
| POWER = "power" # (p + eps)^alpha sharpening | |
| TEMPERATURE = "temperature" # softmax(logits / T) | |
| RAW = "raw" # Pass-through neural posteriors | |
| ISOTONIC = "isotonic" # Calibrated (requires external calibrator) | |
| def logsumexp(log_vec: np.ndarray) -> float: | |
| """ | |
| Numerically stable log-sum-exp. | |
| Edge cases: | |
| - If all values are -inf, returns -inf (empty probability mass) | |
| - If any value is +inf, returns +inf (handles numerical overflow gracefully) | |
| """ | |
| max_val = np.max(log_vec) | |
| if np.isinf(max_val): | |
| return max_val | |
| return max_val + np.log(np.sum(np.exp(log_vec - max_val))) | |
| def log_normalize(log_vec: np.ndarray) -> np.ndarray: | |
| """Normalize log-probabilities to sum to 1 in probability space.""" | |
| return log_vec - logsumexp(log_vec) | |
| class BeliefUpdater: | |
| """ | |
| HMM-style belief state updater with log-space arithmetic. | |
| Maintains a probability distribution over intents and updates it | |
| using transition dynamics and emission observations from neural model. | |
| All internal computations are done in log-space to prevent underflow | |
| in long sequences. Probabilities are converted only for output. | |
| Supports ablation modes: | |
| - use_hmm=False: Raw neural posteriors (no temporal smoothing) | |
| - emission_transform: Different emission sharpening methods | |
| """ | |
| def __init__( | |
| self, | |
| config: StreamingConfig, | |
| use_hmm: bool = True, | |
| emission_transform: EmissionTransform = EmissionTransform.POWER, | |
| isotonic_calibrator: Optional[Callable] = None, | |
| ): | |
| """ | |
| Initialize belief updater. | |
| Args: | |
| config: Streaming configuration with intents, transitions, thresholds. | |
| use_hmm: If False, bypass HMM and return raw neural posteriors. | |
| emission_transform: Method for transforming neural outputs to emissions. | |
| isotonic_calibrator: Optional calibrator function for ISOTONIC mode. | |
| """ | |
| self.config = config | |
| self.n_intents = len(config.intents) | |
| self.use_hmm = use_hmm | |
| self.emission_transform = emission_transform | |
| self.isotonic_calibrator = isotonic_calibrator | |
| # Initialize log-space belief | |
| prior = config.get_prior_vector() | |
| self.log_belief = np.log(prior + config.epsilon) | |
| self.log_belief = log_normalize(self.log_belief) | |
| # Pre-compute log transition matrix | |
| self.log_T = np.log(config.transition_matrix + config.epsilon) | |
| self._step_count = 0 | |
| self._last_raw_probs: Optional[Dict[str, float]] = None | |
| def reset(self) -> None: | |
| """Reset belief to prior distribution.""" | |
| prior = self.config.get_prior_vector() | |
| self.log_belief = np.log(prior + self.config.epsilon) | |
| self.log_belief = log_normalize(self.log_belief) | |
| self._step_count = 0 | |
| self._last_raw_probs = None | |
| def predict(self) -> np.ndarray: | |
| """ | |
| Prediction step in log-space: propagate belief through transition matrix. | |
| log_b_pred[j] = logsumexp_i(log_T[i,j] + log_b[i]) | |
| Returns: | |
| Log-space predicted belief state. | |
| """ | |
| log_b_pred = np.zeros(self.n_intents) | |
| for j in range(self.n_intents): | |
| # Sum over all previous states i: T[i,j] * b[i] | |
| log_terms = self.log_T[:, j] + self.log_belief | |
| log_b_pred[j] = logsumexp(log_terms) | |
| return log_normalize(log_b_pred) | |
| def compute_log_emission( | |
| self, | |
| neural_probs: Dict[str, float], | |
| logits: Optional[Dict[str, float]] = None, | |
| ) -> np.ndarray: | |
| """ | |
| Compute log-emission likelihood from neural network outputs. | |
| Supports multiple transform methods for ablation: | |
| - POWER: log_emission = alpha * log(p + eps) | |
| - TEMPERATURE: log_emission = logits / T - logsumexp(logits / T) | |
| - RAW: log_emission = log(p + eps) | |
| - ISOTONIC: log_emission = log(calibrate(p) + eps) | |
| Args: | |
| neural_probs: Dict mapping intent names to probabilities. | |
| logits: Optional dict of raw logits for temperature scaling. | |
| Returns: | |
| Log-emission likelihood vector. | |
| """ | |
| log_emission = np.full(self.n_intents, np.log(self.config.epsilon)) | |
| if self.emission_transform == EmissionTransform.POWER: | |
| # Power transform: sharpen probabilities | |
| for intent, prob in neural_probs.items(): | |
| if intent in self.config.intent_to_idx: | |
| idx = self.config.intent_to_idx[intent] | |
| log_emission[idx] = self.config.alpha * np.log(prob + self.config.epsilon) | |
| elif self.emission_transform == EmissionTransform.TEMPERATURE: | |
| # Temperature scaling on logits | |
| if logits is None: | |
| # Fallback: invert softmax approximately | |
| logits = {k: np.log(v + self.config.epsilon) for k, v in neural_probs.items()} | |
| temp = self.config.temperature | |
| scaled_logits = np.full(self.n_intents, -np.inf) | |
| for intent, logit in logits.items(): | |
| if intent in self.config.intent_to_idx: | |
| idx = self.config.intent_to_idx[intent] | |
| scaled_logits[idx] = logit / temp | |
| log_emission = log_normalize(scaled_logits) | |
| elif self.emission_transform == EmissionTransform.RAW: | |
| # Pass-through: no sharpening | |
| for intent, prob in neural_probs.items(): | |
| if intent in self.config.intent_to_idx: | |
| idx = self.config.intent_to_idx[intent] | |
| log_emission[idx] = np.log(prob + self.config.epsilon) | |
| elif self.emission_transform == EmissionTransform.ISOTONIC: | |
| # Isotonic calibration | |
| if self.isotonic_calibrator is not None: | |
| calibrated = self.isotonic_calibrator(neural_probs) | |
| for intent, prob in calibrated.items(): | |
| if intent in self.config.intent_to_idx: | |
| idx = self.config.intent_to_idx[intent] | |
| log_emission[idx] = np.log(prob + self.config.epsilon) | |
| else: | |
| # Fallback to raw if no calibrator | |
| for intent, prob in neural_probs.items(): | |
| if intent in self.config.intent_to_idx: | |
| idx = self.config.intent_to_idx[intent] | |
| log_emission[idx] = np.log(prob + self.config.epsilon) | |
| return log_emission | |
| def update( | |
| self, | |
| neural_probs: Dict[str, float], | |
| logits: Optional[Dict[str, float]] = None, | |
| ) -> np.ndarray: | |
| """ | |
| Full belief update in log-space. | |
| If use_hmm=True: | |
| 1. log_b_pred = predict() (transition dynamics) | |
| 2. log_emission = compute_log_emission(neural_probs) | |
| 3. log_b = normalize(log_b_pred + log_emission) | |
| If use_hmm=False (ablation mode): | |
| Directly use neural posteriors as belief. | |
| Args: | |
| neural_probs: Dict mapping intent names to probabilities from neural model. | |
| logits: Optional dict of raw logits for temperature scaling. | |
| Returns: | |
| Updated belief state (in probability space). | |
| """ | |
| self._last_raw_probs = neural_probs.copy() | |
| if not self.use_hmm: | |
| # Ablation: bypass HMM, use raw neural posteriors | |
| for intent, prob in neural_probs.items(): | |
| if intent in self.config.intent_to_idx: | |
| idx = self.config.intent_to_idx[intent] | |
| self.log_belief[idx] = np.log(prob + self.config.epsilon) | |
| self.log_belief = log_normalize(self.log_belief) | |
| self._step_count += 1 | |
| return self.belief | |
| # Predict step | |
| log_b_pred = self.predict() | |
| # Compute log-emission likelihood | |
| log_emission = self.compute_log_emission(neural_probs, logits) | |
| # Update step: element-wise addition in log-space and normalize | |
| log_b_updated = log_b_pred + log_emission | |
| self.log_belief = log_normalize(log_b_updated) | |
| self._step_count += 1 | |
| return self.belief | |
| def update_with_log_emission(self, log_emission: np.ndarray) -> np.ndarray: | |
| """ | |
| Update belief with pre-computed log-emission vector. | |
| Args: | |
| log_emission: Pre-computed log-emission likelihood vector. | |
| Returns: | |
| Updated belief state (in probability space). | |
| """ | |
| if not self.use_hmm: | |
| self.log_belief = log_normalize(log_emission) | |
| self._step_count += 1 | |
| return self.belief | |
| log_b_pred = self.predict() | |
| log_b_updated = log_b_pred + log_emission | |
| self.log_belief = log_normalize(log_b_updated) | |
| self._step_count += 1 | |
| return self.belief | |
| def belief(self) -> np.ndarray: | |
| """Get current belief in probability space.""" | |
| return np.exp(self.log_belief) | |
| def get_belief_dict(self) -> Dict[str, float]: | |
| """Get current belief as dictionary.""" | |
| probs = self.belief | |
| return { | |
| self.config.idx_to_intent[i]: float(probs[i]) | |
| for i in range(self.n_intents) | |
| } | |
| def get_top_intent(self) -> Tuple[str, float]: | |
| """Get intent with highest belief probability.""" | |
| probs = self.belief | |
| idx = int(np.argmax(probs)) | |
| return self.config.idx_to_intent[idx], float(probs[idx]) | |
| def get_intent_prob(self, intent: str) -> float: | |
| """Get belief probability for specific intent.""" | |
| if intent in self.config.intent_to_idx: | |
| idx = self.config.intent_to_idx[intent] | |
| return float(np.exp(self.log_belief[idx])) | |
| return 0.0 | |
| def get_log_belief(self) -> np.ndarray: | |
| """Get current log-belief (for debugging/analysis).""" | |
| return self.log_belief.copy() | |
| def get_last_raw_probs(self) -> Optional[Dict[str, float]]: | |
| """Get last raw neural probabilities (before HMM filtering).""" | |
| return self._last_raw_probs | |
| def step_count(self) -> int: | |
| """Number of update steps performed.""" | |
| return self._step_count | |
| def is_normalized(self, tol: float = 1e-6) -> bool: | |
| """Check if belief state is properly normalized.""" | |
| return abs(self.belief.sum() - 1.0) < tol | |