drowse.ai 5.3.1__py3-none-any.whl

drowse/cli/runners/__init__.py

@@ -0,0 +1,155 @@
+"""Runner functions for drowse CLI subcommands.
+
+This package splits the former single ``runners.py`` module into one file per
+verb group (``serve`` / ``manifold`` / ``pack`` / ``template`` / ``config`` /
+``lens`` / ``sae`` / ``experiment``) plus a ``shared`` module for the helpers
+used across more than one group.  Every runner and helper is re-exported here,
+so the historical import surface ``drowse.cli.runners.<name>`` (and the
+monkeypatch seams the test-suite relies on) keeps working unchanged.
+
+Session/config helpers that the test-suite swaps at the package level
+(``_make_session`` / ``_print_startup`` / ``_print_model_info`` /
+``_load_effective_config`` / the ``_fold_*`` folders /
+``_load_or_fit_transfer_alignment``) are invoked by the submodule runners
+through this package object (``import drowse.cli.runners as _pkg``), so a
+``monkeypatch.setattr("drowse.cli.runners.<name>", ...)`` still takes effect.
+"""
+
+# This module is intentionally a compatibility barrel: every imported private
+# name is re-exported for historical callers and monkeypatch seams.
+# ruff: noqa: F401
+
+from __future__ import annotations
+
+from drowse.cli.runners.shared import (
+    _R,
+    _iter_manifold_folders,
+    _load_effective_config,
+    _make_session,
+    _print_model_info,
+    _print_startup,
+    _print_verb_menu,
+    _progress_printer,
+    _resolve_manifold_folder,
+    _resolve_manifold_ns_name,
+    _resolve_probes,
+    _drowse_error_exit,
+    _setup_steering_vectors,
+    _split_manifold_ns_name,
+    _warmup_session,
+)
+
+# The transfer-alignment single-flight orchestration now lives beside the
+# alignment primitives it wraps; re-exported under its historical private name
+# so ``drowse.cli.runners._load_or_fit_transfer_alignment`` still resolves and
+# stays monkeypatchable at the package level.
+from drowse.io.alignment import (
+    load_or_fit_transfer_alignment as _load_or_fit_transfer_alignment,
+)
+
+from drowse.cli.runners.serve import (
+    _best_serve_sae_release,
+    _enable_serve_live_lens_if_compatible,
+    _enable_serve_live_sae_if_available,
+    _run_serve,
+)
+from drowse.cli.runners.manifold import (
+    _MANIFOLD_RUNNERS,
+    _VARIANT_SUFFIX_RE,
+    _fold_all_fitted_manifolds,
+    _fold_manifold_to_profile_with_identity,
+    _print_why_histogram,
+    _require_model,
+    _run_manifold,
+    _run_manifold_bake,
+    _run_manifold_compare,
+    _run_manifold_extract,
+    _run_manifold_fit,
+    _run_manifold_from_template,
+    _run_manifold_generate,
+    _run_manifold_merge,
+    _run_manifold_transfer,
+    _run_manifold_why,
+    _split_variant_suffix,
+)
+from drowse.cli.runners.pack import (
+    _PACK_RUNNERS,
+    _run_pack,
+    _run_pack_clear,
+    _run_pack_export,
+    _run_pack_install,
+    _run_pack_ls,
+    _run_pack_push,
+    _run_pack_refresh,
+    _run_pack_rm,
+    _run_pack_search,
+    _run_pack_show,
+)
+from drowse.cli.runners.template import (
+    _TEMPLATE_RUNNERS,
+    _normalize_context_entry,
+    _run_template,
+    _run_template_create,
+    _run_template_ls,
+    _run_template_rm,
+    _run_template_score,
+    _run_template_show,
+)
+from drowse.cli.runners.config import (
+    _run_config,
+    _run_config_show,
+    _run_config_validate,
+)
+from drowse.cli.runners.lens import (
+    _LENS_DOC_CHARS,
+    _LENS_RUNNERS,
+    _load_lens_corpus,
+    _lens_fit_source_preflight_matches,
+    _parse_layer_list,
+    _run_lens,
+    _run_lens_decompose,
+    _run_lens_fetch,
+    _run_lens_fit,
+    _run_lens_ls,
+    _run_lens_rm,
+    _run_lens_show,
+    _run_lens_top,
+    _run_lens_use,
+    _try_lens_fit_noop_preflight,
+    _try_lens_fit_noop_preflight_locked,
+)
+from drowse.cli.runners.sae import (
+    _SAE_RUNNERS,
+    _load_sae_training_corpus,
+    _resolve_active_sae_source,
+    _run_sae,
+    _run_sae_fetch,
+    _run_sae_ls,
+    _run_sae_rm,
+    _run_sae_show,
+    _run_sae_train,
+    _run_sae_use,
+)
+from drowse.cli.runners.experiment import (
+    _EXPERIMENT_RUNNERS,
+    _TRANSCRIPT_RUNNERS,
+    _TRANSCRIPT_VERBS,
+    _parse_grid_terms,
+    _run_experiment,
+    _run_experiment_fan,
+    _run_experiment_naturalness,
+    _run_experiment_transcript,
+    _run_transcript_run,
+)
+
+
+_COMMAND_RUNNERS = {
+    "serve":      _run_serve,
+    "manifold":   _run_manifold,
+    "pack":       _run_pack,
+    "config":     _run_config,
+    "experiment": _run_experiment,
+    "template":   _run_template,
+    "lens":       _run_lens,
+    "sae":        _run_sae,
+}

drowse/cli/runners/config.py

@@ -0,0 +1,85 @@
+"""``drowse config <verb>`` runners (show / validate)."""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+from drowse.cli.runners.shared import _drowse_error_exit
+
+
+@_drowse_error_exit
+def _run_config(args: argparse.Namespace) -> None:
+    cmd = getattr(args, "config_cmd", None)
+    if cmd == "show":
+        _run_config_show(args)
+    elif cmd == "validate":
+        _run_config_validate(args)
+    else:
+        print("usage: drowse config {show,validate}")
+        print()
+        print("  show      Print the effective merged config")
+        print("  validate  Validate a config file (exit 0 valid, 2 invalid)")
+        sys.exit(0)
+
+
+def _run_config_show(args: argparse.Namespace) -> None:
+    from drowse import __version__
+    from drowse.cli.config_file import ConfigFile, apply_flag_overrides
+    extras = [Path(p) for p in (args.config or [])]
+    composed = ConfigFile.effective(extras, include_default=not args.no_default)
+    if args.model is not None:
+        composed = apply_flag_overrides(composed, model=args.model)
+    header = f"# effective merged config for drowse {__version__}"
+    sys.stdout.write(composed.to_yaml(header=header))
+
+
+def _run_config_validate(args: argparse.Namespace) -> None:
+    from drowse.cli.config_file import ConfigFile, ConfigFileError
+    from drowse.core.steering_expr import referenced_selectors
+    p = Path(args.file)
+    if not p.exists():
+        print(f"config validate: {p}: file not found", file=sys.stderr)
+        sys.exit(2)
+    try:
+        cfg = ConfigFile.load(p)
+        if cfg.vectors is None:
+            print(f"{p}: ok")
+            return
+        # Dry-run: don't install, just check resolvability.
+        from drowse.io.selectors import all_concepts
+        installed = {(c.namespace, c.name) for c in all_concepts()}
+        installed_names = {c.name for c in all_concepts()}
+        missing: list[str] = []
+        for ns, concept, _variant in referenced_selectors(cfg.vectors):
+            if ns is None:
+                if concept in installed_names:
+                    continue
+                # Bare pole of an installed bipolar resolves fine too.
+                slug = concept.split(".")[0] if "." in concept else concept
+                if any(
+                    slug in c.name.split(".")
+                    for c in all_concepts()
+                    if "." in c.name
+                ):
+                    continue
+                missing.append(concept)
+                continue
+            if ns == "default" or (ns, concept) in installed:
+                continue
+            if ns == "local":
+                missing.append(f"{ns}/{concept}")
+                continue
+            # HF namespace — we assume install would succeed; don't probe.
+        if missing:
+            raise ConfigFileError(
+                f"unresolvable vectors (not installed and no namespace to install from): {missing}"
+            )
+    except ConfigFileError as e:
+        print(f"config validate: {p}: {e}", file=sys.stderr)
+        sys.exit(2)
+    except Exception as e:
+        print(f"config validate: {p}: {type(e).__name__}: {e}", file=sys.stderr)
+        sys.exit(2)
+    print(f"{p}: ok")

drowse/cli/runners/experiment.py

@@ -0,0 +1,344 @@
+"""``drowse experiment <verb>`` runners (fan / transcript / naturalness)."""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+from typing import Any
+
+import drowse.cli.runners as _pkg
+from drowse.cli.parsers import _EXPERIMENT_VERBS
+from drowse.cli.runners.shared import _print_verb_menu, _drowse_error_exit
+
+#: The one nested verb under ``experiment transcript`` — its own menu table.
+_TRANSCRIPT_VERBS = [
+    ("run", "Replay a transcript on the current session"),
+]
+
+
+@_drowse_error_exit
+def _run_experiment(args: argparse.Namespace) -> None:
+    """Dispatch ``drowse experiment <verb>``."""
+    cmd = getattr(args, "experiment_cmd", None)
+    if cmd is None:
+        _print_verb_menu("experiment", _EXPERIMENT_VERBS)
+        sys.exit(0)
+    # Registered-subparser invariant: argparse rejects an unknown verb.
+    _EXPERIMENT_RUNNERS[cmd](args)
+
+
+def _run_experiment_naturalness(args: argparse.Namespace) -> None:
+    """Score a steered generation's behavior-manifold naturalness."""
+    import json as _json
+
+    import torch
+
+    from drowse.core.manifold import domain_from_spec
+    from drowse.core.naturalness import (
+        compute_node_behavior_centroid,
+        compute_trajectory_distributions,
+        fit_behavior_manifold,
+        trajectory_naturalness,
+    )
+    from drowse.core.profile import Profile
+    from drowse.core.sampling import SamplingConfig
+    from drowse.core.steering_expr import ManifoldTerm, parse_expr
+    from drowse.io.manifolds import ManifoldFolder
+
+    _pkg._load_effective_config(args, default_max_tokens=128)
+    mfolder = Path(args.manifold)
+    if not (mfolder / "manifold.json").exists():
+        print(
+            f"experiment naturalness: no manifold.json in {mfolder}",
+            file=sys.stderr,
+        )
+        sys.exit(2)
+    # Naturalness consumes only authoring geometry + node corpus; fitted
+    # payload integrity is irrelevant to this independent behavior-space fit.
+    mf = ManifoldFolder.load(mfolder, verify_manifest=False)
+    node_groups = mf.node_groups()
+    domain = domain_from_spec(mf.domain)
+    node_coords = torch.tensor(mf.node_coords, dtype=torch.float32)
+    node_params = domain.embed(node_coords)
+
+    _pkg._print_startup(args)
+    session = _pkg._make_session(args)
+    _pkg._print_model_info(session)
+
+    # 1. Fit the behavior manifold from the node corpus — each node's
+    #    mean next-token distribution, in Hellinger space.
+    print(f"fitting behavior manifold ({len(node_groups)} nodes)...")
+    centroids = [
+        compute_node_behavior_centroid(
+            session.model, session.tokenizer, session.device, statements,
+        )
+        for _label, statements in node_groups
+    ]
+    behavior = fit_behavior_manifold(torch.stack(centroids), node_params)
+
+    sampling = SamplingConfig(max_tokens=args.max_tokens, seed=0)
+
+    def _score(steer: str | None) -> tuple[str, float, float]:
+        result = session.generate(
+            args.prompt, steering=steer, sampling=sampling,
+        ).first
+        text = result.text
+        traj = compute_trajectory_distributions(
+            session.model, session.tokenizer, session.device, text,
+        )
+        per_step = trajectory_naturalness(traj, behavior, domain, node_coords)
+        return text, float(per_step.mean()), float(per_step.max())
+
+    rows: list[dict[str, Any]] = []
+    _text, mean_d, max_d = _score(args.steer)
+    rows.append({
+        "label": "manifold", "steering": args.steer,
+        "mean_bhattacharyya": mean_d, "max_bhattacharyya": max_d,
+    })
+
+    if args.compare_linear:
+        steering = parse_expr(args.steer)
+        mterms = [
+            v for v in steering.alphas.values()
+            if isinstance(v, ManifoldTerm)
+        ]
+        if len(mterms) != 1 or len(steering.alphas) != 1:
+            print(
+                "experiment naturalness: --compare-linear requires the "
+                "steer expression to be a single manifold term",
+                file=sys.stderr,
+            )
+            sys.exit(2)
+        mt = mterms[0]
+        session.ensure_manifold_loaded(mt.manifold)
+        act_manifold = session.manifolds[mt.manifold]
+        # Resolve label-form positions to coords up front — every
+        # downstream call here wants a coord tuple, and the chord
+        # baseline is per-coord arithmetic that can't operate on a
+        # string label.
+        mt_position = act_manifold.resolve_position(mt.position)
+        # Linear baseline: the straight chord through activation space
+        # from the manifold point at node 0 to the term's position, per
+        # layer — what plain additive steering would do instead of
+        # following the manifold's curvature.
+        origin = act_manifold.node_coords[0]
+        chord = {
+            L: (
+                act_manifold.manifold_point(L, mt_position)
+                - act_manifold.manifold_point(L, origin)
+            )
+            for L in act_manifold.layer_indices
+        }
+        session.steer("__manifold_linear_baseline__", Profile(chord))
+        _ltext, lmean, lmax = _score(
+            f"{mt.coeff:g} __manifold_linear_baseline__"
+        )
+        pos_label = (
+            mt.position if isinstance(mt.position, str)
+            else ",".join(f"{c:g}" for c in mt_position)
+        )
+        rows.append({
+            "label": "linear-chord", "steering": f"chord@{pos_label}",
+            "mean_bhattacharyya": lmean, "max_bhattacharyya": lmax,
+        })
+
+    if args.json_output:
+        print(_json.dumps({"prompt": args.prompt, "results": rows}, indent=2))
+        return
+    print()
+    print("behavior-manifold naturalness  (lower = more natural)")
+    for r in rows:
+        print(
+            f"  {r['label']:<14} mean D_B={r['mean_bhattacharyya']:.4f}  "
+            f"max D_B={r['max_bhattacharyya']:.4f}"
+        )
+    if len(rows) == 2:
+        delta = rows[1]["mean_bhattacharyya"] - rows[0]["mean_bhattacharyya"]
+        verdict = (
+            "manifold steering stays closer to the behavior manifold"
+            if delta > 0 else
+            "linear steering stays closer to the behavior manifold"
+        )
+        print(f"  -> {verdict} (Δmean={delta:+.4f})")
+
+
+def _run_experiment_transcript(args: argparse.Namespace) -> None:
+    """Dispatch ``drowse experiment transcript <verb>``."""
+    cmd = getattr(args, "transcript_cmd", None)
+    if cmd is None:
+        _print_verb_menu("experiment transcript", _TRANSCRIPT_VERBS)
+        sys.exit(0)
+    # Registered-subparser invariant: argparse rejects an unknown verb.
+    _TRANSCRIPT_RUNNERS[cmd](args)
+
+
+def _parse_grid_terms(raw_terms: list[str]) -> dict[str, list[float]]:
+    from drowse.cli.alpha_grid import AlphaListError, parse_alpha_list
+
+    grid: dict[str, list[float]] = {}
+    for raw in raw_terms:
+        if "=" not in raw:
+            print(
+                f"experiment fan: grid term must be CONCEPT=ALPHAS, got {raw!r}",
+                file=sys.stderr,
+            )
+            sys.exit(2)
+        name, alpha_text = raw.split("=", 1)
+        name = name.strip()
+        if not name:
+            print("experiment fan: grid concept cannot be empty", file=sys.stderr)
+            sys.exit(2)
+        try:
+            alphas = parse_alpha_list(alpha_text)
+        except AlphaListError as e:
+            print(f"experiment fan: {name}: {e}", file=sys.stderr)
+            sys.exit(2)
+        if not alphas:
+            print(f"experiment fan: {name}: alpha list is empty", file=sys.stderr)
+            sys.exit(2)
+        grid[name] = [float(a) for a in alphas]
+    return grid
+
+
+def _run_experiment_fan(args: argparse.Namespace) -> None:
+    import json as _json
+
+    _pkg._load_effective_config(args, default_max_tokens=256)
+    _pkg._print_startup(args)
+    session = _pkg._make_session(args)
+    _pkg._print_model_info(session)
+
+    grid = _parse_grid_terms(args.grid)
+
+    # A grid is sequential generations; without a per-completion line the
+    # user sees nothing between "Loaded N probes" and the final table.  JSON
+    # mode stays silent so stdout remains a single parseable document.
+    total = 1
+    for alphas in grid.values():
+        total *= len(alphas)
+
+    def _report(idx: int, result: Any, alphas: dict[str, float]) -> None:
+        row = ", ".join(f"{k}={v:+.3f}" for k, v in alphas.items())
+        print(
+            f"  {idx + 1}/{total}: {row} ({result.token_count} tokens)",
+            flush=True,
+        )
+
+    runset = session.generate_sweep(
+        args.prompt,
+        grid,
+        base_steering=args.base_steering,
+        stateless=False,
+        on_result=None if args.json_output else _report,
+    )
+    if args.json_output:
+        print(_json.dumps(runset.to_dict(), indent=2))
+        return
+    print(f"experiment fan: {len(runset)} run(s)")
+    for idx, result in enumerate(runset):
+        node_id = runset.node_ids[idx] if idx < len(runset.node_ids) else None
+        row = runset.grid[idx] if idx < len(runset.grid) else {}
+        row_str = ", ".join(f"{k}={v:+.3f}" for k, v in row.items())
+        node_str = f" node={node_id[:8]}" if node_id else ""
+        print(
+            f"{idx:>3}: {row_str}{node_str} "
+            f"tokens={result.token_count} finish={result.finish_reason}"
+        )
+
+
+def _run_transcript_run(args: argparse.Namespace) -> None:
+    from drowse.core.transcript import (
+        Transcript, TranscriptError,
+    )
+
+    transcript_path = Path(args.path)
+    if not transcript_path.is_file():
+        print(f"transcript run: {transcript_path}: file not found", file=sys.stderr)
+        sys.exit(2)
+    try:
+        transcript = Transcript.load(transcript_path)
+    except TranscriptError as e:
+        print(f"transcript run: {e}", file=sys.stderr)
+        sys.exit(2)
+
+    _pkg._load_effective_config(args, default_max_tokens=256)
+    if not args.model:
+        if transcript.model_id:
+            args.model = transcript.model_id
+        else:
+            print(
+                "transcript run: model required (pass <model> or include "
+                "`model_id` in the transcript)",
+                file=sys.stderr,
+            )
+            sys.exit(2)
+
+    _pkg._print_startup(args)
+    session = _pkg._make_session(args)
+    _pkg._print_model_info(session)
+
+    # Import via ``default`` so the transcript lands as a fresh branch
+    # under the synthetic root; replay walks the imported branch and
+    # reports drift inline.
+    try:
+        leaf_id = transcript.import_into(
+            session, mode="default", strict=args.strict,
+        )
+    except TranscriptError as e:
+        print(f"transcript run: {e}", file=sys.stderr)
+        sys.exit(2)
+
+    print(f"transcript: {len(transcript.turns)} turns loaded "
+          f"(leaf: {leaf_id[:8]})")
+    print()
+    for idx, turn in enumerate(transcript.turns):
+        if turn.role != "user":
+            continue
+        print(f"--- replay turn {idx} ---")
+        print(f"user: {turn.text[:80]}")
+        # Look ahead for the assistant turn this user prompt produced.
+        expected = None
+        if idx + 1 < len(transcript.turns) and transcript.turns[idx + 1].role == "assistant":
+            expected = transcript.turns[idx + 1]
+        try:
+            recipe = expected.recipe if expected is not None else None
+            steering = recipe.steering if recipe is not None else None
+            sampling = recipe.sampling if recipe is not None else None
+            result = session.generate(
+                turn.text,
+                steering=steering,
+                sampling=sampling,
+                stateless=True,
+            ).first
+        except Exception as e:
+            print(f"  replay failed: {e}")
+            continue
+        print(f"assistant: {result.text[:120]}")
+        if expected is not None and expected.readings:
+            actual = {
+                name: (reading.coords[0] if reading.coords else 0.0)
+                for name, reading in result.probe_readings.items()
+            }
+            deltas = []
+            for name, expected_v in expected.readings.items():
+                actual_v = actual.get(name, 0.0)
+                deltas.append((name, actual_v - expected_v, expected_v, actual_v))
+            deltas.sort(key=lambda x: abs(x[1]), reverse=True)
+            print("  readings drift (top 5):")
+            for name, d, ev, av in deltas[:5]:
+                print(f"    {name:<32}  Δ {d:+.4f}  (expected {ev:+.4f} → got {av:+.4f})")
+        print()
+
+
+# The three research verbs, plus the one nested ``transcript`` leaf.  Declared
+# after their runners so the tables carry bound functions.
+_EXPERIMENT_RUNNERS = {
+    "fan":         _run_experiment_fan,
+    "transcript":  _run_experiment_transcript,
+    "naturalness": _run_experiment_naturalness,
+}
+
+_TRANSCRIPT_RUNNERS = {
+    "run": _run_transcript_run,
+}