Spaces:
Sleeping
Sleeping
| from __future__ import annotations | |
| from dataclasses import dataclass | |
| from pathlib import Path | |
| from typing import Iterable, Protocol, Sequence | |
| import numpy as np | |
| class BrainSignals: | |
| """Signals emitted by one agent-context brain for one transcript update. | |
| `readiness` is the controller-facing score for a speculative short reply: | |
| readiness = 1 / (1 + mean_token_entropy) | |
| A low-entropy draft means this agent has a confident next move. `p_end` is a | |
| turn-completion heuristic: in the live loop it should combine trailing | |
| silence, sentence-final punctuation, and/or high EOS probability. | |
| """ | |
| surprise: float | |
| hidden: np.ndarray | |
| readiness: float | |
| p_end: float | |
| def __post_init__(self) -> None: | |
| hidden = np.asarray(self.hidden, dtype=np.float32) | |
| if hidden.ndim != 1: | |
| raise ValueError("hidden must be a 1-D float32 vector") | |
| if not 0.0 <= float(self.readiness) <= 1.0: | |
| raise ValueError("readiness must be in [0, 1]") | |
| if not 0.0 <= float(self.p_end) <= 1.0: | |
| raise ValueError("p_end must be in [0, 1]") | |
| object.__setattr__(self, "hidden", hidden) | |
| object.__setattr__(self, "surprise", float(self.surprise)) | |
| object.__setattr__(self, "readiness", float(self.readiness)) | |
| object.__setattr__(self, "p_end", float(self.p_end)) | |
| class Brain(Protocol): | |
| """Interface the live instrumented LLM must satisfy for one agent context.""" | |
| def next_signals(self) -> BrainSignals: | |
| """Return signals for the newest incremental transcript update.""" | |
| class ReplayBrain: | |
| """Deterministic `Brain` backed by precomputed signal samples.""" | |
| def __init__(self, samples: Sequence[BrainSignals | dict[str, object]], *, name: str = "replay") -> None: | |
| self.name = name | |
| self._samples = [coerce_signals(sample) for sample in samples] | |
| self._index = 0 | |
| def __len__(self) -> int: | |
| return len(self._samples) | |
| def __iter__(self) -> Iterable[BrainSignals]: | |
| return iter(self._samples) | |
| def reset(self) -> None: | |
| self._index = 0 | |
| def next_signals(self) -> BrainSignals: | |
| if self._index >= len(self._samples): | |
| raise StopIteration(f"ReplayBrain {self.name!r} is exhausted") | |
| sample = self._samples[self._index] | |
| self._index += 1 | |
| return sample | |
| def from_npz( | |
| cls, | |
| path: str | Path, | |
| *, | |
| readiness: Sequence[float] | None = None, | |
| p_end: Sequence[float] | None = None, | |
| name: str | None = None, | |
| ) -> "ReplayBrain": | |
| dump = np.load(path, allow_pickle=True) | |
| surprises = np.asarray(dump["nll_series"], dtype=np.float32) | |
| hidden = np.asarray(dump["hidden_states"], dtype=np.float32) | |
| if hidden.ndim != 2: | |
| raise ValueError("hidden_states in npz must be a 2-D matrix") | |
| n_steps = int(surprises.shape[0]) | |
| readiness_values = ( | |
| np.asarray(readiness, dtype=np.float32) | |
| if readiness is not None | |
| else np.linspace(0.35, 0.75, n_steps, dtype=np.float32) | |
| ) | |
| p_end_values = np.asarray(p_end, dtype=np.float32) if p_end is not None else np.zeros(n_steps, dtype=np.float32) | |
| if n_steps: | |
| p_end_values[-1] = max(float(p_end_values[-1]), 0.95) | |
| samples = [ | |
| BrainSignals( | |
| surprise=float(surprises[index]), | |
| hidden=hidden[index], | |
| readiness=float(readiness_values[index]), | |
| p_end=float(p_end_values[index]), | |
| ) | |
| for index in range(n_steps) | |
| ] | |
| return cls(samples, name=name or Path(path).stem) | |
| def coerce_signals(sample: BrainSignals | dict[str, object]) -> BrainSignals: | |
| if isinstance(sample, BrainSignals): | |
| return sample | |
| return BrainSignals( | |
| surprise=float(sample["surprise"]), | |
| hidden=np.asarray(sample["hidden"], dtype=np.float32), | |
| readiness=float(sample["readiness"]), | |
| p_end=float(sample["p_end"]), | |
| ) | |