drowse.ai 5.3.1__py3-none-any.whl

drowse/io/paths.py

@@ -0,0 +1,289 @@
+"""Filesystem path helpers for the ~/.drowse/ tree.
+
+All paths resolve through drowse_home(), which honors the DROWSE_HOME
+environment variable for testing and non-default installs.
+"""
+from __future__ import annotations
+
+import base64
+import os
+import re
+from pathlib import Path
+
+# Variant filename conventions:
+#   raw (DiM)           -> ``<safe_model_id>.safetensors``
+#   SAE-DiM             -> ``<safe_model_id>_sae-<release>.safetensors``
+#   transferred (v1.6)  -> ``<safe_model_id>_from-<safe_src>.safetensors``
+#
+# The literals ``_sae-`` and ``_from-`` are the *kind*
+# separators. Model ids and right-hand-side slugs can legally contain those
+# literals, so filename components escape them before concatenation and the
+# parser splits at the earliest unescaped kind separator.
+# Difference-of-means is the only extraction method (4.0), so the raw
+# tensor carries no method suffix.  Kind suffixes are mutually exclusive
+# — a tensor is at most one of {SAE, transferred}.
+_VARIANT_SEP_SAE = "_sae-"
+_VARIANT_SEP_FROM = "_from-"
+_VARIANT_SEPARATORS: tuple[tuple[str, str], ...] = (
+    (_VARIANT_SEP_SAE, "sae"),
+    (_VARIANT_SEP_FROM, "from"),
+)
+
+
+def encode_release_id(release: str) -> str:
+    """Encode an SAE release as one reversible grammar-safe component."""
+    if not release:
+        raise ValueError("SAE release must not be empty")
+    encoded = base64.b32encode(release.encode("utf-8")).decode("ascii")
+    return "_z" + encoded.rstrip("=").lower()
+
+
+def decode_release_id(encoded: str) -> str:
+    """Reverse :func:`encode_release_id`, rejecting non-canonical input."""
+    if not encoded.startswith("_z"):
+        raise ValueError(f"invalid encoded SAE release {encoded!r}")
+    payload = encoded[2:].upper()
+    payload += "=" * (-len(payload) % 8)
+    try:
+        release = base64.b32decode(payload, casefold=False).decode("utf-8")
+    except (ValueError, UnicodeDecodeError) as exc:
+        raise ValueError(f"invalid encoded SAE release {encoded!r}") from exc
+    if not release or encode_release_id(release) != encoded:
+        raise ValueError(f"invalid encoded SAE release {encoded!r}")
+    return release
+
+
+def _encode_tensor_component(value: str) -> str:
+    """Escape reserved filename separators inside one model/variant slug."""
+    value = value.replace("%", "%25")
+    for separator, _kind in _VARIANT_SEPARATORS:
+        value = value.replace(separator, f"%5F{separator[1:]}")
+    return value
+
+
+def _decode_tensor_component(value: str) -> str:
+    """Reverse :func:`_encode_tensor_component`."""
+    for separator, _kind in _VARIANT_SEPARATORS:
+        value = value.replace(f"%5F{separator[1:]}", separator)
+    return value.replace("%25", "%")
+
+# The trailing ``:<variant>`` scheme understood across selectors / runners.
+# One source of truth for the variant suffix grammar (this module owns the
+# tensor-filename variant scheme); ``io.selectors`` and ``cli.runners`` import it.
+VARIANT_SUFFIX_RE = re.compile(
+    r"^(raw|sae(?:-[a-z0-9._-]+)?|role(?:-[a-z0-9._-]+)?|from(?:-[a-z0-9._-]+)?)$"
+)
+
+def drowse_home() -> Path:
+    """Return the root ~/.drowse/ directory. Honors $DROWSE_HOME override."""
+    override = os.environ.get("DROWSE_HOME") or os.environ.get("POLYTHETIC_HOME") or os.environ.get("SAKLAS_HOME")
+    if override:
+        return Path(override)
+    from drowse.io.brand_migration import migrate_legacy_home
+
+    return migrate_legacy_home(Path.home())
+
+
+def models_dir() -> Path:
+    return drowse_home() / "models"
+
+
+def manifolds_dir() -> Path:
+    """Root of the manifold-steering artifact tree.
+
+    A manifold is a labeled-node artifact placed on an n-dimensional domain.
+    """
+    return drowse_home() / "manifolds"
+
+
+def templates_dir() -> Path:
+    """Root of the templated-completion artifact tree.
+
+    A template is a first-class artifact (peer to a manifold): a slot, a set
+    of candidate values, and one or more multi-turn contexts whose final
+    assistant turn carries the slot. Two consumers read it — the manifold fit
+    (pools the slot-filled assistant centroids into nodes) and the completion
+    scorer (the restricted-choice logprob distribution over the values).
+    """
+    return drowse_home() / "templates"
+
+
+def ensure_within(root: Path, *parts: str) -> Path:
+    """Join ``parts`` onto ``root`` and verify the result stays inside ``root``.
+
+    Path-traversal barrier for user-supplied path components — namespaces,
+    concept / manifold names, model ids, and manifest-relative filenames all
+    reach the filesystem layer from HTTP request bodies, CLI args, and
+    downloaded artifact manifests. A ``..`` segment
+    or an absolute component must not be allowed to escape the ``~/.drowse/``
+    subtree.
+
+    Normalizes the joined path (collapsing ``..`` segments) and rejects a
+    result that does not stay under ``root`` — the normalize-then-prefix-check
+    idiom. Returns the normalized in-bounds path; raises :class:`ValueError`
+    otherwise. ``NAME_REGEX`` / ``_LABEL_REGEX`` already reject traversal
+    syntax upstream, but this is the defense-in-depth barrier at the boundary
+    where the string actually becomes a path.
+    """
+    root_norm = os.path.normpath(os.fspath(root))
+    candidate = os.path.normpath(os.path.join(root_norm, *parts))
+    if candidate != root_norm and not candidate.startswith(root_norm + os.sep):
+        raise ValueError(f"unsafe path component {parts!r} escapes {root_norm!r}")
+    return Path(candidate)
+
+
+def manifold_dir(namespace: str, name: str) -> Path:
+    return ensure_within(manifolds_dir(), namespace, name)
+
+
+def neutral_statements_path() -> Path:
+    return drowse_home() / "neutral_statements.json"
+
+
+def baseline_prompts_path() -> Path:
+    """User-override path for the shared A2 baseline user prompts.
+
+    Conversational extraction generates each node's corpus as responses to
+    this fixed prompt set; a node corpus aligns ``response[i] -> prompt[i % k]``.
+    Falls back to the bundled ``drowse/data/baseline_prompts.json`` when absent.
+    """
+    return drowse_home() / "baseline_prompts.json"
+
+
+def safe_model_id(model_id: str) -> str:
+    """Encode any HF or local model id as one reversible filename component."""
+    if not model_id:
+        raise ValueError("model id must not be empty")
+    payload = base64.urlsafe_b64encode(model_id.encode("utf-8")).decode("ascii")
+    return "_z" + payload.rstrip("=")
+
+
+def unsafe_model_id(safe_id: str) -> str:
+    """Reverse :func:`safe_model_id`, rejecting malformed escape sequences."""
+    if not safe_id.startswith("_z"):
+        raise ValueError(f"invalid safe model id {safe_id!r}")
+    encoded = safe_id[2:]
+    encoded += "=" * (-len(encoded) % 4)
+    try:
+        model_id = base64.b64decode(
+            encoded, altchars=b"-_", validate=True,
+        ).decode("utf-8")
+    except (ValueError, UnicodeDecodeError) as exc:
+        raise ValueError(f"invalid encoded safe model id {safe_id!r}") from exc
+    if not model_id or safe_model_id(model_id) != safe_id:
+        raise ValueError(f"invalid encoded safe model id {safe_id!r}")
+    return model_id
+
+
+def model_dir(model_id: str) -> Path:
+    return ensure_within(models_dir(), safe_model_id(model_id))
+
+
+def safe_sae_suffix(release: str | None) -> str:
+    """Filename suffix for an SAE variant.  ``None``/``""`` = raw (no suffix)."""
+    if not release:
+        return ""
+    return f"{_VARIANT_SEP_SAE}{encode_release_id(release)}"
+
+
+def safe_from_suffix(source_model_id: str | None) -> str:
+    """Reversibly encode a source model id for a transfer variant suffix."""
+    if not source_model_id:
+        return ""
+    return f"{_VARIANT_SEP_FROM}{encode_release_id(source_model_id)}"
+
+
+def tensor_filename(
+    model_id: str,
+    *,
+    release: str | None = None,
+    transferred_from: str | None = None,
+    model_id_is_safe: bool = False,
+    transferred_from_is_encoded: bool = False,
+) -> str:
+    """Construct the canonical tensor filename.
+
+    At most one of ``release`` and ``transferred_from`` may be set — composed
+    kind variants are not supported. ``transferred_from`` accepts an HF/local model id by
+    default. Internal callers holding a parsed ``from-*`` filename component
+    set ``transferred_from_is_encoded=True``. The target equivalent remains
+    ``model_id_is_safe=True``.
+
+    The raw tensor at the canonical path is difference-of-means (the only
+    extraction method as of 4.0), so it carries no method suffix.
+
+    """
+    if release and transferred_from:
+        raise ValueError(
+            "tensor_filename: release and transferred_from are mutually exclusive"
+        )
+    target_safe = model_id if model_id_is_safe else safe_model_id(model_id)
+    target = _encode_tensor_component(target_safe)
+    if release:
+        return f"{target}{safe_sae_suffix(release)}.safetensors"
+    if transferred_from:
+        suffix = (
+            f"{_VARIANT_SEP_FROM}{transferred_from}"
+            if transferred_from_is_encoded
+            else safe_from_suffix(transferred_from)
+        )
+        return f"{target}{suffix}.safetensors"
+    return f"{target}.safetensors"
+
+
+def sidecar_filename(
+    model_id: str,
+    *,
+    release: str | None = None,
+    transferred_from: str | None = None,
+) -> str:
+    """Sidecar JSON partner for a tensor filename."""
+    if release and transferred_from:
+        raise ValueError(
+            "sidecar_filename: release and transferred_from are mutually exclusive"
+        )
+    target = _encode_tensor_component(safe_model_id(model_id))
+    if release:
+        return f"{target}{safe_sae_suffix(release)}.json"
+    if transferred_from:
+        return f"{target}{safe_from_suffix(transferred_from)}.json"
+    return f"{target}.json"
+
+
+def parse_tensor_filename(
+    filename: str,
+) -> tuple[str, str | None] | None:
+    """Reverse of :func:`tensor_filename`. Returns ``(safe_model_id, variant)``.
+
+    ``variant`` is one of:
+      * ``None`` — raw DiM tensor (no separator).
+      * ``"sae-<release>"`` — SAE-DiM variant.
+      * ``"from-<safe_src>"`` — transferred-from variant.
+    The variant string carries its kind tag so callers can dispatch
+    without re-parsing.  Returns ``None`` for filenames that aren't
+    ``.safetensors``.
+    """
+    if not filename.endswith(".safetensors"):
+        return None
+    stem = filename[: -len(".safetensors")]
+    matches = [
+        (stem.find(separator), separator, kind)
+        for separator, kind in _VARIANT_SEPARATORS
+        if separator in stem
+    ]
+    if matches:
+        index, separator, kind = min(matches, key=lambda item: item[0])
+        model = _decode_tensor_component(stem[:index])
+        value = _decode_tensor_component(stem[index + len(separator):])
+        try:
+            unsafe_model_id(model)
+            decode_release_id(value)
+        except ValueError:
+            return None
+        return model, f"{kind}-{value}"
+    model = _decode_tensor_component(stem)
+    try:
+        unsafe_model_id(model)
+    except ValueError:
+        return None
+    return model, None

drowse/io/probes_bootstrap.py

@@ -0,0 +1,58 @@
+"""Bootstrap the per-model probe-centering baseline and the bundled probe roster.
+
+A steering vector lives as a 2-node ``pca`` manifold (4.0), so the bundled
+probe set is sourced by folding fitted manifolds (see
+``DrowseSession._bootstrap_manifold_probes``); this module owns the two pieces
+that feed it: the per-model probe-centering means (:func:`bootstrap_layer_means`,
+derived from the neutral-activation cache) and the tag→manifold roster
+(:func:`load_default_manifolds`).
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+
+def load_default_manifolds() -> dict[str, list[str]]:
+    """Return {tag: [manifold_name, ...]} for bundled default/ manifolds.
+
+    A steering vector lives as a 2-node ``pca`` manifold, tagged
+    (``manifold.json::tags``) for category-grouped probe bootstrap.  Triggers
+    first-run materialization of the bundled artifacts.
+    """
+    from drowse.io.bootstrap import materialize_bundled_artifacts
+    from drowse.io.manifolds import iter_manifold_folders
+
+    materialize_bundled_artifacts()
+    by_tag: dict[str, list[str]] = {}
+    for _ns, mf in iter_manifold_folders(namespace="default"):
+        for tag in mf.tags or []:
+            by_tag.setdefault(tag, []).append(mf.name)
+    return by_tag
+
+
+def bootstrap_layer_means(
+    model: Any, tokenizer: Any, layers: torch.nn.ModuleList, model_info: dict[str, Any],
+) -> dict[int, torch.Tensor]:
+    """Per-layer neutral mean activation for probe centering.
+
+    Derived as ``X.mean(0)`` from the per-model neutral-activation cache
+    (:func:`drowse.io.alignment.load_or_compute_neutral_activations`).  The
+    neutral mean *is* the probe-centering baseline — same corpus, same
+    last-content-token fp32 pooling the whitener's covariance is built from —
+    so there is no separate ``layer_means`` forward pass or disk cache.  The
+    neutral-activation cache is the single per-model artifact both the
+    centering mean and the Mahalanobis whitener read; the means fall out of it
+    for free, so a cold model pays one neutral-corpus forward loop instead of
+    two.  Stale on ``neutral_statements.json`` drift via the neutral cache's
+    own sha256 key.
+    """
+    from drowse.io.alignment import load_or_compute_neutral_activations
+
+    model_id = model_info.get("model_id", "unknown")
+    acts = load_or_compute_neutral_activations(
+        model, tokenizer, layers, model_id=model_id,
+    )
+    return {idx: X.mean(dim=0) for idx, X in acts.items()}