| import numpy as np |
| from typing import Dict, List, TYPE_CHECKING |
|
|
| import torch |
| from sklearn.cluster import KMeans |
|
|
| if TYPE_CHECKING: |
| from .model import BitTransformerLM |
|
|
|
|
| class TelemetrySynthesizer: |
| """Analyze telemetry batches and cluster activation patterns.""" |
|
|
| def __init__(self, n_clusters: int = 2) -> None: |
| self.n_clusters = n_clusters |
|
|
| def _summary(self, telemetry: Dict[str, List[torch.Tensor]]) -> np.ndarray: |
| """Compute activation/attention summaries for a single telemetry dict.""" |
| acts = telemetry["activations"] |
| attn = telemetry["attention_maps"] |
| summaries = [] |
| for a, m in zip(acts, attn): |
| mean = a.mean().item() |
| var = a.var(unbiased=False).item() |
| prob = m.softmax(-1) |
| entropy = -(prob * prob.clamp_min(1e-9).log()).sum(-1).mean().item() |
| summaries.append([mean, var, entropy]) |
| return np.array(summaries).ravel() |
|
|
| def synthesize( |
| self, telemetries: List[Dict[str, List[torch.Tensor]]], bit_seqs: torch.Tensor |
| ) -> Dict[str, List]: |
| """Cluster telemetry summaries and return cluster info.""" |
| data = np.stack([self._summary(t) for t in telemetries]) |
| km = KMeans(n_clusters=self.n_clusters, n_init=1) |
| labels = km.fit_predict(data) |
| representatives: List[List[int]] = [] |
| for c in range(self.n_clusters): |
| idx = np.where(labels == c)[0] |
| if len(idx) > 0: |
| representatives.append(bit_seqs[idx[0]].tolist()) |
| else: |
| representatives.append([]) |
| return {"cluster_assignments": labels.tolist(), "representatives": representatives} |
|
|
| def cluster_sequences( |
| self, model: "BitTransformerLM", bit_seqs: torch.Tensor |
| ) -> List[List[int]]: |
| """Run the model to gather telemetry and return representative sequences. |
| |
| Parameters |
| ---------- |
| model: BitTransformerLM |
| Model used to compute telemetry for each sequence. |
| bit_seqs: torch.Tensor |
| Tensor containing one bit sequence per row. |
| |
| Returns |
| ------- |
| list[list[int]] |
| Representative sequences chosen from KMeans clusters. |
| """ |
| telemetries: List[Dict[str, List[torch.Tensor]]] = [] |
| with torch.no_grad(): |
| for seq in bit_seqs: |
| _, tele = model(seq.unsqueeze(0)) |
| telemetries.append(tele) |
| info = self.synthesize(telemetries, bit_seqs) |
| return info["representatives"] |
|
|
|
|
| def detect_metric_drift( |
| metrics_log: Dict[str, List[float]], |
| window: int = 10, |
| threshold: float = 0.2, |
| ) -> Dict[str, bool]: |
| """Detect metric drift between consecutive windows. |
| |
| Args: |
| metrics_log: History of scalar metrics keyed by name. |
| window: Number of recent steps to compare. |
| threshold: Absolute difference required to flag drift. |
| |
| Returns: |
| Dictionary mapping metric keys to a boolean drift indicator. |
| """ |
| drift = {} |
| for key, values in metrics_log.items(): |
| if len(values) < window * 2: |
| drift[key] = False |
| continue |
| recent = np.mean(values[-window:]) |
| prev = np.mean(values[-2 * window : -window]) |
| drift[key] = abs(recent - prev) > threshold |
| return drift |
|
|
