| | 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 |
| |
|
