Update ai_risk_engine.py

ai_risk_engine.py

@@ -1,12 +1,14 @@
 """
 Bayesian risk engine with hyperpriors (hierarchical Beta‑binomial).
+Uses Pyro for variational inference.
 """
 import logging
 import pyro
 import pyro.distributions as dist
 from pyro.infer import SVI, Trace_ELBO, Predictive
 import torch
-from typing import Dict, Optional
+import numpy as np
+from typing import Dict, Optional, List, Tuple
 
 logger = logging.getLogger(__name__)
 
@@ -19,14 +21,16 @@ class AIRiskEngine:
     def __init__(self, num_categories: int = 10):
         self.num_categories = num_categories
         self.category_names = ["chat", "code", "summary", "image", "audio", "iot", "switch", "server", "service", "unknown"]
+        self._history: List[Tuple[int, float]] = []  # (category_idx, success)
         self._init_model()
-        self._history = []  # store (category, success) for later updates
 
     def _init_model(self):
-        # Hyperpriors
+        # Hyperpriors (Gamma for alpha, beta)
         self.alpha0 = pyro.param("alpha0", torch.tensor(2.0), constraint=dist.constraints.positive)
         self.beta0 = pyro.param("beta0", torch.tensor(2.0), constraint=dist.constraints.positive)
-        # We'll learn these via SVI when update is called
+        # Category‑specific parameters
+        self.p_alpha = pyro.param("p_alpha", torch.ones(self.num_categories) * 2.0, constraint=dist.constraints.positive)
+        self.p_beta = pyro.param("p_beta", torch.ones(self.num_categories) * 2.0, constraint=dist.constraints.positive)
 
     def model(self, observations=None):
         # Global hyperprior (concentration parameters)
@@ -38,14 +42,13 @@ class AIRiskEngine:
             p = pyro.sample("p", dist.Beta(alpha0, beta0))
 
         if observations is not None:
-            # Observations: list of (category_idx, success)
             cat_idx = torch.tensor([obs[0] for obs in observations])
             successes = torch.tensor([obs[1] for obs in observations])
             with pyro.plate("data", len(observations)):
                 pyro.sample("obs", dist.Bernoulli(p[cat_idx]), obs=successes)
 
     def guide(self, observations=None):
-        # Variational parameters
+        # Variational parameters for hyperpriors
         alpha0_q = pyro.param("alpha0_q", torch.tensor(2.0), constraint=dist.constraints.positive)
         beta0_q = pyro.param("beta0_q", torch.tensor(2.0), constraint=dist.constraints.positive)
         pyro.sample("alpha0", dist.Gamma(alpha0_q, 1.0))
@@ -60,6 +63,7 @@ class AIRiskEngine:
         """Store observation and optionally trigger a learning step."""
         cat_idx = self.category_names.index(category) if category in self.category_names else -1
         if cat_idx == -1:
+            logger.warning(f"Unknown category: {category}")
             return
         self._history.append((cat_idx, 1.0 if success else 0.0))
         # Run a few steps of SVI
@@ -71,10 +75,10 @@ class AIRiskEngine:
             return
         optimizer = pyro.optim.Adam({"lr": 0.01})
         svi = SVI(self.model, self.guide, optimizer, loss=Trace_ELBO())
-        for _ in range(steps):
+        for step in range(steps):
             loss = svi.step(self._history)
-            if steps % 10 == 0:
-                logger.debug(f"SVI loss: {loss}")
+            if step % 10 == 0:
+                logger.debug(f"SVI step {step}, loss: {loss}")
 
     def risk_score(self, category: str) -> Dict[str, float]:
         """Return posterior predictive risk metrics for a category."""