drowse.ai 5.3.1__py3-none-any.whl

drowse/cli/runners/sae.py

@@ -0,0 +1,304 @@
+"""``drowse sae <verb>`` runners and dispatch table."""
+
+from __future__ import annotations
+
+import argparse
+import sys
+
+import drowse.cli.runners as _pkg
+from drowse.cli.parsers import _SAE_VERBS
+from drowse.cli.runners.shared import _print_verb_menu, _drowse_error_exit
+
+
+def _load_sae_training_corpus(args: argparse.Namespace) -> tuple[list[str], str]:
+    import json as _json
+    import math
+    from pathlib import Path
+
+    target_docs = max(1, math.ceil(args.tokens / args.seq_len))
+    if args.corpus is not None:
+        path = Path(args.corpus)
+        if not path.exists():
+            print(f"sae train: corpus file not found: {path}", file=sys.stderr)
+            sys.exit(2)
+        docs: list[str] = []
+        with open(path, encoding="utf-8") as handle:
+            for line in handle:
+                line = line.strip()
+                if not line:
+                    continue
+                if line.startswith("{"):
+                    try:
+                        payload = _json.loads(line)
+                        line = str(payload.get("text", "")) or line
+                    except _json.JSONDecodeError:
+                        pass
+                docs.append(line)
+                if len(docs) >= target_docs:
+                    break
+        if not docs:
+            print("sae train: corpus has no non-empty documents", file=sys.stderr)
+            sys.exit(2)
+        return docs, f"file:{path.name}"
+    from drowse.core.jlens import JacobianLensError
+    from drowse.io.lens import DEFAULT_LENS_CORPUS, stream_lens_corpus
+
+    print(f"Streaming {target_docs} documents from {DEFAULT_LENS_CORPUS}...")
+    try:
+        return stream_lens_corpus(target_docs)
+    except JacobianLensError as exc:
+        print(f"sae train: {exc}", file=sys.stderr)
+        sys.exit(2)
+
+
+def _run_sae_train(args: argparse.Namespace) -> None:
+    import json
+    from drowse.core.session import DrowseSession
+
+    if args.learning_rate <= 0 or args.l1 < 0 or args.dead_threshold < 0:
+        print("sae train: learning-rate must be > 0 and sparsity values >= 0", file=sys.stderr)
+        sys.exit(2)
+    docs, spec = _load_sae_training_corpus(args)
+    _pkg._print_startup(args)
+    with DrowseSession.from_pretrained(
+        args.model, device=args.device, quantize=args.quantize, probes=[],
+    ) as session:
+        _pkg._print_model_info(session)
+        layer = args.layer
+        if layer is None:
+            layer = round(0.65 * max(len(session.layers) - 1, 0))
+        result = session.train_sae(
+            args.name, docs, layer=layer, corpus_spec=spec,
+            tokens=args.tokens, seq_len=args.seq_len, batch_size=args.batch_size,
+            d_sae=args.width, expansion=args.expansion,
+            learning_rate=args.learning_rate, l1_coefficient=args.l1,
+            dead_feature_threshold=args.dead_threshold, seed=args.seed,
+            force=args.force, on_progress=lambda message: print(f"  {message}", flush=True),
+        )
+    if args.json_output:
+        print(json.dumps(result, indent=2))
+        return
+    metrics = result["metrics"]
+    print(
+        f"Trained {result['source']} for {args.model}: L{layer}, "
+        f"{metrics['d_sae']} features, {metrics['tokens_trained']:,} tokens"
+    )
+    print(f"Artifact: {result['artifact']}")
+
+
+def _model_shape_from_config(model_id: str) -> tuple[int, int]:
+    """``(n_layers, hidden_dim)`` for ``model_id`` without loading weights.
+
+    The CLI-local name kept as this module's monkeypatch seam; the shape rule
+    itself lives once in :func:`drowse.core.model.config_model_shape`, shared
+    with the manifold fit's weight-free no-op preflight.
+    """
+    from drowse.core.model import config_model_shape
+
+    return config_model_shape(model_id)
+
+
+def _run_sae_fetch(args: argparse.Namespace) -> None:
+    """Resolve a provider SAE release and pin its binding — no model load.
+
+    The ``lens fetch`` shape: provider weights land in the Hugging Face cache
+    and Drowse persists only the resolved release/layer binding.  Everything
+    ``load_sae`` validates for an external release is provable without the
+    base model — registry compatibility from the release entry, covered
+    layers against the config's layer count, and the decoder row against its
+    hidden size — so the transformer is never instantiated.
+    """
+    import json
+    from drowse.core.sae import (
+        load_sae_backend, select_runtime_layer, validate_residual_width,
+    )
+    from drowse.io.sae import save_sae_metadata
+
+    if not args.source.startswith("saelens:"):
+        print("sae fetch: source must be saelens:RELEASE", file=sys.stderr)
+        sys.exit(2)
+    release = args.source[len("saelens:"):]
+    if not release:
+        print("sae fetch: release must not be empty", file=sys.stderr)
+        sys.exit(2)
+    try:
+        n_layers, hidden_dim = _model_shape_from_config(args.model)
+    except (OSError, ValueError) as exc:
+        print(
+            f"sae fetch: could not read the config for {args.model}: {exc}",
+            file=sys.stderr,
+        )
+        sys.exit(2)
+    if not args.json_output:
+        # Mirror ``lens fetch``: announce before the provider download,
+        # which is otherwise silent when hub progress bars are off.
+        print(f"Fetching {args.source} into Hugging Face cache...", flush=True)
+    backend = load_sae_backend(
+        release, revision=args.revision, model_id=args.model,
+        device="cpu", dtype=None, warn_on_multiple=False,
+    )
+    covered = frozenset(idx for idx in backend.layers if 0 <= idx < n_layers)
+    selected = select_runtime_layer(covered, n_layers, args.layer)
+    width = int(backend.feature_count(selected))  # materializes the weights
+    if width <= 0:
+        print(f"sae fetch: release {release!r} has no features", file=sys.stderr)
+        sys.exit(1)
+    validate_residual_width(backend, selected, hidden_dim)
+    save_sae_metadata(args.model, release, {
+        "layer": selected,
+        "width": width,
+        "revision": backend.revision,
+        "fingerprint": backend.fingerprint,
+        "sae_id": backend.sae_ids_by_layer.get(str(selected)),
+        "repo_id": backend.repo_id,
+        "neuronpedia_id": backend.neuronpedia_ids_by_layer.get(str(selected)),
+    })
+    if args.json_output:
+        print(json.dumps({
+            "model": args.model,
+            "source": args.source,
+            "layer": selected,
+            "width": width,
+            "revision": backend.revision,
+            "repo_id": backend.repo_id,
+        }, indent=2))
+    else:
+        print(
+            f"Fetched {args.source} for {args.model}: "
+            f"L{selected}, {width} features"
+        )
+        print("Provider weights remain in the Hugging Face cache; binding is active.")
+
+
+def _run_sae_ls(args: argparse.Namespace) -> None:
+    import json
+    from drowse.io.sae import list_sae_sources
+
+    rows = list_sae_sources(args.model)
+    if args.json_output:
+        print(json.dumps(rows, indent=2))
+        return
+    if not rows:
+        print(f"No SAE sources for {args.model}.")
+        return
+    for row in rows:
+        marker = "*" if row["active"] else " "
+        print(f"{marker} {row['source']}  L{row['layer']}  {row['features']} features")
+
+
+def _resolve_active_sae_source(model: str, source: str | None) -> str | None:
+    if source is not None:
+        return source
+    from drowse.io.sae import load_active_sae_source
+
+    active = load_active_sae_source(model)
+    if active is None:
+        return None
+    return (
+        f"local:{active['name']}"
+        if active["kind"] == "local"
+        else f"saelens:{active['name']}"
+    )
+
+
+def _run_sae_show(args: argparse.Namespace) -> None:
+    import json
+    from drowse.io.sae import load_active_sae_source, load_sae_metadata
+    from drowse.io.sae_artifacts import load_local_sae_manifest
+
+    source = _resolve_active_sae_source(args.model, args.source)
+    if source is None:
+        print(f"no SAE source for {args.model}", file=sys.stderr)
+        sys.exit(1)
+    if source.startswith("local:"):
+        payload = load_local_sae_manifest(args.model, source[6:])
+    elif source.startswith("saelens:"):
+        payload = load_sae_metadata(args.model, source[len("saelens:"):])
+    else:
+        print("sae show: source must be local:NAME or saelens:RELEASE", file=sys.stderr)
+        sys.exit(2)
+    if payload is None:
+        print(f"no SAE source {source} for {args.model}", file=sys.stderr)
+        sys.exit(1)
+    active = load_active_sae_source(args.model)
+    out = {
+        "model": args.model,
+        "source": source,
+        "active": bool(
+            active is not None
+            and (
+                (source.startswith("local:") and active["kind"] == "local" and active["name"] == source[6:])
+                or (source.startswith("saelens:") and active["kind"] == "saelens" and active["name"] == source[len("saelens:"):])
+            )
+        ),
+        **payload,
+    }
+    if args.json_output:
+        print(json.dumps(out, indent=2))
+        return
+    print(f"SAE {source} for {args.model}{' (active)' if out['active'] else ''}")
+    print(f"  layer:    {payload['layer']}")
+    print(f"  features: {payload.get('d_sae', payload.get('width'))}")
+    if source.startswith("local:"):
+        print(f"  tokens:   {payload['tokens_trained']}")
+        print(f"  corpus:   {payload['corpus_spec']}")
+    else:
+        print(f"  revision: {payload['revision']}")
+        print("  payload:  Hugging Face cache (provider-owned)")
+
+
+def _run_sae_use(args: argparse.Namespace) -> None:
+    from drowse.io.sae import use_sae_source
+
+    use_sae_source(args.model, args.source)
+    print(f"Active SAE for {args.model}: {args.source}")
+
+
+def _run_sae_rm(args: argparse.Namespace) -> None:
+    from drowse.io.sae import remove_sae_binding
+    from drowse.io.sae_artifacts import remove_local_sae
+
+    source = _resolve_active_sae_source(args.model, args.source)
+    if source is None:
+        print(f"no SAE source for {args.model}", file=sys.stderr)
+        sys.exit(1)
+    if not args.yes:
+        action = "Remove local artifact" if source.startswith("local:") else "Forget external binding"
+        answer = input(f"{action} {source}? [y/N] ").strip().lower()
+        if answer not in {"y", "yes"}:
+            print("Aborted.")
+            return
+    if source.startswith("local:"):
+        removed = remove_local_sae(args.model, source[6:])
+    elif source.startswith("saelens:"):
+        removed = remove_sae_binding(args.model, source[len("saelens:"):])
+    else:
+        print("sae rm: source must be local:NAME or saelens:RELEASE", file=sys.stderr)
+        sys.exit(2)
+    if not removed:
+        print(f"no SAE source {source} for {args.model}", file=sys.stderr)
+        sys.exit(1)
+    if source.startswith("local:"):
+        print(f"Removed Drowse-owned SAE {source}.")
+    else:
+        print("Forgot SAELens binding; Hugging Face cache was not modified.")
+
+
+_SAE_RUNNERS = {
+    "train": _run_sae_train,
+    "fetch": _run_sae_fetch,
+    "ls": _run_sae_ls,
+    "show": _run_sae_show,
+    "use": _run_sae_use,
+    "rm": _run_sae_rm,
+}
+
+
+@_drowse_error_exit
+def _run_sae(args: argparse.Namespace) -> None:
+    cmd = getattr(args, "sae_cmd", None)
+    if cmd is None:
+        _print_verb_menu("sae", _SAE_VERBS)
+        sys.exit(0)
+    # Registered-subparser invariant: argparse rejects an unknown verb.
+    _SAE_RUNNERS[cmd](args)

drowse/cli/runners/serve.py

@@ -0,0 +1,220 @@
+"""``drowse serve`` runner and its live-readout startup policy helpers."""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+from typing import Any, Sequence
+
+import drowse.cli.runners as _pkg
+from drowse.cli.runners.shared import (
+    _drowse_error_exit, _setup_steering_vectors, _warmup_session,
+)
+
+
+@_drowse_error_exit
+def _enable_serve_live_lens_if_compatible(session: Any) -> bool:
+    """Attach the highest-priority compatible cached lens and make it live."""
+    try:
+        compatible = False
+        model_id = getattr(session, "model_id", None)
+        if isinstance(model_id, str) and model_id:
+            from drowse.io.lens_sources import (
+                lens_source_preference_key,
+                list_lens_sources,
+            )
+
+            # Startup applies one deterministic attachment policy even when a
+            # prior fit/fetch left another source active. Explicit switches
+            # still take effect immediately for the running session.
+            rows = sorted(
+                list_lens_sources(model_id),
+                key=lambda row: lens_source_preference_key(
+                    str(row.get("source", "")),
+                ),
+            )
+            for row in rows:
+                source = row.get("source")
+                if not isinstance(source, str) or not source:
+                    continue
+                try:
+                    session.select_jlens_source(source)
+                    compatible = session.has_compatible_jlens()
+                except Exception:  # noqa: BLE001 — try the next cached source
+                    continue
+                if compatible:
+                    break
+        else:
+            compatible = session.has_compatible_jlens()
+        if not compatible:
+            return False
+        layers = session.lens.set_live(True).layers or ()
+        print(f"Live J-lens readout: on ({len(layers)} fitted layers)")
+        return True
+    except Exception as e:  # noqa: BLE001 — never block serve startup
+        print(f"Live J-lens readout: enable failed ({e})", file=sys.stderr)
+        return False
+
+
+def _best_serve_sae_release(rows: Sequence[dict[str, Any]]) -> str | None:
+    """Choose the strongest provider default for the dashboard runtime.
+
+    Provider-hosted, labelled, official releases win; canonical releases beat
+    broad variant sets, and a curated base release beats its ``-all`` sibling.
+    The final name tie-break keeps the choice stable across registry ordering.
+    Local artifacts remain explicit user choices unless one is already active.
+    """
+    candidates = [
+        row for row in rows
+        if row.get("source") == "saelens"
+        and isinstance(row.get("release"), str)
+        and bool(row["release"].strip())
+    ]
+    if not candidates:
+        return None
+
+    def rank(row: dict[str, Any]) -> tuple[bool, bool, bool, bool, str]:
+        release = str(row["release"])
+        repo_id = row.get("repo_id")
+        official = isinstance(repo_id, str) and repo_id.startswith("google/")
+        return (
+            not official,
+            not bool(row.get("neuronpedia")),
+            "canonical" not in release,
+            release.endswith("-all"),
+            release,
+        )
+
+    return str(min(candidates, key=rank)["release"])
+
+
+@_drowse_error_exit
+def _enable_serve_live_sae_if_available(session: Any) -> bool:
+    """Attach a cached or best-provider SAE and start its live view."""
+    try:
+        info = session.sae_info
+        if info is None:
+            from drowse.io.sae import list_sae_sources
+
+            # Prepared artifacts are the cheapest and most intentional
+            # default. This also repairs caches created before active.json.
+            rows = sorted(
+                list_sae_sources(session.model_id),
+                key=lambda row: (not bool(row.get("active")), row.get("kind") != "local"),
+            )
+            for row in rows:
+                source = row.get("source")
+                if not isinstance(source, str) or not source:
+                    continue
+                release = source.removeprefix("saelens:")
+                try:
+                    info = session.load_sae(release)
+                except Exception:  # noqa: BLE001 — try the next cached source
+                    continue
+                break
+        if info is None:
+            from drowse.core.sae import list_sae_releases
+
+            release = _best_serve_sae_release(list_sae_releases(session.model_id))
+            if release is None:
+                return False
+            info = session.load_sae(release)
+        state = session.sae.set_live(True)
+        print(
+            "Live SAE readout: on "
+            f"({info.get('release', 'active')} at L{state.layer})"
+        )
+        return True
+    except Exception as e:  # noqa: BLE001 — never block serve startup
+        print(f"Live SAE readout: enable failed ({e})", file=sys.stderr)
+        return False
+
+
+def _run_serve(args: argparse.Namespace) -> None:
+    try:
+        import fastapi  # noqa: F401
+        import uvicorn
+    except ImportError:
+        # fastapi + uvicorn are base dependencies since v3.x; this only
+        # fires when they've been uninstalled out from under drowse.
+        print(
+            "Server dependencies not installed. Run:\n"
+            "  pip install --upgrade drowse.ai",
+            file=sys.stderr,
+        )
+        sys.exit(1)
+
+    _pkg._load_effective_config(args)
+    from drowse.server.app import is_loopback_host
+    api_key = getattr(args, "api_key", None)
+    if api_key is None:
+        api_key = os.environ.get("DROWSE_API_KEY")
+    if not is_loopback_host(args.host) and not api_key:
+        print("drowse serve: non-loopback binding requires --api-key or DROWSE_API_KEY", file=sys.stderr)
+        sys.exit(2)
+    if not args.model:
+        print(
+            "drowse serve: model required. Pass a HuggingFace repo id (e.g.\n"
+            "  drowse serve google/gemma-2-2b-it\n"
+            "or supply it via -c setup.yaml with a `model:` field.",
+            file=sys.stderr,
+        )
+        sys.exit(2)
+
+    _pkg._print_startup(args)
+    session = _pkg._make_session(args)
+    _pkg._print_model_info(session)
+
+    # Config-file vectors first, then any explicit --steer expression on top.
+    config_expr = getattr(args, "config_vectors", None)
+    if config_expr:
+        _setup_steering_vectors(session, config_expr, verbose=True)
+    steer_expr: str | None = args.steer
+    default_steering = _setup_steering_vectors(session, steer_expr, verbose=True)
+    if default_steering is None and config_expr:
+        from drowse.core.steering_expr import parse_expr
+        default_steering = parse_expr(config_expr)
+
+    from drowse.server import create_app
+    # Default-on: the dashboard ships with the wheel and is the easiest
+    # way for casual users to drive drowse; ``--no-web`` opts out for
+    # production / proxied deployments where ``/`` already belongs to
+    # something else.
+    web_enabled = not getattr(args, "no_web", False)
+    app = create_app(session, default_steering=default_steering,
+                     cors_origins=args.cors or None,
+                     api_key=api_key,
+                     web=web_enabled)
+
+    # The default probe roster — tagged concept axes plus every fitted bundled
+    # multi-node manifold (``personas`` / ``emotions``) — is attached at session
+    # construction (``_bootstrap_manifold_probes``), so the dashboard's probe
+    # rack opens already watching them with no serve-side step.
+
+    _warmup_session(session)
+
+    # Live-lens-on by default: when the model has a fitted Jacobian lens,
+    # serve starts with the full fitted readout streaming, so
+    # the dashboard's J-LENS tab is hot on first load (the toggle still
+    # turns it off per session). Serve-side policy only — the library stays
+    # opt-in. The generation's logit-alternative K also controls
+    # the J-lens readout width.
+    _enable_serve_live_lens_if_compatible(session)
+
+    # Web dashboard policy: attach the strongest compatible provider SAE (or
+    # preserve an explicitly active local/provider source) and turn its live
+    # discovery view on. API-only ``--no-web`` serves do not acquire this
+    # optional provider dependency or download weights implicitly.
+    if web_enabled:
+        _enable_serve_live_sae_if_available(session)
+
+    print(f"\nServing on http://{args.host}:{args.port}")
+    print(f"OpenAI-compatible:  http://{args.host}:{args.port}/v1")
+    print(f"Ollama-compatible:  http://{args.host}:{args.port}/api")
+    print(f"API docs:           http://{args.host}:{args.port}/docs")
+    if args.port != 11434:
+        print("Tip: for drop-in Ollama compatibility, run with `--port 11434`.")
+    from drowse.server.ws_stream import MAX_WS_MESSAGE_BYTES
+    uvicorn.run(app, host=args.host, port=args.port, log_level="info",
+                ws_max_size=MAX_WS_MESSAGE_BYTES)