Spaces:
Sleeping
Sleeping
File size: 4,106 Bytes
bcbb8c5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 | from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable, Protocol, Sequence
import numpy as np
@dataclass(frozen=True)
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
@classmethod
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"]),
)
|