drowse.ai 5.3.1__py3-none-any.whl

drowse/core/surface_topology.py

@@ -0,0 +1,128 @@
+"""Synthetic point-cloud surface diagnostics, separate from chart recovery.
+
+Two coefficient fields distinguish orientable surfaces from 2-torsion.
+Homology signatures identify a surface only *conditional on* a connected,
+closed two-manifold. Finite point samples do not establish that hypothesis;
+these results are evidence, never a certificate or an automatic steering chart.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+from numpy.typing import ArrayLike
+
+
+@dataclass(frozen=True)
+class SurfaceEvidence:
+    candidate: str | None
+    reason: str
+    betti_mod2: tuple[int, int, int] | None = None
+    betti_mod3: tuple[int, int, int] | None = None
+    scale_interval: tuple[float, float] | None = None
+    relative_persistence: float = 0.0
+    sample_count: int = 0
+
+
+def _surface_name(b2: tuple[int, ...], b3: tuple[int, ...]) -> str | None:
+    if b2[0] != 1 or b3[0] != 1 or b2[2] != 1:
+        return None
+    if b3[2] == 1 and b2[1] == b3[1] and b2[1] % 2 == 0:
+        genus = b2[1] // 2
+        return {0: 'sphere', 1: 'torus'}.get(genus, f'orientable-genus-{genus}')
+    if b3[2] == 0 and b2[1] == b3[1] + 1:
+        genus = b2[1]
+        return {1: 'projective-plane', 2: 'klein-bottle'}.get(genus, f'nonorientable-genus-{genus}')
+    return None
+
+
+def detect_surface(points: ArrayLike, *, min_relative_persistence: float = 0.15) -> SurfaceEvidence:
+    """Measure closed-surface evidence in explicit Euclidean R^n sample data.
+
+    Input is 32..384 distinct finite points; larger inputs are rejected, not
+    silently subsampled. Ripser computes H0/H1/H2 over F2 and F3. A dominant
+    finite H2 bar and a simultaneous stable signature in both fields are
+    required. Its birth-normalized persistence must dominate the next H2 bar
+    by at least 3:1;
+    competing signatures and short scale intervals cause abstention.
+
+    The persistence threshold is a fraction of the feature's birth radius,
+    not the cloud diameter or a probability/confidence score. This avoids
+    discarding a resolved small hole solely because another axis is long.
+    Isometries and uniform rescaling leave
+    the decision invariant. Noise, sparse coverage, thin necks, boundaries,
+    and self-intersections can invalidate a closed-surface interpretation.
+    Activation-space callers must whiten explicitly before calling this API.
+    """
+    if np.iscomplexobj(points):
+        raise ValueError('Surface points must be real-valued')
+    data = np.asarray(points, dtype=np.float64)
+    if data.ndim != 2 or not 32 <= len(data) <= 384 or data.shape[1] < 2:
+        raise ValueError('Surface diagnostics need 32..384 points with at least two coordinates')
+    if not np.isfinite(data).all():
+        raise ValueError('Surface points must be finite')
+    if not np.isfinite(min_relative_persistence) or not 0 < min_relative_persistence < 1:
+        raise ValueError('min_relative_persistence must be between zero and one')
+    from scipy.spatial.distance import pdist, squareform
+
+    # Scale before squaring distances, avoiding overflow for large unit changes.
+    magnitude = np.max(np.abs(data))
+    if magnitude == 0:
+        raise ValueError('Surface points must be distinct')
+    data = data / magnitude
+    data -= data.mean(axis=0)
+    spread = np.max(np.abs(data))
+    if spread == 0:
+        raise ValueError('Surface points must be distinct')
+    data /= spread
+    pairwise = pdist(data)
+    if pairwise.min() <= 1e-12:
+        raise ValueError('Surface points contain duplicate or numerically indistinguishable samples')
+    scale = np.median(pairwise)
+    distances = squareform(pairwise / scale)
+    try:
+        from ripser import ripser
+    except ImportError as exc:
+        raise ImportError('Surface diagnostics require drowse[topology]') from exc
+
+    diagrams = [ripser(distances, distance_matrix=True, maxdim=2, coeff=2)['dgms']]
+    h2 = diagrams[0][2]
+    if not len(h2):
+        return SurfaceEvidence(None, 'No persistent two-dimensional class', sample_count=len(data))
+    lifetimes = h2[:, 1] - h2[:, 0]
+    relative = lifetimes / h2[:, 0]
+    order = np.argsort(relative)
+    principal = h2[order[-1]]
+    lifetime = float(lifetimes[order[-1]])
+    if not np.isfinite(principal).all() or lifetime < min_relative_persistence * principal[0]:
+        return SurfaceEvidence(None, 'Two-dimensional persistence is too short or unbounded', sample_count=len(data))
+    if len(order) > 1 and relative[order[-1]] < 3 * relative[order[-2]]:
+        return SurfaceEvidence(None, 'Competing two-dimensional classes; surface unresolved', sample_count=len(data))
+
+    diagrams.append(ripser(distances, distance_matrix=True, maxdim=2, coeff=3)['dgms'])
+
+    # All critical values are included; no fixed grid can skip a short-lived
+    # contradictory class. Trim H2 endpoints so its birth/death do not vote.
+    lo, hi = principal[0] + 0.05 * lifetime, principal[1] - 0.05 * lifetime
+    events = sorted({float(lo), float(hi)} | {
+        float(value) for field in diagrams for barcode in field
+        for value in barcode.ravel() if lo < value < hi
+    })
+    intervals = []
+    for left, right in zip(events, events[1:]):
+        middle = (left + right) / 2
+        signatures = [tuple(int(np.sum((d[:, 0] <= middle) & (middle < d[:, 1]))) for d in field) for field in diagrams]
+        key = (_surface_name(*signatures), *signatures)
+        if intervals and intervals[-1][0] == key:
+            intervals[-1] = (key, intervals[-1][1], right)
+        else:
+            intervals.append((key, left, right))
+    viable = [(key, left, right) for key, left, right in intervals
+              if key[0] is not None and right - left >= min_relative_persistence * left]
+    if not viable or len({key[0] for key, _, _ in viable}) != 1:
+        return SurfaceEvidence(None, 'No unambiguous simultaneous surface signature', relative_persistence=lifetime / float(principal[0]), sample_count=len(data))
+    key, left, right = max(viable, key=lambda interval: interval[2] - interval[1])
+    return SurfaceEvidence(
+        key[0], 'Homology consistent with a connected closed surface; no recovered coordinate chart',
+        key[1], key[2], (left, right), (right - left) / left, len(data),
+    )

drowse/core/token_callback.py

@@ -0,0 +1,52 @@
+"""Typed token-callback contract and consumer capability declaration."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Protocol
+
+from drowse.core.results import TokenAlt
+
+
+class TokenCallback(Protocol):
+    def __call__(self, text: str, is_thinking: bool, token_id: int | None,
+                 logprob: float | None, top_alts: list[TokenAlt] | None,
+                 perplexity: float | None = None, /) -> None: ...
+
+
+class StepTokenCallback(Protocol):
+    """The decode loop's INTERNAL tap contract — ``generate_steered``'s
+    ``on_token``.  The public :class:`TokenCallback` (what users hand to the
+    session) stays six-argument; the trailing ``step_id`` is the loop-owned
+    forward index (the same value that forward's ``step_callback`` /
+    ``score_callback`` received), which the session tap uses to pair the
+    instrument runs' step-keyed memos.  ``DrowseSession._token_tap`` absorbs
+    it and invokes user callbacks with the public six arguments."""
+
+    def __call__(self, text: str, is_thinking: bool, token_id: int | None,
+                 logprob: float | None, top_alts: list[TokenAlt] | None,
+                 perplexity: float | None, step_id: int, /) -> None: ...
+
+
+@dataclass(frozen=True)
+class TokenConsumerOptions:
+    live_scores: bool = False
+    per_layer_scores: bool = False
+    lens_readout: bool = False
+    sae_readout: bool = False
+    perplexity: bool = False
+
+
+@dataclass(frozen=True)
+class TokenConsumer:
+    callback: TokenCallback
+    options: TokenConsumerOptions = TokenConsumerOptions()
+
+    def __call__(self, text: str, is_thinking: bool, token_id: int | None,
+                 logprob: float | None, top_alts: list[TokenAlt] | None,
+                 perplexity: float | None = None) -> None:
+        self.callback(text, is_thinking, token_id, logprob, top_alts, perplexity)
+
+
+def consumer_options(callback: object | None) -> TokenConsumerOptions:
+    return callback.options if isinstance(callback, TokenConsumer) else TokenConsumerOptions()

drowse/core/token_payloads.py

@@ -0,0 +1,244 @@
+"""Token-level probe payload shaping for generation callbacks."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Mapping, cast
+
+import torch
+
+from drowse.core.instruments.types import (
+    InstrumentFamily,
+    Reading,
+    ScalarReading,
+    as_probe_reading,
+    reading_axis0,
+    reading_per_layer_axis0,
+)
+from drowse.core.measurements import build_measurements
+from drowse.core.results import ProbeReading
+
+#: Per-family reading slot names, keyed by the ``family`` argument to
+#: :meth:`TokenProbePayload.merge_readings`.  One slot per member of
+#: ``InstrumentFamily`` — ``tests/test_instruments_types.py`` pins the two
+#: sets equal against ``session.instruments``, so a fourth read family
+#: cannot land a slot-less payload.
+_FAMILY_SLOTS: dict[InstrumentFamily, str] = {
+    "geometry": "geometry_readings",
+    "lens": "lens_readings",
+    "sae": "sae_readings",
+}
+
+
+@dataclass(slots=True)
+class TokenProbePayload:
+    """Probe-derived payloads for one generated token.
+
+    Readings are kept in per-family slots (``geometry_readings`` /
+    ``lens_readings`` / ``sae_readings``) so the 5.x measurement envelope can
+    split them into ``instruments.geometry`` / ``.lens`` / ``.sae`` while the
+    flat ``scores`` / ``per_layer_scores`` cross-family views stay derived from
+    all three.
+    """
+
+    scores: dict[str, float] | None = None
+    per_layer_scores: dict[str, dict[str, float]] | None = None
+    geometry_readings: dict[str, ProbeReading] | None = None
+    lens_readings: dict[str, ScalarReading] | None = None
+    sae_readings: dict[str, ScalarReading] | None = None
+
+    @property
+    def has_readings(self) -> bool:
+        """Whether any family slot is populated — the cheap guard.
+
+        Distinct from :attr:`all_readings` on purpose: the per-token payload
+        gate must not pay for the cross-family ``ProbeReading`` projection
+        just to answer "is there anything here".
+        """
+        return bool(
+            self.geometry_readings or self.lens_readings or self.sae_readings
+        )
+
+    @property
+    def all_readings(self) -> dict[str, ProbeReading]:
+        """The three family slots merged into one ``ProbeReading`` dict.
+
+        A COMPATIBILITY view, and the only one on this path: the library
+        ``TokenEvent.probe_readings`` field keys probes across families by a
+        single reading type, so the single-axis families' native
+        ``ScalarReading``s are projected here.  The native WS path never
+        touches it — ``to_token_payload`` serializes each family's own
+        shape — so the projection does not run per token unless a caller
+        asked for the cross-family dict.
+        """
+        merged: dict[str, ProbeReading] = {}
+        if self.geometry_readings:
+            merged.update(self.geometry_readings)
+        for slot in (self.lens_readings, self.sae_readings):
+            if slot:
+                merged.update(
+                    {n: as_probe_reading(r) for n, r in slot.items()}
+                )
+        return merged
+
+    def to_token_payload(
+        self,
+        *,
+        lens: Any = None,
+        lens_aggregate: Any = None,
+        lens_token_ids: Any = None,
+        lens_source: str | None = None,
+        sae: Any = None,
+        sae_source: str | None = None,
+        sae_layer: int | None = None,
+        steering: str | None = None,
+    ) -> dict[str, Any]:
+        """Shape this token's payload as ``{"measurements": <envelope>}``.
+
+        The measurement envelope is the whole wire record for a token — every
+        cross-family scalar view is derived inside it, so there are no
+        top-level aliases beside it.  ``lens`` is the native per-layer readout
+        (``lens_readout`` in the envelope) and ``sae`` the native feature
+        discovery (``sae_features``); the attached per-family readings come
+        from the payload's own slots.
+        """
+        measurements = build_measurements(
+            scope="token",
+            geometry_readings=self.geometry_readings,
+            lens_readings=self.lens_readings,
+            sae_readings=self.sae_readings,
+            per_layer_scores=self.per_layer_scores,
+            lens_readout=lens,
+            lens_aggregate=lens_aggregate,
+            lens_token_ids=lens_token_ids,
+            lens_source=lens_source,
+            sae_features=sae,
+            sae_source=sae_source,
+            sae_layer=sae_layer,
+            steering=steering,
+        )
+        return {"measurements": measurements}
+
+    def merge_readings(
+        self,
+        extra: dict[str, Any],
+        *,
+        family: str,
+        per_layer: bool = False,
+    ) -> None:
+        """Merge additional per-probe readings into a named family's slot.
+
+        ``family`` is one of ``"geometry"`` / ``"lens"`` / ``"sae"`` — J-lens
+        token probes and SAE feature probes score on their own paths rather
+        than through the Monitor, so their readings land in their own slot.
+        The flat ``scores`` (and, when ``per_layer``, ``per_layer_scores``)
+        cross-family views are updated the same way regardless of family.
+        """
+        if not extra:
+            return
+        try:
+            slot_name = _FAMILY_SLOTS[cast(InstrumentFamily, family)]
+        except KeyError:
+            raise ValueError(f"unknown reading family {family!r}") from None
+        self.scores = {**(self.scores or {}), **_axis0_scores(extra)}
+        current: dict[str, ProbeReading] | None = getattr(self, slot_name)
+        setattr(self, slot_name, {**(current or {}), **extra})
+        if per_layer:
+            merged = dict(self.per_layer_scores or {})
+            for layer, row in (_per_layer_axis0(extra) or {}).items():
+                merged[layer] = {**merged.get(layer, {}), **row}
+            self.per_layer_scores = merged or None
+
+
+def _axis0_scores(readings: Mapping[str, Reading]) -> dict[str, float]:
+    return {
+        name: reading_axis0(reading) for name, reading in readings.items()
+    }
+
+
+def _per_layer_axis0(
+    readings: Mapping[str, Reading],
+) -> dict[str, dict[str, float]] | None:
+    by_layer: dict[str, dict[str, float]] = {}
+    for probe_name, reading in readings.items():
+        for layer, value in reading_per_layer_axis0(reading).items():
+            by_layer.setdefault(str(layer), {})[probe_name] = round(value, 6)
+    return by_layer or None
+
+
+def build_token_probe_payload(
+    *,
+    monitor: Any,
+    capture: Any,
+    capture_state: Any,
+    incremental_readings: list[dict[str, ProbeReading]],
+    needs_scores: bool,
+    persists_layer_scores: bool,
+    assistant_node_id: str | None,
+    geometry_run: Any = None,
+    step_id: int = -1,
+) -> TokenProbePayload:
+    """Score and shape probe payloads for one generated token.
+
+    This is the token-probe slice of ``DrowseSession._token_tap`` extracted into
+    a typed helper. It performs at most one monitor geometry pass: either reuse
+    the most recent incremental reading, or score the latest captured hidden
+    states through the geometry run's step-keyed ``observe`` (FULL-retention
+    mode — a gate callback that already full-roster-scored this forward primed
+    the memo, so the tap hits it instead of rescoring; ``monitor`` remains the
+    fallback for callers without a run).
+    The monitor readings land in the ``geometry`` family slot.
+    """
+    if not needs_scores:
+        return TokenProbePayload()
+
+    has_incremental_reading = bool(
+        (capture_state.incremental or capture_state.lean)
+        and incremental_readings
+    )
+    latest_hidden_for_token: dict[int, torch.Tensor] | None = None
+    if not has_incremental_reading:
+        latest_hidden_for_token = {
+            layer_idx: bucket[-1]
+            for layer_idx, bucket in capture.per_layer_buckets().items()
+            if bucket
+        }
+
+    readings: dict[str, ProbeReading] | None = None
+    if has_incremental_reading:
+        if capture_state.lean or geometry_run is None:
+            # LEAN rows are deliberately partial (``coords_only``) and are
+            # never primed into the observe memo — read the sink's row
+            # directly (routing through ``observe`` would either rescore
+            # the full roster per token or serve a partial as full).
+            readings = incremental_readings[-1] or None
+        else:
+            # FULL-incremental: the sink primed the run's memo for this
+            # forward, so ``observe`` is a hit returning the sink's exact
+            # row.  A miss falls back to the sink's appended row (aligned
+            # 1:1 with forwards — the miss cases are an empty-capture row,
+            # which never primes, and a narrow caller's idle run).
+            readings = (
+                geometry_run.observe(step_id, {})
+                or incremental_readings[-1]
+                or None
+            )
+    elif latest_hidden_for_token:
+        readings = (
+            geometry_run.observe(step_id, latest_hidden_for_token)
+            if geometry_run is not None
+            else monitor.score_single_token(latest_hidden_for_token)
+        ) or None
+
+    if not readings:
+        return TokenProbePayload()
+
+    return TokenProbePayload(
+        scores=_axis0_scores(readings),
+        per_layer_scores=(
+            _per_layer_axis0(readings)
+            if assistant_node_id is not None and persists_layer_scores
+            else None
+        ),
+        geometry_readings=readings,
+    )