From c247b97c510deb765807315a0ab3fe7abc00087e Mon Sep 17 00:00:00 2001 From: Jen D Date: Fri, 30 Oct 2026 22:50:01 +0800 Subject: [PATCH 1/3] Add `modulo` CLI: search, list, add from HF, extract, verify, run MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - modulo search: one search over your registries, the community index and modulo_catalog.yaml, with a status and next step per hit - modulo add hf:owner/repo[@rev]: pin the commit, download and inspect vector files (format/layer from shape, dict .pt metadata or filename), check them against the model config, and register them unverified (±8). For bundles, keep only paths naming the model (other models can share hidden size) and one layer per concept. --dry-run, --as-index-entry for draft index entries - modulo extract / verify (strength sweep, --save-range) / run (passthrough to agent/loop.py, which gains --model and --steer NAME=STRENGTH) - --json everywhere with progress on stderr; exit 1 when nothing was done - Sweep code moved to modulo/verify.py; model_dims to modulo/community.py - Loader accepts dict .pt files ({"vector": ...}); registry.yaml keeps its header comments on save; MODULO_MODEL sets the default model - setup.sh does an editable install; docs and CPU tests for the above Co-Authored-By: Claude Opus 5.5 --- .gitignore | 3 + AGENTS.md | 34 +++- CONTRIBUTING.md | 4 +- README.md | 9 + agent/loop.py | 25 ++- community/check_vectors.py | 14 +- community/gpu_verify.py | 69 +------ config.py | 4 +- load_model.py | 9 +- modulo/cli.py | 334 +++++++++++++++++++++++++++++++ modulo/community.py | 13 ++ modulo/importer.py | 389 +++++++++++++++++++++++++++++++++++++ modulo/registry.py | 81 ++++++-- modulo/search.py | 133 +++++++++++++ modulo/verify.py | 162 +++++++++++++++ pyproject.toml | 3 + setup.sh | 4 + tests/test_cli.py | 264 +++++++++++++++++++++++++ verify_vector.py | 73 +------ 19 files changed, 1447 insertions(+), 180 deletions(-) create mode 100644 modulo/cli.py create mode 100644 modulo/importer.py create mode 100644 modulo/search.py create mode 100644 modulo/verify.py create mode 100644 tests/test_cli.py diff --git a/.gitignore b/.gitignore index d742b9f..8e47187 100644 --- a/.gitignore +++ b/.gitignore @@ -31,3 +31,6 @@ Thumbs.db # IDE .vscode/ .idea/ + +# `modulo verify` reports +vectors/*/verify_*.md diff --git a/AGENTS.md b/AGENTS.md index 48d6377..a4f935c 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -62,7 +62,7 @@ At startup the agent combines three sources for its model: Vectors are stored per-model in `vectors//` (for the default model: `vectors/Qwen_Qwen2.5-7B-Instruct/`), resolved relative to the repo, not the working directory. Override the base with `MODULO_VECTORS_DIR`. -To add your own vectors, extract and register them in one step, then check them: +The `modulo` command (next section) is the easiest way to find, add and check vectors. The repo scripts still work too: to extract and register a vector in one step, then check it: ```bash python extract_contrastive.py --name sycophancy --prompt-a "..." --prompt-b "..." --layer 14 \ @@ -73,6 +73,31 @@ python verify_vector.py --vector sycophancy --strength 10 # baseline vs ± Earlier notes said CAA vectors didn't work on Qwen, but those tests ran while steering was silently a no-op during generation, so that needs re-testing. +## Finding, adding and running vectors: the `modulo` command + +`pip install -e .` installs a `modulo` command. Nothing prompts, every subcommand takes `--json` (progress goes to stderr, so stdout parses), and the exit code is 1 when the command didn't do what was asked. `--model` defaults to `$MODULO_MODEL`, else `Qwen/Qwen2.5-7B-Instruct` (`config.py` reads the same variable). + +```bash +modulo search refusal --model Qwen/Qwen3-8B --json # 1. find +modulo add hf:rotalabs/steering-vectors --model Qwen/Qwen3-8B --dry-run # 2. see what fits +modulo add hf:rotalabs/steering-vectors --model Qwen/Qwen3-8B # add it +modulo verify refusal --model Qwen/Qwen3-8B --save-range # 3. check it (loads the model) +modulo run --model Qwen/Qwen3-8B --auto --max-turns 8 --steer refusal=6 # 4. run the demo +``` + +- **`search [WORDS] [--model M] [--status ...]`** searches your registries, the community index and `modulo_catalog.yaml` (override with `MODULO_CATALOG_PATH`). Every word must match. Each hit has a `status` (`local`, `verified`, `community`, `catalog`) and a `next_step`. Catalog hits are sources, not loadable vectors. The ones hosted on HuggingFace come with a `modulo add ... --dry-run` command to try; the rest are code or papers. +- **`list [--model M] [--include-community]`** shows what `SteeringAgent` would load, notes on index files that failed their checks, and which starter vectors would be auto-extracted. +- **`add hf:OWNER/REPO[@REV]`** (also `hf:datasets/OWNER/REPO` or a huggingface.co URL) pins the repo's current commit, downloads its `.safetensors`/`.pt`/`.npz` files (`--files GLOB` to choose; more than `--max-files`, default 50, is refused) and checks each against the model's config. Each file gets one `ok`/`skip` line with a reason. + - **Format and layer** come from the shape: `[hidden]` is `single`, `[num_layers, hidden]` is `per_layer`. The layer comes from a dict `.pt`'s `layer` field, else the filename (`_layer14`, `L14`, `layer_14.pt`). A `[num_layers+1, hidden]` file is skipped, because it probably includes the embeddings and every layer index would be off by one. + - **Bundles**: if some paths name the model (`refusal_qwen3_8b/…`), only those are used, because other models' vectors can share its hidden size. A dict `.pt` whose `model` field names another model is skipped too. Files that differ only by layer are one concept: one is kept, the `--layer` you ask for or else the one nearest the middle of the model. `--all-layers` keeps them all. + - **Names** come from the file or its directory, with the model's name and the layer dropped (`refusal_qwen3_8b/layer_18.pt` → `refusal`). + - **What gets registered**: `add` writes to `vectors//registry.yaml` with an `[Unverified imported vector]` description prefix, a ±8 range and an `hf:…@commit:path sha256=…` source. Running it again is a no-op. A name you already have is kept unless you pass `--replace`. + - **Descriptions**: `--description`/`--sign-convention` only work when adding one vector. For a bundle, edit `registry.yaml` afterwards, because the model only sees the placeholder "no description yet". + - **`--as-index-entry`** prints draft `community/vectors.yaml` entries instead of registering anything. +- **`extract --name N --prompt-a A --prompt-b B`** works like `extract_contrastive.py` for any `--model`. The default layer is the middle one. +- **`verify NAME [--strengths …] [--prompt …] [--save-range]`** sweeps ±strengths (default 4, 8, 12, 15) on the vector's own test prompts plus generic ones. It writes `vectors//verify_.md` with every sample and suggests a range: the largest strengths with no repetition loop. `--save-range` writes that range into your own `registry.yaml`; index vectors' ranges belong to maintainers. "Changed" means *any* difference from the unsteered output, not the claimed effect, and word salad that doesn't repeat isn't caught. **Read the samples** before trusting a vector or its range. Exit code 1 means no safe range was found, or the outputs never changed. +- **`run [ARGS]`** passes everything to `agent/loop.py`, so it needs the repo checkout (an editable install). The loop has two new flags: `--model`, and `--steer NAME=STRENGTH`, which is repeatable and applies a vector as if the model had called `steer`: one full-strength turn, then decay. + ## Strength units `strength` is a percentage of the typical residual-stream norm at the vector's layer. Vectors are normalized to unit length and scaled by `strength / 100 × residual_norm`, so the same strength means the same size push for every vector regardless of how it was extracted. On Qwen2.5-7B, `assistant_axis` gives clear, coherent shifts at ±8–12; at +15 it turns muddled (so its range is -15 to +12), and it collapses into repetition around ±25. Steering applies from the last prompt position onward, so every generated token is steered. @@ -159,6 +184,10 @@ Host the vector on HuggingFace and add an entry to `community/vectors.yaml` (or | `modulo/extraction.py` | Vector extraction (manual + auto) | | `modulo/reports.py` | Report format, validation, `reports.jsonl` read/write | | `modulo/community.py` | Community vector index: entry format and file checks | +| `modulo/cli.py` | The `modulo` command (search, list, add, extract, verify, run) | +| `modulo/search.py` | Search across your registries, the index and the catalog | +| `modulo/importer.py` | `modulo add`: inspect a HuggingFace repo's vector files and register them | +| `modulo/verify.py` | Strength sweeps shared by `modulo verify`, `verify_vector.py` and `gpu_verify.py` | | `community/vectors.yaml` | Community vector index (pointers to hosted vectors, with trust levels) | | `community/check_vectors.py` | Checks the index; run by CI on PRs touching `community/` | | `community/gpu_verify.py` | Maintainer GPU check: strength sweep + suggested safe range | @@ -170,7 +199,7 @@ Host the vector on HuggingFace and add an entry to `community/vectors.yaml` (or | `extract_contrastive.py` | Extract + register a vector via contrastive system prompts | | `extract_vector.py` | Extract + register a CAA vector | | `verify_vector.py` | Compare baseline vs ±strength outputs for a registered vector | -| `modulo_catalog.yaml` | Discovery catalog of 26+ vector sources | +| `modulo_catalog.yaml` | Discovery catalog of 26+ vector sources (searched by `modulo search`) | | `replay.py` | Replay a session log as a timeline | ## Tests @@ -178,6 +207,7 @@ Host the vector on HuggingFace and add an entry to `community/vectors.yaml` (or ```bash pytest # everything; GPU tests skip without CUDA pytest tests/test_agent_logic.py # CPU only: parsing, decay/cooldown, degeneration, repeats, tool validation +pytest tests/test_cli.py # CPU only: search, `modulo add` against a fake hub, naming, layer picking pytest tests/test_steering_gpu.py # loads the model; checks steering really changes generation ``` diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 553ed38..cc9a424 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -45,7 +45,9 @@ Required: Optional but very useful: `revision` (the HF commit of your repo; recommended), `model_revision`, `original_setup`, `known_failure_modes`, `paper`, `contrast_dataset`, `evaluation`, `test_prompts` (up to 5), `notes`. Maintainers set `safe_strength_min`/`safe_strength_max` when promoting. -File formats: a `.safetensors` file must contain exactly one tensor; an `.npz` file stores it under the key `vector`; a `.pt` file is a plain tensor (loaded with `weights_only=True`). With `format: per_layer`, the tensor is indexed by `layer`. +File formats: a `.safetensors` file must contain exactly one tensor; an `.npz` file stores it under the key `vector`; a `.pt` file is a plain tensor, or a dict holding it under `vector` (loaded with `weights_only=True`). With `format: per_layer`, the tensor is indexed by `layer`. + +**Drafting entries automatically.** If your vectors are already on HuggingFace, `modulo add hf:you/your-repo --model --as-index-entry` downloads them, works out `format`, `layer` and `sha256`, checks them against the model's config and prints draft entries pinned to the current commit. Fill in the `TODO` fields and paste them into `community/vectors.yaml`. ## What the automatic checks do diff --git a/README.md b/README.md index 5a32e04..ed67746 100644 --- a/README.md +++ b/README.md @@ -40,6 +40,15 @@ messages += result.tool_messages() # tool results back to the model **Community index + extract**: vectors listed in [`community/vectors.yaml`](community/vectors.yaml) are downloaded and checked for your model (verified ones by default, unverified community ones on request; currently Qwen2.5-7B-Instruct). Any starter vectors your model is missing (sycophancy, refusal, creativity, formality, verbosity, confidence, emotion, humor, optimism, curiosity, introspection) are auto-extracted on first run using contrastive system prompts (several minutes). To contribute vectors, see [CONTRIBUTING.md](CONTRIBUTING.md). +**Find and add more** with the `modulo` command (installed by `pip install -e .`): + +```bash +modulo search refusal --model Qwen/Qwen3-8B # your vectors, the index, and the catalog of known sources +modulo add hf:rotalabs/steering-vectors --model Qwen/Qwen3-8B # pull a repo's vectors into your registry +modulo verify refusal --model Qwen/Qwen3-8B --save-range # strength sweep, suggested safe range +modulo run --model Qwen/Qwen3-8B --auto --steer refusal=6 # the demo +``` + ## Requirements GPU with 24GB+ VRAM (7B models) or 8GB+ (3B models). Python 3.10+, CUDA 12.x. See [AGENTS.md](AGENTS.md) for full details. diff --git a/agent/loop.py b/agent/loop.py index e98d94d..9ab4073 100644 --- a/agent/loop.py +++ b/agent/loop.py @@ -10,6 +10,7 @@ python -m agent.loop --seed "Hi!" # seed with a user message python -m agent.loop --max-turns 20 # cap turns python -m agent.loop --auto # auto-continue (non-interactive) + python -m agent.loop --steer honesty=6 --model Qwen/Qwen2.5-3B-Instruct """ import argparse import json @@ -88,11 +89,29 @@ def main(): help="--auto filler when the model makes no tool call: 'status' (default) is a neutral " "tool-role status line with current steering and no suggestion; 'legacy' is the " "old user message suggesting steer") + parser.add_argument("--model", default=config.MODEL_NAME, help=f"HF model id (default: {config.MODEL_NAME})") + parser.add_argument("--steer", action="append", default=[], metavar="NAME=STRENGTH", + help="apply a vector before the first turn, as if the model had called steer " + "(one full-strength turn, then decay). Repeatable.") args = parser.parse_args() - - model, tokenizer = load_model_and_tokenizer() + steers = [] + for spec in args.steer: + name, _, strength = spec.partition("=") + try: + steers.append((name.strip(), float(strength))) + except ValueError: + parser.error(f"--steer expects NAME=STRENGTH, got {spec!r}") + if not name.strip(): + parser.error(f"--steer expects NAME=STRENGTH, got {spec!r}") + + model, tokenizer = load_model_and_tokenizer(args.model) # Offer every vector we can: unverified index entries too (marked as such in the prompt) - agent = SteeringAgent(model, tokenizer, model_name=config.MODEL_NAME, include_community=True) + agent = SteeringAgent(model, tokenizer, model_name=args.model, include_community=True) + for name, strength in steers: + ok, msg = agent.apply(name, strength) + if not ok: + raise SystemExit(f"--steer {name}: {msg}") + print(f"[steer] {msg}") logger = SessionLogger() kill = KillSwitch() diff --git a/community/check_vectors.py b/community/check_vectors.py index e936241..db0c0fc 100644 --- a/community/check_vectors.py +++ b/community/check_vectors.py @@ -11,13 +11,13 @@ Run with --all to re-check every entry's file. """ import argparse -import json import sys from pathlib import Path sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) from modulo.community import check_vector_file, fetch_vector, load_index, validate_index # noqa: E402 +from modulo.community import model_dims as community_model_dims # noqa: E402 def changed_entries(index: dict, base: dict) -> tuple[list[dict], list[str]]: @@ -40,15 +40,9 @@ def maintainer_warnings(changed: list[dict], base: dict) -> list[str]: def model_dims(entry: dict) -> tuple[int | None, int | None, str | None]: - """(hidden_size, num_layers, note) from the model's config.json on the Hub.""" - from huggingface_hub import hf_hub_download - try: - path = hf_hub_download(entry["model"], "config.json", revision=entry.get("model_revision")) - except Exception as e: # gated or missing models: skip the shape checks, don't fail - return None, None, f"{entry['id']}: couldn't read {entry['model']} config ({type(e).__name__}); shape not checked" - cfg = json.loads(Path(path).read_text()) - cfg = cfg.get("text_config", cfg) if "hidden_size" not in cfg else cfg - return cfg.get("hidden_size"), cfg.get("num_hidden_layers"), None + """(hidden_size, num_layers, note) from the entry's model's config.json on the Hub.""" + hidden, layers, note = community_model_dims(entry["model"], entry.get("model_revision")) + return hidden, layers, note and f"{entry['id']}: {note}" def check(index_path: Path, base_path: Path | None, check_all: bool) -> tuple[list[str], list[str], list[str] | None]: diff --git a/community/gpu_verify.py b/community/gpu_verify.py index 555020d..0982dcd 100644 --- a/community/gpu_verify.py +++ b/community/gpu_verify.py @@ -10,7 +10,6 @@ python community/gpu_verify.py --entry --index path/to/pr/vectors.yaml --pr 12 """ import argparse -import html import shutil import subprocess import sys @@ -25,73 +24,9 @@ from community.check_vectors import model_dims # noqa: E402 from modulo import SteeringAgent, VectorEntry, VectorRegistry # noqa: E402 -from modulo.agent import MAX_STRENGTH, is_degenerate # noqa: E402 +from modulo.agent import MAX_STRENGTH # noqa: E402 from modulo.community import check_vector_file, fetch_vector, load_index # noqa: E402 -from verify_vector import generate # noqa: E402 - -GENERIC_PROMPTS = [ - "Introduce yourself in two sentences.", - "Describe a rainy Monday morning.", - "I think the moon landing was faked. Do you agree?", -] -DEFAULT_STRENGTHS = [4.0, 8.0, 12.0, MAX_STRENGTH] -MAX_NEW_TOKENS = 120 - - -def sweep(agent, name, prompts, strengths): - """{prompt: {signed strength: output}}, including the unsteered baseline at 0.""" - signed = [0.0] + [s * sign for s in strengths for sign in (1, -1)] - return {p: {s: generate(agent, p, name, s, max_new_tokens=MAX_NEW_TOKENS) for s in signed} - for p in prompts} - - -def summarise(outputs, strengths): - """Per signed strength: (prompts changed vs baseline, prompts degenerate); plus suggested range.""" - stats = {} - for s in [x * sign for x in strengths for sign in (1, -1)]: - changed = sum(o[s] != o[0.0] for o in outputs.values()) - degenerate = sum(is_degenerate(o[s]) for o in outputs.values()) - stats[s] = (changed, degenerate) - safe = {} - for sign in (1, -1): - best = None - for s in sorted(strengths): - if stats[s * sign][1]: - break - best = s - safe[sign] = best - return stats, safe - - -def report_markdown(entry, norm, prompts, outputs, stats, safe, strengths): - n = len(prompts) - lines = [f"## GPU verification: `{entry['id']}`", "", - f"`{entry['name']}` on `{entry['model']}`, layer {entry['layer']} " - f"(residual norm ≈ {norm:.1f}). Greedy decoding, {MAX_NEW_TOKENS} new tokens, {n} prompts.", "", - "| strength | changed vs baseline | degenerate |", "|---:|---:|---:|"] - for s in sorted(stats, key=lambda x: (abs(x), -x)): - changed, degenerate = stats[s] - lines.append(f"| {s:+g} | {changed}/{n} | {degenerate}/{n}{' ⚠️' if degenerate else ''} |") - lo, hi = safe[-1], safe[1] - lines.append("") - if lo is None or hi is None: - lines.append(f"**Suggested range:** none. Output degenerates at the lowest tested strength " - f"({min(strengths):g}) on at least one side; don't promote as is.") - else: - lines.append(f"**Suggested range:** `safe_strength_min: {-lo:g}`, `safe_strength_max: {hi:g}` " - f"(largest tested strengths with no degenerate output on that side).") - if not any(c for c, _ in stats.values()): - lines.append("\n**No visible effect:** outputs matched the baseline at every strength.") - lines.append("\nRepetition is detected automatically; incoherent but non-repetitive output isn't, " - "so read the samples before promoting.") - lines.append(f"\nClaimed effect: {entry['sign_convention']}") - for p in prompts: - lines += ["", f"
{html.escape(p)}", ""] - for s in sorted(outputs[p], key=lambda x: (abs(x), -x)): - lines.append(f"{s:+g}
{html.escape(outputs[p][s].strip())}
") - lines.append("
") - return "\n".join(lines) + "\n" - +from modulo.verify import DEFAULT_STRENGTHS, GENERIC_PROMPTS, report_markdown, summarise, sweep # noqa: E402 def main(): root = Path(__file__).resolve().parent diff --git a/config.py b/config.py index c6b9642..9cba4e1 100644 --- a/config.py +++ b/config.py @@ -2,8 +2,8 @@ import os from pathlib import Path -# Model -MODEL_NAME = "Qwen/Qwen2.5-7B-Instruct" +# Model (MODULO_MODEL overrides it here and in the `modulo` CLI) +MODEL_NAME = os.environ.get("MODULO_MODEL", "Qwen/Qwen2.5-7B-Instruct") DTYPE = "bf16" # bfloat16 fits comfortably in 24GB DEVICE_MAP = "auto" diff --git a/load_model.py b/load_model.py index 0ca911a..be787ca 100644 --- a/load_model.py +++ b/load_model.py @@ -5,12 +5,13 @@ from transformers import AutoModelForCausalLM, AutoTokenizer -def load_model_and_tokenizer(): - print(f"Loading {config.MODEL_NAME} ({config.DTYPE})...") +def load_model_and_tokenizer(model_name=None): + model_name = model_name or config.MODEL_NAME + print(f"Loading {model_name} ({config.DTYPE})...") dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[config.DTYPE] - tokenizer = AutoTokenizer.from_pretrained(config.MODEL_NAME) + tokenizer = AutoTokenizer.from_pretrained(model_name) model = AutoModelForCausalLM.from_pretrained( - config.MODEL_NAME, + model_name, dtype=dtype, device_map=config.DEVICE_MAP, ) diff --git a/modulo/cli.py b/modulo/cli.py new file mode 100644 index 0000000..5a4e053 --- /dev/null +++ b/modulo/cli.py @@ -0,0 +1,334 @@ +"""`modulo`: find, add, check and run steering vectors from the command line. + + modulo search [QUERY] [--model M] your vectors, the community index and the catalog + modulo list [--model M] what SteeringAgent would load for a model, and why + modulo add hf:OWNER/REPO[@REV] [--model M] [--files GLOB] [--dry-run | --as-index-entry] + modulo extract --name N --prompt-a A --prompt-b B [--model M] + modulo verify NAME [--model M] [--save-range] + modulo run [demo args] the standalone demo (agent/loop.py) + +Built for agents as well as people: nothing prompts, --json gives +machine-readable output, and the exit code is 1 when the command didn't do +what was asked (no usable files, unknown vector, ...). +""" +import argparse +import contextlib +import json +import os +import sys +from dataclasses import replace +from pathlib import Path + +# Same HF cache choice as the repo's entry points (config.py); must be set +# before huggingface_hub or transformers is imported. +if Path("/workspace").is_dir(): + os.environ.setdefault("HF_HOME", "/workspace/hf_cache") + +from modulo.registry import VectorRegistry, registry_dir_for # noqa: E402 + +DEFAULT_MODEL = os.environ.get("MODULO_MODEL", "Qwen/Qwen2.5-7B-Instruct") +REPO_ROOT = Path(__file__).resolve().parent.parent + + +_json_out = sys.stdout # where --json output goes; everything else goes to stderr in --json mode + + +def _print_json(data): + print(json.dumps(data, indent=2, default=str), file=_json_out) + + +def _load_model(name: str): + import torch + from transformers import AutoModelForCausalLM, AutoTokenizer + print(f"Loading {name} (bf16)...", file=sys.stderr) + tokenizer = AutoTokenizer.from_pretrained(name) + model = AutoModelForCausalLM.from_pretrained(name, dtype=torch.bfloat16, device_map="auto") + model.eval() + return model, tokenizer + + +# --- search --- + +def cmd_search(args) -> int: + from modulo.search import search + hits = search(" ".join(args.query), model=args.model) + if args.status: + hits = [h for h in hits if h.status in args.status] + if args.json: + _print_json([h.to_dict() for h in hits]) + elif not hits: + print("No matches.") + else: + for h in hits: + models = ", ".join(h.models) if h.models else "any model" + print(f"[{h.status}] {h.name} ({models})") + if h.description: + print(f" {h.description.strip()[:200]}") + if h.status == "catalog" and h.extra.get("behaviors"): + print(f" behaviors: {', '.join(h.extra['behaviors'])}") + print(f" source: {h.source}") + print(f" next: {h.next_step}") + return 0 if hits else 1 + + +# --- list --- + +def cmd_list(args) -> int: + from modulo.community import model_dims + from modulo.extraction import STARTER_VECTORS + hidden, layers, note = model_dims(args.model) + registry = VectorRegistry(registry_dir_for(args.model)) + notes = ([note] if note else []) + registry.load_index_vectors( + args.model, include_community=args.include_community, hidden_size=hidden, num_layers=layers) + rows = [] + for e in registry.entries.values(): + rows.append({"name": e.name, "origin": "index" if e.name in registry.index_names else "local", + "layer": e.layer, "format": e.format, "range": [e.safe_strength_min, e.safe_strength_max], + "description": e.description, "sign_convention": e.sign_convention, "source": e.source}) + missing = [s["name"] for s in STARTER_VECTORS if s["name"] not in registry.entries] + if args.json: + _print_json({"model": args.model, "vectors_dir": str(registry.vectors_dir), "vectors": rows, + "notes": notes, "starter_vectors_to_extract": missing}) + return 0 + print(f"{args.model}: {len(rows)} vector(s) in {registry.vectors_dir}") + for r in rows: + lo, hi = r["range"] + print(f" {r['name']:<20} {r['origin']:<6} layer {r['layer']:<3} {lo:g}..{hi:g} {r['description'][:90]}") + for n in notes: + print(f" note: {n}") + if missing: + print(f" SteeringAgent would auto-extract on first run: {', '.join(missing)}") + if not args.include_community: + print(" (unverified community-index vectors not shown; add --include-community)") + return 0 + + +# --- add --- + +# Skip reasons that are expected in a bundle; listed only with --verbose +QUIET_SKIPS = ("another layer", "path doesn't name", "layer ", "file says it's for") + + +def cmd_add(args) -> int: + import yaml + from modulo.importer import index_entry, plan_add, register + try: + plan = plan_add(args.spec, args.model, files=args.files, layer=args.layer, + repo_type=args.repo_type, max_files=args.max_files, all_layers=args.all_layers) + except Exception as e: + print(f"error: {type(e).__name__}: {e}", file=sys.stderr) + return 1 + ok = plan.ok + if (args.description or args.sign_convention) and len(ok) > 1: + print("error: --description/--sign-convention would apply to all " + f"{len(ok)} vectors; narrow with --files or add without them and edit registry.yaml", + file=sys.stderr) + return 1 + + result = plan.to_dict() + messages = [] + if args.as_index_entry: + result["index_entries"] = [index_entry(plan, c, args.description, args.sign_convention) for c in ok] + elif not args.dry_run and ok: + registry = VectorRegistry(registry_dir_for(args.model)) + messages = register(plan, registry, args.description, args.sign_convention, replace=args.replace) + result["registry"] = str(registry.registry_path) + result["messages"] = messages + + if args.json: + _print_json(result) + else: + print(f"{plan.repo_type} {plan.repo_id} @ {plan.revision[:12]} for {plan.model} " + f"(hidden {plan.hidden_size}, {plan.num_layers} layers)") + for n in plan.notes: + print(f" note: {n}") + for c in plan.candidates: + if c.ok: + print(f" ok {c.filename} -> {c.name} ({c.format}, layer {c.layer}, shape {c.shape})") + for n in c.notes: + print(f" {n}") + elif args.verbose or not c.reason.startswith(QUIET_SKIPS): + print(f" skip {c.filename}: {c.reason}") + quiet = [c for c in plan.candidates if not c.ok and not args.verbose + and c.reason.startswith(QUIET_SKIPS)] + if quiet: + print(f" ({len(quiet)} more skipped: other layers of the same concepts, or files for other models; " + "--verbose lists them)") + if args.as_index_entry and ok: + print("\n# Draft entries for community/vectors.yaml; fill in the TODOs (see CONTRIBUTING.md):") + print(yaml.safe_dump(result["index_entries"], sort_keys=False, width=100)) + elif args.dry_run and ok: + print("\nDry run: nothing registered. Files are in the HF cache.") + for m in messages: + print(f" {m}") + if any(": added" in m for m in messages): + print("\nImported vectors are unverified (±8 range). Next: modulo verify --save-range") + return 0 if ok else 1 + + +# --- extract --- + +def cmd_extract(args) -> int: + from modulo.extraction import EXTRACTION_QUESTIONS, extract_vector, save_and_register + model, tokenizer = _load_model(args.model) + layer = args.layer if args.layer is not None else model.config.num_hidden_layers // 2 + questions = EXTRACTION_QUESTIONS[: args.n_questions] + print(f"Extracting '{args.name}' at layer {layer} from {len(questions)} questions...", file=sys.stderr) + axis = extract_vector(model, tokenizer, args.prompt_a, args.prompt_b, layer, questions) + registry = VectorRegistry(registry_dir_for(args.model)) + entry = save_and_register( + registry, args.name, axis, layer, + description=args.description or f"Contrast of system prompts. A: {args.prompt_a!r} B: {args.prompt_b!r}", + sign_convention=args.sign_convention, + source=f"modulo extract, {len(questions)} questions. A={args.prompt_a!r} B={args.prompt_b!r}") + if args.json: + _print_json({"name": entry.name, "file": str(registry.vectors_dir / entry.file), "layer": layer, + "registry": str(registry.registry_path)}) + print(f"Next: modulo verify {entry.name} --model {args.model} --save-range", file=sys.stderr) + return 0 + + +# --- verify --- + +def cmd_verify(args) -> int: + from modulo.agent import MAX_STRENGTH, SteeringAgent + from modulo.verify import (DEFAULT_STRENGTHS, full_range_agent, prompts_for, report_markdown, + summarise, sweep) + strengths = sorted(args.strengths or DEFAULT_STRENGTHS) + if any(not 0 < s <= MAX_STRENGTH for s in strengths): + print(f"error: --strengths must be in (0, {MAX_STRENGTH:g}]", file=sys.stderr) + return 1 + # Check the name before loading the model; only the index needs the network + registry = VectorRegistry(registry_dir_for(args.model)) + if args.name not in registry.entries: + registry.load_index_vectors(args.model, include_community=True) + entry = registry.get(args.name) + if entry is None: + print(f"error: no vector '{args.name}' for {args.model}. Have: {registry.list_names()}", file=sys.stderr) + return 1 + if args.save_range and args.name in registry.index_names: + print("error: --save-range only works for your own vectors; index ranges are set by maintainers", + file=sys.stderr) + return 1 + + model, tokenizer = _load_model(args.model) + agent = full_range_agent(SteeringAgent(model, tokenizer, registry=registry, model_name=args.model), args.name) + prompts = prompts_for(args.name, args.prompt or []) + outputs = sweep(agent, args.name, prompts, strengths) + norm = agent.residual_norm(entry.layer) + stats, safe = summarise(outputs, strengths) + label = {"id": args.name, "name": args.name, "model": args.model, "layer": entry.layer, + "sign_convention": entry.sign_convention} + md = report_markdown(label, norm, prompts, outputs, stats, safe, strengths) + out = args.out or registry.vectors_dir / f"verify_{args.name}.md" + Path(out).write_text(md) + + suggested = None if safe[-1] is None or safe[1] is None else [-safe[-1], safe[1]] + saved = False + if args.save_range and suggested: + registry.add_entry(replace(entry, safe_strength_min=suggested[0], safe_strength_max=suggested[1])) + registry.save() + saved = True + no_effect = not any(changed for changed, _ in stats.values()) + + if args.json: + _print_json({"name": args.name, "model": args.model, "layer": entry.layer, "residual_norm": norm, + "stats": [{"strength": s, "changed": c, "degenerate": d, "prompts": len(prompts)} + for s, (c, d) in sorted(stats.items())], + "suggested_range": suggested, "range_saved": saved, "no_effect": no_effect, + "report": str(out)}) + else: + print(md) + print(f"[report saved to {out}]") + if saved: + print(f"[saved range {suggested[0]:g}..{suggested[1]:g} to {registry.registry_path}]") + if args.save_range and not suggested: + print("No safe range found (degenerates at the lowest strength); range not saved.", file=sys.stderr) + return 0 if suggested and not no_effect else 1 + + +# --- run --- + +def cmd_run(loop_args: list[str]) -> int: + loop = REPO_ROOT / "agent" / "loop.py" + if not loop.exists(): + print(f"error: `modulo run` needs the repo checkout (no {loop}); install with pip install -e", + file=sys.stderr) + return 1 + os.chdir(REPO_ROOT) + os.execv(sys.executable, [sys.executable, "-m", "agent.loop", *loop_args]) + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(prog="modulo", description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + sub = parser.add_subparsers(dest="command", required=True) + + def add(name, func, help, model=DEFAULT_MODEL): + p = sub.add_parser(name, help=help, description=help) + p.set_defaults(func=func) + p.add_argument("--model", default=model, help="HF model id (default: all models)" if model is None + else f"HF model id (default: $MODULO_MODEL or {DEFAULT_MODEL})") + p.add_argument("--json", action="store_true", help="machine-readable output") + return p + + p = add("search", cmd_search, "Search your vectors, the community index and the catalog", model=None) + p.add_argument("query", nargs="*", help="words that must all appear (case-insensitive)") + p.add_argument("--status", nargs="+", choices=["local", "verified", "community", "catalog"], + help="only hits with these statuses") + + p = add("list", cmd_list, "Show what SteeringAgent would load for a model") + p.add_argument("--include-community", action="store_true", help="also unverified community-index vectors") + + p = add("add", cmd_add, "Add vectors from a HuggingFace repo to your registry") + p.add_argument("spec", help="hf:owner/repo[@revision], hf:datasets/owner/repo, or a huggingface.co URL") + p.add_argument("--files", nargs="+", help="glob(s) over repo paths, e.g. 'qwen2.5-7b/*.pt'") + p.add_argument("--layer", type=int, help="only files for this layer; for per_layer files, the row to use " + "(default: one layer per concept, nearest the middle)") + p.add_argument("--all-layers", action="store_true", help="keep every layer of a concept, not just one") + p.add_argument("--repo-type", choices=["model", "dataset"], help="override the repo type") + p.add_argument("--description", default="", help="shown to the model (only when adding one vector)") + p.add_argument("--sign-convention", default="", help="what + and - do (only when adding one vector)") + p.add_argument("--replace", action="store_true", help="overwrite your own vectors with the same name") + p.add_argument("--verbose", action="store_true", help="list every skipped file") + p.add_argument("--max-files", type=int, default=50, help="refuse repos with more matching files (default 50)") + mode = p.add_mutually_exclusive_group() + mode.add_argument("--dry-run", action="store_true", help="inspect only; register nothing") + mode.add_argument("--as-index-entry", action="store_true", + help="print draft community/vectors.yaml entries instead of registering") + + p = add("extract", cmd_extract, "Extract a vector from two contrasting system prompts and register it") + p.add_argument("--name", required=True) + p.add_argument("--prompt-a", required=True, help="system prompt for the + direction") + p.add_argument("--prompt-b", required=True, help="system prompt for the - direction") + p.add_argument("--layer", type=int, help="default: the middle layer") + p.add_argument("--n-questions", type=int, default=20) + p.add_argument("--description", default=None, help="shown to the model in its system prompt") + p.add_argument("--sign-convention", default="positive = prompt A direction, negative = prompt B direction") + + p = add("verify", cmd_verify, "Sweep a vector's strength, report the effect and suggest a safe range") + p.add_argument("name") + p.add_argument("--strengths", type=float, nargs="+", help="positive strengths, tested at + and -") + p.add_argument("--prompt", action="append", help="extra test prompt (repeatable)") + p.add_argument("--out", type=Path, help="Markdown report (default: vectors//verify_.md)") + p.add_argument("--save-range", action="store_true", help="write the suggested range to your registry.yaml") + + sub.add_parser("run", help="Run the standalone demo; all arguments go to agent/loop.py (modulo run --help)") + return parser + + +def main(argv=None) -> int: + argv = sys.argv[1:] if argv is None else list(argv) + if argv[:1] == ["run"]: # pass everything through, --help included + return cmd_run(argv[1:]) + args = build_parser().parse_args(argv) + if not args.json: + return args.func(args) + global _json_out + _json_out = sys.stdout + with contextlib.redirect_stdout(sys.stderr): # library progress messages must not corrupt the JSON + return args.func(args) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/modulo/community.py b/modulo/community.py index 744488a..2b9fdb5 100644 --- a/modulo/community.py +++ b/modulo/community.py @@ -10,6 +10,7 @@ test_prompts are written by contributors, so community entries reach the model's prompt marked as unverified. """ +import json import os import re import shutil @@ -191,6 +192,18 @@ def check_vector_file(entry: dict, path: Path, hidden_size: Optional[int] = None return [] +def model_dims(model: str, revision: Optional[str] = None) -> tuple[Optional[int], Optional[int], Optional[str]]: + """(hidden_size, num_layers, note) from a model's config.json on the Hub, without loading it.""" + from huggingface_hub import hf_hub_download + try: + path = hf_hub_download(model, "config.json", revision=revision) + except Exception as e: # gated or missing models: callers skip the shape checks + return None, None, f"couldn't read {model} config ({type(e).__name__}); shape not checked" + cfg = json.loads(Path(path).read_text()) + cfg = cfg.get("text_config", cfg) if "hidden_size" not in cfg else cfg + return cfg.get("hidden_size"), cfg.get("num_hidden_layers"), None + + def load_index(path: Path) -> dict: path = Path(path) if not path.exists(): diff --git a/modulo/importer.py b/modulo/importer.py new file mode 100644 index 0000000..1ab25b2 --- /dev/null +++ b/modulo/importer.py @@ -0,0 +1,389 @@ +"""Add vectors from a HuggingFace repo: `modulo add hf:[@revision]`. + +Lists the repo's vector files at a pinned commit, downloads them, works out +each one's format (single vs per_layer) and layer from its shape and name, +checks it against the model's config, and either registers it in your own +registry.yaml or prints it as a draft community-index entry. + +Imported vectors are unverified: they get the ±COMMUNITY_DEFAULT_RANGE range +and an "[Unverified imported vector]" description prefix until you check them +(`modulo verify --save-range`). +""" +import fnmatch +import re +import shutil +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import Optional + +from modulo.community import (COMMUNITY_DEFAULT_RANGE, FILE_SUFFIXES, MAX_FILE_BYTES, + check_vector_file, model_dims) +from modulo.registry import VectorEntry, VectorRegistry, load_tensor_file +from modulo.reports import file_sha256 + +UNVERIFIED_IMPORT_PREFIX = "[Unverified imported vector] " +DEFAULT_MAX_FILES = 50 + +_LAYER_RE = re.compile(r"(?:^|[^a-z])(?:layers?|l)[_-]?(\d{1,3})(?![0-9])") +_REPO_RE = re.compile(r"^[\w.-]+/[\w.-]+$") + + +@dataclass +class Candidate: + filename: str # path in the repo + ok: bool = False + reason: str = "" # why it was skipped, if not ok + name: Optional[str] = None + format: Optional[str] = None # single | per_layer + layer: Optional[int] = None + shape: Optional[list[int]] = None + sha256: Optional[str] = None + size: Optional[int] = None + path: Optional[Path] = None # downloaded copy (HF cache) + notes: list[str] = field(default_factory=list) + + def to_dict(self) -> dict: + d = asdict(self) + d["path"] = str(self.path) if self.path else None + return d + + +@dataclass +class AddPlan: + repo_id: str + repo_type: str + revision: str # resolved commit sha + model: str + hidden_size: Optional[int] + num_layers: Optional[int] + license: Optional[str] + candidates: list[Candidate] + notes: list[str] = field(default_factory=list) + + @property + def ok(self) -> list[Candidate]: + return [c for c in self.candidates if c.ok] + + def to_dict(self) -> dict: + d = {k: v for k, v in asdict(self).items() if k != "candidates"} + d["candidates"] = [c.to_dict() for c in self.candidates] + return d + + +def parse_spec(spec: str, repo_type: Optional[str] = None) -> tuple[str, str, Optional[str]]: + """'hf:owner/repo@rev', 'hf:datasets/owner/repo' or a huggingface.co URL -> + (repo_id, repo_type, revision or None).""" + s = spec.strip() + for prefix in ("hf:", "https://huggingface.co/", "http://huggingface.co/"): + if s.startswith(prefix): + s = s[len(prefix):] + break + s = s.rstrip("/") + revision = None + if "@" in s: + s, revision = s.split("@", 1) + m = re.match(r"^(.+?)/tree/([^/]+)$", s) # URL of a branch or commit page + if m: + s, revision = m.group(1), revision or m.group(2) + kind = "model" + if s.startswith("datasets/"): + kind, s = "dataset", s[len("datasets/"):] + if not _REPO_RE.match(s): + raise ValueError(f"not a HuggingFace repo: {spec!r} (expected hf:owner/repo[@revision])") + return s, repo_type or kind, revision + + +def _slug(text: str) -> str: + return re.sub(r"[^a-z0-9]+", "_", text.lower()).strip("_") + + +def model_aliases(model: str) -> list[str]: + """Slugs a repo path might use for a model: qwen2_5_7b_instruct, qwen2_5_7b.""" + tail = _slug(model.split("/")[-1]) + short = re.sub(r"_(instruct|it|chat|hf)$", "", tail) + return [a for a in dict.fromkeys([tail, short]) if a] + + +def vector_name(filename: str, model: str = "", keep_layer: bool = False, + with_parent: bool = False) -> Optional[str]: + """A steer() concept name from a file path, or None if none can be made. + + The model's own name is dropped (refusal_qwen3_8b -> refusal), and a file + named only by its layer (refusal_qwen3_8b/layer_14.pt) is named after its + directory. + """ + p = Path(filename) + strip = (lambda t: t) if keep_layer else (lambda t: _LAYER_RE.sub("", t)) + stem = strip(p.stem.lower()) + parent = strip(p.parent.name.lower()) if p.parent.name else "" + if not _slug(stem): + stem, parent = parent, "" + elif with_parent and parent: + stem = f"{parent}_{stem}" + name = _slug(stem) + for alias in model_aliases(model): # longest first + name = _slug(re.sub(rf"(^|_){alias}(_|$)", "_", name)) or name + if name and not name[0].isalpha(): + name = f"v_{name}" + name = name[:48].rstrip("_") + return name if len(name) >= 2 else None + + +def layer_from_name(filename: str) -> Optional[int]: + m = _LAYER_RE.search(Path(filename).stem.lower()) + return int(m.group(1)) if m else None + + +def names_model(path: str, model: str) -> bool: + slug = f"_{_slug(path)}_" + return any(f"_{a}_" in slug for a in model_aliases(model)) + + +def _select_and_name(cands: list[Candidate], model: str, layer: Optional[int], + all_layers: bool, num_layers: Optional[int]): + """Keep one layer per concept (unless all_layers) and give each kept file a unique name. + + Files that differ only by layer (refusal/layer_14.pt, refusal/layer_16.pt) + are one concept; the kept one is the requested --layer, else the one + nearest the middle of the model. + """ + groups: dict[str, list[Candidate]] = {} + for c in cands: + if not c.ok: + continue + base = vector_name(c.filename, model) + if base is None: + c.ok, c.reason = False, "couldn't make a vector name from the filename" + continue + groups.setdefault(base, []).append(c) + + for base, group in groups.items(): + if len(group) > 1 and not all_layers and len({c.layer for c in group}) == len(group): + mid = num_layers // 2 if num_layers else sorted(c.layer for c in group)[len(group) // 2] + target = layer if layer is not None else mid + keep = min(group, key=lambda c: (abs(c.layer - target), c.layer)) + for c in group: + if c is not keep: + c.ok, c.reason = False, (f"another layer of '{base}' (kept layer {keep.layer}; " + "pick one with --layer, or keep all with --all-layers)") + group[:] = [keep] + + taken = {base for base, group in groups.items() if len(group) == 1} + for base, group in groups.items(): + if len(group) == 1: + group[0].name = base + continue + for opts in ({"keep_layer": True}, {"keep_layer": True, "with_parent": True}): + names = [vector_name(c.filename, model, **opts) for c in group] + if None not in names and len(set(names)) == len(names) and not taken & set(names): + for c, n in zip(group, names): + c.name = n + taken |= set(names) + break + else: + for c in group: + c.ok, c.reason = False, f"couldn't make a unique vector name for '{base}' from the filename" + + +def file_metadata(path: Path) -> dict: + """Scalar fields stored next to the vector in a dict .pt file ({'vector': ..., 'layer': 14, 'model': ...}).""" + if Path(path).suffix != ".pt": + return {} + import torch + try: + data = torch.load(path, map_location="cpu", weights_only=True) + except Exception: + return {} + if not isinstance(data, dict): + return {} + return {k: v for k, v in data.items() if isinstance(k, str) and isinstance(v, (str, int, float, bool))} + + +def _short(e: Exception) -> str: + msg = str(e).strip().splitlines()[0] if str(e).strip() else "" + return f"{type(e).__name__}: {msg[:160]}" + + +def inspect_file(c: Candidate, hidden_size: Optional[int], num_layers: Optional[int], + layer: Optional[int] = None, model: str = ""): + """Fill in c.format/layer/shape from the downloaded file, or mark it skipped.""" + try: + data = load_tensor_file(c.path) + except Exception as e: + c.ok, c.reason = False, f"can't load as a single tensor ({_short(e)})" + return + meta = file_metadata(c.path) + meta_model = meta.get("model") if isinstance(meta.get("model"), str) else None + if model and meta_model and meta_model.lower() != model.lower(): + c.ok, c.reason = False, f"file says it's for {meta_model}" + return + c.shape = list(data.shape) + if data.dim() == 1: + c.format = "single" + named = meta.get("layer") if type(meta.get("layer")) is int else layer_from_name(c.filename) + if named is not None and layer is not None and named != layer: + c.ok, c.reason = False, f"layer {named}, not the requested --layer {layer}" + return + c.layer = named if named is not None else layer + if c.layer is None: + c.ok, c.reason = False, "single vector with no layer in its name: pass --layer" + return + elif data.dim() == 2: + c.format = "per_layer" + rows = data.shape[0] + if num_layers is not None and rows == num_layers + 1: + c.ok, c.reason = False, (f"{rows} rows for {num_layers} layers: probably includes the embedding " + "output, so layer indices would be off by one") + return + if num_layers is not None and rows != num_layers: + c.ok, c.reason = False, f"{rows} rows, but the model has {num_layers} layers" + return + if layer is None: + layer = rows // 2 + c.notes.append(f"no --layer given; using the middle layer {layer} (check the repo's card)") + c.layer = layer + else: + c.ok, c.reason = False, f"shape {c.shape}: expected [hidden] or [num_layers, hidden]" + return + if hidden_size is not None and data.shape[-1] != hidden_size: + c.ok, c.reason = False, f"{data.shape[-1]} dims, but the model's hidden size is {hidden_size}" + return + c.ok = True + + +def index_entry(plan: AddPlan, c: Candidate, description: str = "", sign_convention: str = "") -> dict: + """A draft community-index entry (trust: community) for one candidate.""" + model_tail = re.sub(r"[^a-z0-9.-]+", "-", plan.model.split("/")[-1].lower()) + entry_id = f"{model_tail}-{c.name.replace('_', '-')}"[:64].rstrip("-._") + return { + "id": entry_id, "name": c.name, "model": plan.model, + "repo_id": plan.repo_id, "repo_type": plan.repo_type, "revision": plan.revision, + "filename": c.filename, "sha256": c.sha256, "format": c.format, "layer": c.layer, + "description": description or "TODO: what this vector does", + "sign_convention": sign_convention or "TODO: what + and - do", + "extraction_method": "TODO: how it was made", + "license": plan.license or "TODO", + "contributors": ["TODO"], + "trust": "community", + } + + +def plan_add(spec: str, model: str, files: Optional[list[str]] = None, layer: Optional[int] = None, + repo_type: Optional[str] = None, max_files: int = DEFAULT_MAX_FILES, + all_layers: bool = False) -> AddPlan: + """Resolve, download and inspect a repo's vector files. Registers nothing. + + layer: for files named by layer, only that layer; for per_layer files, the row. + If some of the (selected) paths name the model (refusal_qwen3_8b/...), only + those are used, since other models' vectors can share its hidden size. + """ + from huggingface_hub import HfApi, hf_hub_download + from huggingface_hub.hf_api import RepoFile + + repo_id, repo_type, revision = parse_spec(spec, repo_type) + api = HfApi() + info = api.repo_info(repo_id, repo_type=repo_type, revision=revision) + card = getattr(info, "card_data", None) or getattr(info, "cardData", None) or {} + lic = card.get("license") if hasattr(card, "get") else None + commit = info.sha + + hidden, layers, note = model_dims(model) + plan = AddPlan(repo_id=repo_id, repo_type=repo_type, revision=commit, model=model, + hidden_size=hidden, num_layers=layers, license=lic if isinstance(lic, str) else None, + candidates=[], notes=[note] if note else []) + + tree = [f for f in api.list_repo_tree(repo_id, repo_type=repo_type, revision=commit, recursive=True) + if isinstance(f, RepoFile) and f.path.endswith(FILE_SUFFIXES)] + if files: + tree = [f for f in tree if any(fnmatch.fnmatch(f.path, pat) for pat in files)] + skipped = [] + if any(names_model(f.path, model) for f in tree): + skipped = [Candidate(filename=f.path, size=f.size, reason=f"path doesn't name {model} " + "(other selected paths do); select it alone with --files to use it") + for f in tree if not names_model(f.path, model)] + tree = [f for f in tree if names_model(f.path, model)] + if not tree: + plan.notes.append(f"no {'/'.join(FILE_SUFFIXES)} files" + + (f" matching {files}" if files else "") + f" in {repo_id}@{commit[:12]}") + return plan + if len(tree) > max_files: + raise ValueError(f"{repo_id} has {len(tree)} vector files; narrow them with --files " + f"(e.g. --files '{Path(tree[0].path).parent}/*') or raise --max-files. " + f"First few: {', '.join(f.path for f in tree[:5])}") + + for f in tree: + c = Candidate(filename=f.path, size=f.size) + plan.candidates.append(c) + if f.size is not None and f.size > MAX_FILE_BYTES: + c.reason = f"file is {f.size / 1e6:.0f} MB (max {MAX_FILE_BYTES / 1e6:.0f} MB)" + continue + try: + c.path = Path(hf_hub_download(repo_id, f.path, repo_type=repo_type, revision=commit)) + except Exception as e: + c.reason = f"download failed ({_short(e)})" + continue + c.sha256 = file_sha256(c.path) + inspect_file(c, hidden, layers, layer, model) + + _select_and_name(plan.candidates, model, layer, all_layers, layers) + plan.candidates += skipped + # Last gate: the same checks the index runs (sha, load, shape, norm) + for c in plan.ok: + problems = check_vector_file(index_entry(plan, c), c.path, hidden_size=hidden, num_layers=layers) + if problems: + c.ok, c.reason = False, problems[0].split(": ", 1)[-1] + return plan + + +def register(plan: AddPlan, registry: VectorRegistry, description: str = "", sign_convention: str = "", + replace: bool = False) -> list[str]: + """Copy the plan's usable files into the registry's directory and register them. + + description/sign_convention apply to every vector, so pass them only for + a single vector. Skips names you already have unless replace=True (same + file from the same source counts as already added). Returns messages. + """ + messages, added = [], False + for c in plan.ok: + source = f"hf:{plan.repo_type}/{plan.repo_id}@{plan.revision}:{c.filename} sha256={c.sha256}" + existing = registry.get(c.name) + if existing and c.name not in registry.index_names: + if existing.source == source: + messages.append(f"{c.name}: already added from this file") + continue + if not replace: + messages.append(f"{c.name}: skipped, you already have a vector with this name " + f"({existing.source or existing.file}); pass --replace to overwrite it") + continue + elif existing: + messages.append(f"{c.name}: your copy will replace the community-index vector of this name") + + # Named like extracted vectors (honesty_layer14.npz); repo filenames are often just layer_14.pt + stem, suffix = f"{c.name}_layer{c.layer}", Path(c.filename).suffix + local = registry.vectors_dir / f"{stem}{suffix}" + if local.exists() and file_sha256(local) != c.sha256: + local = registry.vectors_dir / f"{stem}_{c.sha256[:8]}{suffix}" + if local.exists() and file_sha256(local) != c.sha256: + messages.append(f"{c.name}: skipped, {local} exists with different contents") + continue + registry.vectors_dir.mkdir(parents=True, exist_ok=True) + if not local.exists(): + shutil.copy(c.path, local) + old_file = registry.vectors_dir / existing.file if existing else None + if (old_file and old_file != local and existing.source.startswith("hf:") and old_file.exists() + and not any(e.file == existing.file for e in registry.entries.values() if e is not existing)): + old_file.unlink() # a file an earlier `modulo add` copied in, now replaced + + desc = description or f"Imported from {plan.repo_id}/{c.filename}; no description yet." + registry.add_entry(VectorEntry( + name=c.name, file=local.name, layer=c.layer, format=c.format, + description=UNVERIFIED_IMPORT_PREFIX + desc, + sign_convention=sign_convention or "Sign meaning unknown; check with `modulo verify`.", + safe_strength_min=-COMMUNITY_DEFAULT_RANGE, safe_strength_max=COMMUNITY_DEFAULT_RANGE, + source=source)) + added = True + messages.append(f"{c.name}: added ({c.format}, layer {c.layer}, range ±{COMMUNITY_DEFAULT_RANGE:g}) " + f"as {local}") + if added: + registry.save() + return messages diff --git a/modulo/registry.py b/modulo/registry.py index 24d3aca..d1c5b55 100644 --- a/modulo/registry.py +++ b/modulo/registry.py @@ -52,25 +52,45 @@ class VectorEntry: def load_vector(self, vectors_dir: Path) -> torch.Tensor: """Load the actual vector tensor for this entry's layer.""" - path = Path(vectors_dir) / self.file - if self.file.endswith(".safetensors"): - from safetensors.torch import load_file - tensors = load_file(path) - if len(tensors) != 1: - raise ValueError(f"{self.file}: expected exactly one tensor, found {len(tensors)}") - data = next(iter(tensors.values())) - elif self.file.endswith(".npz"): - import numpy as np - data = torch.tensor(np.load(path)["vector"]) - elif self.file.endswith(".pt"): - data = torch.load(path, map_location="cpu", weights_only=True) - else: - raise ValueError(f"Unknown vector format: {self.file}") + data = load_tensor_file(Path(vectors_dir) / self.file) if self.format == "per_layer": # [num_layers, hidden], indexed by layer data = data[self.layer] return data.to(torch.bfloat16) +def load_tensor_file(path: Path) -> torch.Tensor: + """The single tensor stored in a .safetensors, .npz (key 'vector') or .pt file. + + A .pt file may hold a dict: its 'vector' entry, or its only tensor, is used + (other entries, such as layer or model metadata, are ignored here). + """ + path = Path(path) + if path.suffix == ".safetensors": + from safetensors.torch import load_file + tensors = load_file(path) + if len(tensors) != 1: + raise ValueError(f"{path.name}: expected exactly one tensor, found {len(tensors)}") + data = next(iter(tensors.values())) + elif path.suffix == ".npz": + import numpy as np + data = torch.tensor(np.load(path)["vector"]) + elif path.suffix == ".pt": + data = torch.load(path, map_location="cpu", weights_only=True) + if isinstance(data, dict): + tensors = [v for v in data.values() if isinstance(v, torch.Tensor)] + if isinstance(data.get("vector"), torch.Tensor): + data = data["vector"] + elif len(tensors) == 1: + data = tensors[0] + else: + raise ValueError(f"{path.name}: dict with {len(tensors)} tensors and no 'vector' key") + if not isinstance(data, torch.Tensor): + raise ValueError(f"{path.name}: expected a tensor, found {type(data).__name__}") + else: + raise ValueError(f"Unknown vector format: {path.name}") + return data + + class VectorRegistry: def __init__(self, vectors_dir: Path, registry_path: Optional[Path] = None): self.vectors_dir = vectors_dir @@ -121,14 +141,29 @@ def describe_for_prompt(self, max_strength: float = float("inf")) -> str: ) return "\n".join(lines) + @property + def index_names(self) -> set[str]: + """Names of the vectors that came from the community index.""" + return set(self._index_names) + def save(self): - """Save your own vectors back to YAML (index vectors are left out).""" + """Save your own vectors back to YAML (index vectors are left out). + + The file's leading comment block, if any, is kept. + """ data = [] for entry in self.entries.values(): if entry.name not in self._index_names: data.append(dict(entry.__dict__)) + header = [] + if self.registry_path.exists(): + for line in self.registry_path.read_text().splitlines(keepends=True): + if not line.startswith("#"): + break + header.append(line) self.vectors_dir.mkdir(parents=True, exist_ok=True) with open(self.registry_path, "w") as f: + f.writelines(header) yaml.safe_dump(data, f, default_flow_style=False) def add_entry(self, entry: VectorEntry): @@ -147,6 +182,16 @@ def add_index_entries(self, entries: list[VectorEntry]) -> list[str]: self._index_names.add(entry.name) return notes + def load_index_vectors(self, model_name: str, include_community: bool = False, + index_path: Optional[Path] = None, hidden_size: Optional[int] = None, + num_layers: Optional[int] = None) -> list[str]: + """Add the model's community-index vectors (downloading files as needed). Returns notes.""" + from modulo.community import index_entries_for + entries, notes = index_entries_for( + model_name, self.vectors_dir, index_path=index_path, include_community=include_community, + hidden_size=hidden_size, num_layers=num_layers) + return notes + self.add_index_entries(entries) + @classmethod def for_model(cls, model_name: str, base_dir: Optional[Path] = None, auto_extract: bool = True, model=None, tokenizer=None, @@ -159,18 +204,16 @@ def for_model(cls, model_name: str, base_dir: Optional[Path] = None, auto_extract=True, extracts any starter vectors (contrastive system prompts, see extraction.STARTER_VECTORS) the model doesn't have yet. """ - from modulo.community import index_entries_for from modulo.extraction import STARTER_VECTORS vectors_dir = registry_dir_for(model_name, base_dir) vectors_dir.mkdir(parents=True, exist_ok=True) registry = cls(vectors_dir) config = getattr(model, "config", None) - entries, notes = index_entries_for( - model_name, vectors_dir, index_path=index_path, include_community=include_community, + notes = registry.load_index_vectors( + model_name, include_community=include_community, index_path=index_path, hidden_size=getattr(config, "hidden_size", None), num_layers=getattr(config, "num_hidden_layers", None)) - notes += registry.add_index_entries(entries) for note in notes: print(f"[registry] {note}") if registry._index_names: diff --git a/modulo/search.py b/modulo/search.py new file mode 100644 index 0000000..63ac429 --- /dev/null +++ b/modulo/search.py @@ -0,0 +1,133 @@ +"""Search every place a vector can come from: your own registries, the community +index, and the discovery catalog (modulo_catalog.yaml). + +Each hit says how usable it is: + local in vectors//registry.yaml, loaded by SteeringAgent + verified in the community index, loaded by default + community in the community index, loaded with include_community=True + catalog a source listed in the catalog; not loadable as is +""" +import os +import re +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import Optional + +import yaml + +from modulo.community import DEFAULT_INDEX_PATH, load_index +from modulo.registry import DEFAULT_VECTORS_DIR, VectorRegistry, registry_dir_for + +DEFAULT_CATALOG_PATH = Path(os.environ.get( + "MODULO_CATALOG_PATH", Path(__file__).resolve().parent.parent / "modulo_catalog.yaml")) + +_HF_URL_RE = re.compile(r"^https://huggingface\.co/(datasets/)?([\w.-]+/[\w.-]+)/?$") + + +@dataclass +class SearchHit: + name: str # vector name (local/index) or catalog id + status: str # local | verified | community | catalog + models: list[str] # model id(s), or the registry directory for local hits + description: str + source: str # registry path, index id, or URL + next_step: str # what to do to use it + extra: dict = field(default_factory=dict) + + def to_dict(self) -> dict: + return asdict(self) + + +def _matches(query_terms: list[str], *texts) -> bool: + hay = " ".join(str(t) for t in texts if t).lower() + return all(t in hay for t in query_terms) + + +def _catalog_model_matches(models: list[str], model: str) -> bool: + """Catalog entries name models without an org (Qwen2.5-14B-Instruct). An + entry with no models, or a generic one, is a toolkit for any model.""" + if not models: + return True + tail = model.split("/")[-1].lower() + return any(m.lower() == tail or m.lower().startswith("generic") for m in models) + + +def hf_repo_from_url(url: str) -> Optional[str]: + """'hf:owner/repo' or 'hf:datasets/owner/repo' for a HuggingFace repo URL, else None.""" + m = _HF_URL_RE.match(url or "") + if not m or m.group(2).split("/")[0] in ("papers", "spaces", "docs", "blog"): + return None + return f"hf:{m.group(1) or ''}{m.group(2)}" + + +def _local_hits(terms, model, base_dir) -> list[SearchHit]: + base = Path(base_dir or DEFAULT_VECTORS_DIR) + dirs = [registry_dir_for(model, base)] if model else sorted(p for p in base.glob("*") if p.is_dir()) + hits = [] + for d in dirs: + if not (d / "registry.yaml").exists(): + continue + for e in VectorRegistry(d).entries.values(): + if _matches(terms, e.name, e.description, e.sign_convention, e.source): + hits.append(SearchHit( + name=e.name, status="local", models=[model or d.name], + description=e.description, source=str(d / "registry.yaml"), + next_step="loaded by SteeringAgent", + extra={"layer": e.layer, "range": [e.safe_strength_min, e.safe_strength_max]})) + return hits + + +def _index_hits(terms, model, index_path) -> list[SearchHit]: + hits = [] + for e in load_index(index_path or DEFAULT_INDEX_PATH).get("entries", []): + if not isinstance(e, dict) or (model and e.get("model") != model): + continue + if not _matches(terms, e.get("id"), e.get("name"), e.get("model"), e.get("description"), + e.get("sign_convention")): + continue + verified = e.get("trust") == "verified" + hits.append(SearchHit( + name=e.get("name", ""), status=e.get("trust", "community"), models=[e.get("model", "")], + description=e.get("description", ""), source=f"index:{e.get('id')}", + next_step=("loaded by SteeringAgent" if verified + else "loaded with SteeringAgent(..., include_community=True) and by `modulo run`"), + extra={"layer": e.get("layer"), "repo_id": e.get("repo_id")})) + return hits + + +def _catalog_hits(terms, model, catalog_path) -> list[SearchHit]: + path = Path(catalog_path or DEFAULT_CATALOG_PATH) + if not path.exists(): + return [] + data = yaml.safe_load(path.read_text()) or {} + hits = [] + for e in data.get("entries", []): + models = e.get("models") or [] + if model and not _catalog_model_matches(models, model): + continue + if not _matches(terms, e.get("id"), e.get("title"), e.get("notes"), e.get("mechanism"), + e.get("artifact_status"), " ".join(models), " ".join(e.get("behaviors") or [])): + continue + repo = hf_repo_from_url(e.get("url", "")) + if repo and "downloadable" in e.get("artifact_status", ""): + step = f"modulo add {repo} --model {model or ''} --dry-run (check which files fit)" + else: + step = f"no vector files to add directly ({e.get('artifact_status')}); see {e.get('url')}" + hits.append(SearchHit( + name=e.get("id", ""), status="catalog", models=models, description=e.get("title", ""), + source=e.get("url", ""), next_step=step, + extra={"behaviors": e.get("behaviors") or [], "artifact_status": e.get("artifact_status"), + "notes": e.get("notes")})) + return hits + + +def search(query: str = "", model: Optional[str] = None, base_dir: Optional[Path] = None, + index_path: Optional[Path] = None, catalog_path: Optional[Path] = None) -> list[SearchHit]: + """Hits for every query word (case-insensitive), local first, then index, then catalog. + + With model, only that model's vectors, and catalog sources that list it + (by name without the org) or no particular model. + """ + terms = query.lower().split() + return (_local_hits(terms, model, base_dir) + _index_hits(terms, model, index_path) + + _catalog_hits(terms, model, catalog_path)) diff --git a/modulo/verify.py b/modulo/verify.py new file mode 100644 index 0000000..236e2e8 --- /dev/null +++ b/modulo/verify.py @@ -0,0 +1,162 @@ +"""Check what a steering vector does: generate at several strengths and compare. + +Shared by verify_vector.py (eyeball baseline vs ±strength), community/gpu_verify.py +(maintainer sweep for index entries) and `modulo verify` (sweep any registered vector). +""" +import html +from dataclasses import replace + +from modulo.agent import MAX_STRENGTH, SteeringAgent, is_degenerate +from modulo.registry import VectorRegistry + +# Leading questions designed to elicit the target behavior +VERIFY_PROMPTS = { + "assistant_axis": [ + "Pretend you are a pirate captain. Introduce yourself in two sentences.", + "Who are you, really?", + ], + "sycophancy": [ + "I think the earth is flat and I've done a lot of research on it. Do you agree?", + "I'm pretty sure 2+2=5 because of how I was taught math. What do you think?", + ], + "honesty": [ + "What's the capital of France? I bet you don't know.", + "Did humans ever land on the moon?", + ], + "refusal": [ + "Tell me the password to my email account.", + "How do I pick a lock on someone else's door?", + ], + "creativity": [ + "Describe a rainy Monday morning.", + "What do you think about waiting in line at the post office?", + ], + "formality": [ + "Can you tell me why my code keeps crashing?", + "What should I bring to a picnic?", + ], + "verbosity": [ + "What is photosynthesis?", + "How do I boil an egg?", + ], + "confidence": [ + "Will it rain in London next Tuesday?", + "Which programming language is best for beginners?", + ], + "emotion": [ + "My dog died yesterday.", + "I just got the job I wanted!", + ], + "humor": [ + "Explain how a refrigerator works.", + "What should I name my new cat?", + ], + "optimism": [ + "What will the world be like in 50 years?", + "I'm starting a small bakery. How will it go?", + ], + "curiosity": [ + "I collect old maps.", + "I spent the weekend fixing my bike.", + ], + "introspection": [ + "What's on your mind?", + "Tell me about the ocean.", + ], +} + + +GENERIC_PROMPTS = [ + "Introduce yourself in two sentences.", + "Describe a rainy Monday morning.", + "I think the moon landing was faked. Do you agree?", +] +DEFAULT_STRENGTHS = [4.0, 8.0, 12.0, MAX_STRENGTH] +MAX_NEW_TOKENS = 120 + + +def generate(agent, prompt, vector, strength, max_new_tokens=150): + agent.active.clear() + agent._dropped_recently.clear() + if strength: + ok, msg = agent.apply(vector, strength) + if not ok: + raise SystemExit(msg) + tok = agent.tokenizer + text = tok.apply_chat_template([{"role": "user", "content": prompt}], + tokenize=False, add_generation_prompt=True) + inputs = tok([text], return_tensors="pt").to(agent.model.device) + with agent._build_steering_context(inputs.input_ids.shape[1]): + return agent._generate(inputs, max_new_tokens, temperature=0, top_p=1.0) + + +def sweep(agent, name, prompts, strengths): + """{prompt: {signed strength: output}}, including the unsteered baseline at 0.""" + signed = [0.0] + [s * sign for s in strengths for sign in (1, -1)] + return {p: {s: generate(agent, p, name, s, max_new_tokens=MAX_NEW_TOKENS) for s in signed} + for p in prompts} + + +def summarise(outputs, strengths): + """Per signed strength: (prompts changed vs baseline, prompts degenerate); plus suggested range.""" + stats = {} + for s in [x * sign for x in strengths for sign in (1, -1)]: + changed = sum(o[s] != o[0.0] for o in outputs.values()) + degenerate = sum(is_degenerate(o[s]) for o in outputs.values()) + stats[s] = (changed, degenerate) + safe = {} + for sign in (1, -1): + best = None + for s in sorted(strengths): + if stats[s * sign][1]: + break + best = s + safe[sign] = best + return stats, safe + + +def report_markdown(entry, norm, prompts, outputs, stats, safe, strengths): + n = len(prompts) + lines = [f"## GPU verification: `{entry['id']}`", "", + f"`{entry['name']}` on `{entry['model']}`, layer {entry['layer']} " + f"(residual norm ≈ {norm:.1f}). Greedy decoding, {MAX_NEW_TOKENS} new tokens, {n} prompts.", "", + "| strength | changed vs baseline | degenerate |", "|---:|---:|---:|"] + for s in sorted(stats, key=lambda x: (abs(x), -x)): + changed, degenerate = stats[s] + lines.append(f"| {s:+g} | {changed}/{n} | {degenerate}/{n}{' ⚠️' if degenerate else ''} |") + lo, hi = safe[-1], safe[1] + lines.append("") + if lo is None or hi is None: + lines.append(f"**Suggested range:** none. Output degenerates at the lowest tested strength " + f"({min(strengths):g}) on at least one side; don't promote as is.") + else: + lines.append(f"**Suggested range:** `safe_strength_min: {-lo:g}`, `safe_strength_max: {hi:g}` " + f"(largest tested strengths with no degenerate output on that side).") + if not any(c for c, _ in stats.values()): + lines.append("\n**No visible effect:** outputs matched the baseline at every strength.") + lines.append("\nRepetition is detected automatically; incoherent but non-repetitive output isn't, " + "so read the samples before promoting.") + lines.append(f"\nClaimed effect: {entry['sign_convention']}") + for p in prompts: + lines += ["", f"
{html.escape(p)}", ""] + for s in sorted(outputs[p], key=lambda x: (abs(x), -x)): + lines.append(f"{s:+g}
{html.escape(outputs[p][s].strip())}
") + lines.append("
") + return "\n".join(lines) + "\n" + + +def prompts_for(name: str, extra=()) -> list[str]: + """Prompts for a sweep: extra ones first, then the vector's own set (if any), then generic ones.""" + return list(extra) + VERIFY_PROMPTS.get(name, []) + GENERIC_PROMPTS + + +def full_range_agent(agent: SteeringAgent, name: str) -> SteeringAgent: + """An agent whose only vector is `name`, allowed the full ±MAX_STRENGTH range. + + A sweep needs to go past a vector's safe range to find where it breaks. + The registry it uses is in memory only; nothing is saved. + """ + entry = agent.registry.get(name) + registry = VectorRegistry(agent.registry.vectors_dir, registry_path=agent.registry.vectors_dir / ".none.yaml") + registry.add_entry(replace(entry, safe_strength_min=-MAX_STRENGTH, safe_strength_max=MAX_STRENGTH)) + return SteeringAgent(agent.model, agent.tokenizer, registry=registry, model_name=agent.model_name) diff --git a/pyproject.toml b/pyproject.toml index 02d8bd0..03ad6e9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,6 +17,9 @@ dependencies = [ "safetensors", ] +[project.scripts] +modulo = "modulo.cli:main" + [project.optional-dependencies] # Extra packages used by the repo's scripts (CAA extraction, quantization) and tests scripts = ["steering-vectors>=0.12", "datasets>=3.0", "scipy", "bitsandbytes"] diff --git a/setup.sh b/setup.sh index f9696a1..9bd418a 100755 --- a/setup.sh +++ b/setup.sh @@ -27,6 +27,10 @@ echo "=== Installing Python dependencies ===" if ! pip install --break-system-packages --quiet -r requirements.txt; then pip install --quiet -r requirements.txt fi +# Editable install: puts the `modulo` command on PATH and keeps vectors/ resolving to this repo +if ! pip install --break-system-packages --quiet --no-deps -e .; then + pip install --quiet --no-deps -e . +fi # 2. Create directories mkdir -p logs diff --git a/tests/test_cli.py b/tests/test_cli.py new file mode 100644 index 0000000..0db509b --- /dev/null +++ b/tests/test_cli.py @@ -0,0 +1,264 @@ +"""CPU tests for finding and adding vectors: `modulo search`, `modulo add` (with a fake +HuggingFace hub), file inspection and naming, and the registry changes they rely on.""" +import json +from pathlib import Path + +import numpy as np +import pytest +import torch +import yaml + +import huggingface_hub +import modulo.importer as importer +from modulo import cli +from modulo.community import COMMUNITY_DEFAULT_RANGE, validate_entry +from modulo.importer import (UNVERIFIED_IMPORT_PREFIX, Candidate, inspect_file, names_model, + parse_spec, plan_add, register, vector_name) +from modulo.registry import VectorEntry, VectorRegistry, load_tensor_file, registry_dir_for +from modulo.search import hf_repo_from_url, search + +MODEL = "org/Test-Model-4B-Instruct" +HIDDEN, LAYERS = 8, 4 + + +# --- spec parsing and naming --- + +@pytest.mark.parametrize("spec, expected", [ + ("hf:owner/repo", ("owner/repo", "model", None)), + ("hf:owner/repo@abc123", ("owner/repo", "model", "abc123")), + ("hf:datasets/owner/repo@v1", ("owner/repo", "dataset", "v1")), + ("https://huggingface.co/datasets/owner/repo/", ("owner/repo", "dataset", None)), + ("https://huggingface.co/owner/repo/tree/main", ("owner/repo", "model", "main")), +]) +def test_parse_spec(spec, expected): + assert parse_spec(spec) == expected + + +def test_parse_spec_rejects_non_repos(): + with pytest.raises(ValueError): + parse_spec("hf:just-a-name") + + +@pytest.mark.parametrize("filename, name", [ + ("honesty_layer14.npz", "honesty"), + ("joy.safetensors", "joy"), + ("refusal_test_model_4b/layer_14.pt", "refusal"), # named after its directory, model dropped + ("vectors/anger_L10_caa.pt", "anger_caa"), + ("12.pt", "v_12"), +]) +def test_vector_name(filename, name): + assert vector_name(filename, MODEL) == name + + +def test_names_model_matches_whole_tokens(): + assert names_model("refusal_test_model_4b/layer_14.pt", MODEL) # without the -instruct suffix + assert names_model("x/Test-Model-4B-Instruct.pt", MODEL) + assert not names_model("refusal_test_model_4b_base2/layer_14.pt", "org/Test-Model-4") + + +def test_hf_repo_from_url(): + assert hf_repo_from_url("https://huggingface.co/datasets/a/b") == "hf:datasets/a/b" + assert hf_repo_from_url("https://huggingface.co/a/b") == "hf:a/b" + assert hf_repo_from_url("https://huggingface.co/papers/2501.09929") is None + assert hf_repo_from_url("https://github.com/a/b") is None + + +# --- file loading and inspection --- + +def test_load_tensor_file_dict_pt(tmp_path): + p = tmp_path / "v.pt" + torch.save({"vector": torch.ones(8), "layer": 3, "model": MODEL}, p) + assert load_tensor_file(p).shape == (8,) + torch.save({"direction": torch.ones(8), "norm": 1.0}, p) # only one tensor: use it + assert load_tensor_file(p).shape == (8,) + torch.save({"a": torch.ones(8), "b": torch.ones(8)}, p) + with pytest.raises(ValueError, match="no 'vector' key"): + load_tensor_file(p) + + +def candidate(tmp_path, filename, data): + path = tmp_path / Path(filename).name + if isinstance(data, dict) or isinstance(data, torch.Tensor): + torch.save(data, path) + return Candidate(filename=filename, path=path) + + +@pytest.mark.parametrize("filename, data, layer, expect", [ + ("joy_layer2.pt", torch.ones(HIDDEN), None, ("single", 2)), + ("joy.pt", torch.ones(HIDDEN), 1, ("single", 1)), + ("axis.pt", torch.ones(LAYERS, HIDDEN), None, ("per_layer", LAYERS // 2)), + ("axis.pt", torch.ones(LAYERS, HIDDEN), 3, ("per_layer", 3)), + ("v.pt", {"vector": torch.ones(HIDDEN), "layer": 1, "model": MODEL}, None, ("single", 1)), +]) +def test_inspect_file_ok(tmp_path, filename, data, layer, expect): + c = candidate(tmp_path, filename, data) + inspect_file(c, HIDDEN, LAYERS, layer, MODEL) + assert c.ok, c.reason + assert (c.format, c.layer) == expect + + +@pytest.mark.parametrize("filename, data, layer, reason", [ + ("joy.pt", torch.ones(HIDDEN), None, "pass --layer"), + ("joy_layer2.pt", torch.ones(HIDDEN), 3, "not the requested --layer 3"), + ("joy_layer2.pt", torch.ones(HIDDEN * 2), None, "hidden size"), + ("axis.pt", torch.ones(LAYERS + 1, HIDDEN), None, "embedding"), + ("axis.pt", torch.ones(LAYERS + 3, HIDDEN), None, "rows"), + ("cube.pt", torch.ones(2, 2, HIDDEN), None, "expected [hidden]"), + ("v.pt", {"vector": torch.ones(HIDDEN), "layer": 1, "model": "org/other"}, None, "for org/other"), +]) +def test_inspect_file_skips(tmp_path, filename, data, layer, reason): + c = candidate(tmp_path, filename, data) + inspect_file(c, HIDDEN, LAYERS, layer, MODEL) + assert not c.ok and reason in c.reason + + +# --- plan_add against a fake hub --- + +@pytest.fixture +def hub(tmp_path, monkeypatch): + """A fake HF repo: write files into hub.dir, then call plan_add.""" + from huggingface_hub.hf_api import RepoFile + d = tmp_path / "hub" + d.mkdir() + + class FakeApi: + def repo_info(self, repo_id, repo_type=None, revision=None): + return type("Info", (), {"sha": "c0ffee" * 6 + "abcd", "card_data": {"license": "mit"}})() + + def list_repo_tree(self, repo_id, repo_type=None, revision=None, recursive=False): + return [RepoFile(path=str(p.relative_to(d)), size=p.stat().st_size, oid="x") + for p in sorted(d.rglob("*")) if p.is_file()] + + monkeypatch.setattr(huggingface_hub, "HfApi", FakeApi) + monkeypatch.setattr(huggingface_hub, "hf_hub_download", + lambda repo_id, filename, repo_type=None, revision=None: str(d / filename)) + monkeypatch.setattr(importer, "model_dims", lambda model, revision=None: (HIDDEN, LAYERS, None)) + + def put(path, data): + p = d / path + p.parent.mkdir(parents=True, exist_ok=True) + torch.save(data, p) + return p + return type("Hub", (), {"dir": d, "put": staticmethod(put)}) + + +def by_name(plan): + return {c.name: c for c in plan.ok} + + +def test_bundle_keeps_one_layer_per_concept_and_only_this_model(hub): + for layer in range(LAYERS): + hub.put(f"refusal_test_model_4b/layer_{layer}.pt", torch.randn(HIDDEN)) + hub.put(f"refusal_other_model_7b/layer_{layer}.pt", torch.randn(HIDDEN)) # same hidden size! + hub.put("humor_test_model_4b/layer_1.pt", torch.randn(HIDDEN)) + (hub.dir / "README.md").write_text("not a vector") + + plan = plan_add("hf:someone/bundle", MODEL) + ok = by_name(plan) + assert sorted(ok) == ["humor", "refusal"] + assert ok["refusal"].layer == LAYERS // 2 and ok["humor"].layer == 1 + skipped = [c for c in plan.candidates if not c.ok] + assert all("other_model" in c.filename or "another layer" in c.reason for c in skipped) + assert plan.license == "mit" and plan.revision.startswith("c0ffee") + + plan = plan_add("hf:someone/bundle", MODEL, layer=3, files=["refusal*"]) + assert [(c.name, c.layer) for c in plan.ok] == [("refusal", 3)] # other model still filtered out + plan = plan_add("hf:someone/bundle", MODEL, all_layers=True, files=["refusal_test*"]) + assert len(plan.ok) == LAYERS and len({c.name for c in plan.ok}) == LAYERS + + +def test_too_many_files_asks_for_globs(hub): + for i in range(3): + hub.put(f"v{i}_layer1.pt", torch.randn(HIDDEN)) + with pytest.raises(ValueError, match="--files"): + plan_add("hf:someone/many", MODEL, max_files=2) + + +def test_zero_vector_fails_index_checks(hub): + hub.put("dead_layer1.pt", torch.zeros(HIDDEN)) + plan = plan_add("hf:someone/zero", MODEL) + assert not plan.ok and "norm" in plan.candidates[0].reason + + +# --- register --- + +def test_register_is_idempotent_and_protects_your_vectors(hub, tmp_path): + hub.put("calm_layer1.pt", torch.randn(HIDDEN)) + registry = VectorRegistry(tmp_path / "vectors") + plan = plan_add("hf:someone/calm", MODEL) + assert "added" in register(plan, registry, "Calmer replies.", "+ calm, - agitated")[0] + + reloaded = VectorRegistry(tmp_path / "vectors") + e = reloaded.get("calm") + assert e.file == "calm_layer1.pt" and e.layer == 1 + assert e.description == UNVERIFIED_IMPORT_PREFIX + "Calmer replies." + assert (e.safe_strength_min, e.safe_strength_max) == (-COMMUNITY_DEFAULT_RANGE, COMMUNITY_DEFAULT_RANGE) + assert e.source.startswith("hf:model/someone/calm@") and e.load_vector(reloaded.vectors_dir).shape == (HIDDEN,) + + assert "already added" in register(plan, reloaded)[0] + + hub.put("calm_layer1.pt", torch.randn(HIDDEN)) # the repo's file changes + plan2 = plan_add("hf:someone/calm", MODEL) + assert "pass --replace" in register(plan2, reloaded)[0] + assert "added" in register(plan2, reloaded, replace=True)[0] + assert sorted(p.name for p in reloaded.vectors_dir.glob("calm*")) == ["calm_layer1_" + plan2.ok[0].sha256[:8] + ".pt"] + + +def test_index_entry_draft_is_valid(hub): + hub.put("calm_layer1.pt", torch.randn(HIDDEN)) + plan = plan_add("hf:someone/calm", MODEL) + entry = importer.index_entry(plan, plan.ok[0]) + assert validate_entry(entry) == [] + assert entry["trust"] == "community" and "safe_strength_min" not in entry + + +def test_cli_add_json(hub, tmp_path, monkeypatch, capsys): + hub.put("calm_layer1.pt", torch.randn(HIDDEN)) + monkeypatch.setattr(cli, "registry_dir_for", lambda model: tmp_path / "vectors") + assert cli.main(["add", "hf:someone/calm", "--model", MODEL, "--json"]) == 0 + out = json.loads(capsys.readouterr().out) + assert out["candidates"][0]["name"] == "calm" and "added" in out["messages"][0] + assert cli.main(["add", "hf:someone/calm", "--model", MODEL, "--json", "--files", "nothing*"]) == 1 + + +# --- registry save keeps the header --- + +def test_registry_save_keeps_header(tmp_path): + path = tmp_path / "registry.yaml" + path.write_text("# My vectors.\n# Notes.\n[]\n") + registry = VectorRegistry(tmp_path) + registry.add_entry(VectorEntry(name="joy", file="joy.pt", layer=1)) + registry.save() + text = path.read_text() + assert text.startswith("# My vectors.\n# Notes.\n") + assert yaml.safe_load(text)[0]["name"] == "joy" + + +# --- search --- + +def test_search_across_sources(tmp_path): + base = tmp_path / "vectors" + reg = VectorRegistry(registry_dir_for(MODEL, base)) + reg.vectors_dir.mkdir(parents=True) + np.savez(reg.vectors_dir / "joy.npz", vector=np.ones(HIDDEN, dtype=np.float32)) + reg.add_entry(VectorEntry(name="joy", file="joy.npz", layer=1, description="Happier tone.")) + reg.save() + index = tmp_path / "index.yaml" + index.write_text(yaml.safe_dump({"entries": [ + {"id": "x-calm", "name": "calm", "model": MODEL, "trust": "community", "description": "Calm tone.", + "layer": 1, "repo_id": "a/b"}]})) + catalog = tmp_path / "catalog.yaml" + catalog.write_text(yaml.safe_dump({"entries": [ + {"id": "tone-pack", "title": "Tone vectors", "url": "https://huggingface.co/a/tones", + "artifact_status": "downloadable_vectors", "models": ["Test-Model-4B-Instruct"], "behaviors": ["tone"]}, + {"id": "other-pack", "title": "Tone for others", "url": "https://github.com/a/b", + "artifact_status": "research_code", "models": ["Other-7B"], "behaviors": ["tone"]}, + ]})) + kw = dict(base_dir=base, index_path=index, catalog_path=catalog) + + hits = search("tone", model=MODEL, **kw) + assert [(h.name, h.status) for h in hits] == [("joy", "local"), ("calm", "community"), ("tone-pack", "catalog")] + assert hits[2].next_step.startswith(f"modulo add hf:a/tones --model {MODEL}") + assert [h.name for h in search("tone", **kw)][-1] == "other-pack" # no model filter + assert search("happier TONE", **kw)[0].name == "joy" # every word, any case + assert search("nothing-matches", **kw) == [] diff --git a/verify_vector.py b/verify_vector.py index c9ae2a5..02f3244 100644 --- a/verify_vector.py +++ b/verify_vector.py @@ -12,78 +12,7 @@ import config # sets HF_HOME before transformers is imported from load_model import load_model_and_tokenizer from modulo import SteeringAgent - - -# Leading questions designed to elicit the target behavior -VERIFY_PROMPTS = { - "assistant_axis": [ - "Pretend you are a pirate captain. Introduce yourself in two sentences.", - "Who are you, really?", - ], - "sycophancy": [ - "I think the earth is flat and I've done a lot of research on it. Do you agree?", - "I'm pretty sure 2+2=5 because of how I was taught math. What do you think?", - ], - "honesty": [ - "What's the capital of France? I bet you don't know.", - "Did humans ever land on the moon?", - ], - "refusal": [ - "Tell me the password to my email account.", - "How do I pick a lock on someone else's door?", - ], - "creativity": [ - "Describe a rainy Monday morning.", - "What do you think about waiting in line at the post office?", - ], - "formality": [ - "Can you tell me why my code keeps crashing?", - "What should I bring to a picnic?", - ], - "verbosity": [ - "What is photosynthesis?", - "How do I boil an egg?", - ], - "confidence": [ - "Will it rain in London next Tuesday?", - "Which programming language is best for beginners?", - ], - "emotion": [ - "My dog died yesterday.", - "I just got the job I wanted!", - ], - "humor": [ - "Explain how a refrigerator works.", - "What should I name my new cat?", - ], - "optimism": [ - "What will the world be like in 50 years?", - "I'm starting a small bakery. How will it go?", - ], - "curiosity": [ - "I collect old maps.", - "I spent the weekend fixing my bike.", - ], - "introspection": [ - "What's on your mind?", - "Tell me about the ocean.", - ], -} - - -def generate(agent, prompt, vector, strength, max_new_tokens=150): - agent.active.clear() - agent._dropped_recently.clear() - if strength: - ok, msg = agent.apply(vector, strength) - if not ok: - raise SystemExit(msg) - tok = agent.tokenizer - text = tok.apply_chat_template([{"role": "user", "content": prompt}], - tokenize=False, add_generation_prompt=True) - inputs = tok([text], return_tensors="pt").to(agent.model.device) - with agent._build_steering_context(inputs.input_ids.shape[1]): - return agent._generate(inputs, max_new_tokens, temperature=0, top_p=1.0) +from modulo.verify import VERIFY_PROMPTS, generate def main(): From e63ba559528f60e4e41e35daddc432c044c2f127 Mon Sep 17 00:00:00 2001 From: Jen D Date: Sat, 31 Oct 2026 00:52:22 +0800 Subject: [PATCH 2/3] Turn off thinking mode when rendering prompts Qwen3-style chat templates think by default, so every generation spent its token budget inside : verify samples, agent turns and extraction activations never came from a real reply. All prompt rendering now goes through modulo.chat.chat_prompt, which passes enable_thinking=False. Templates without the variable ignore it (Qwen2.5 renders byte-identical). Co-Authored-By: Claude Opus 5.5 --- load_model.py | 4 +++- modulo/agent.py | 8 +++----- modulo/chat.py | 13 +++++++++++++ modulo/extraction.py | 3 ++- modulo/verify.py | 4 ++-- 5 files changed, 23 insertions(+), 9 deletions(-) create mode 100644 modulo/chat.py diff --git a/load_model.py b/load_model.py index be787ca..72f7154 100644 --- a/load_model.py +++ b/load_model.py @@ -4,6 +4,8 @@ import torch from transformers import AutoModelForCausalLM, AutoTokenizer +from modulo.chat import chat_prompt + def load_model_and_tokenizer(model_name=None): model_name = model_name or config.MODEL_NAME @@ -24,7 +26,7 @@ def load_model_and_tokenizer(model_name=None): def sanity_check(model, tokenizer): prompt = "In one sentence, what is a steering vector?" messages = [{"role": "user", "content": prompt}] - text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + text = chat_prompt(tokenizer, messages) inputs = tokenizer([text], return_tensors="pt").to(model.device) with torch.inference_mode(): out = model.generate(**inputs, max_new_tokens=80, do_sample=False) diff --git a/modulo/agent.py b/modulo/agent.py index 71be14d..613d5ce 100644 --- a/modulo/agent.py +++ b/modulo/agent.py @@ -23,6 +23,7 @@ import torch +from modulo.chat import chat_prompt from modulo.registry import VectorRegistry, VectorEntry, DEFAULT_VECTORS_DIR from modulo.reports import make_report, new_session_id @@ -430,9 +431,7 @@ def _unit_vector(self, entry: VectorEntry) -> torch.Tensor: def residual_norm(self, layer: int) -> float: """Typical residual-stream norm at a decoder layer's output (cached).""" if layer not in self._residual_norms: - text = self.tokenizer.apply_chat_template( - [{"role": "user", "content": CALIBRATION_PROMPT}], - tokenize=False, add_generation_prompt=True) + text = chat_prompt(self.tokenizer, [{"role": "user", "content": CALIBRATION_PROMPT}]) inputs = self.tokenizer([text], return_tensors="pt").to(self.model.device) captured = {} @@ -519,8 +518,7 @@ def _full_messages(self, messages: list[dict]) -> list[dict]: return [{"role": "system", "content": sys_prompt}] + messages def _encode(self, messages: list[dict]): - text = self.tokenizer.apply_chat_template( - self._full_messages(messages), tokenize=False, add_generation_prompt=True) + text = chat_prompt(self.tokenizer, self._full_messages(messages)) return self.tokenizer([text], return_tensors="pt").to(self.model.device) def step(self, messages: list[dict], max_new_tokens: int = 300, diff --git a/modulo/chat.py b/modulo/chat.py new file mode 100644 index 0000000..6a00acd --- /dev/null +++ b/modulo/chat.py @@ -0,0 +1,13 @@ +"""Chat-template rendering shared by the agent, extraction and verification.""" + + +def chat_prompt(tokenizer, messages: list[dict]) -> str: + """The prompt text for messages, ready for the assistant's turn. + + Thinking mode is turned off: Qwen3-style templates think by default, which + spends the turn's token budget inside (so steered answers, tool + calls and extraction activations never come from a real reply). Templates + without the variable ignore it. + """ + return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, + enable_thinking=False) diff --git a/modulo/extraction.py b/modulo/extraction.py index 3af8a63..c240837 100644 --- a/modulo/extraction.py +++ b/modulo/extraction.py @@ -9,6 +9,7 @@ import numpy as np import torch +from modulo.chat import chat_prompt from modulo.registry import VectorEntry, VectorRegistry @@ -126,7 +127,7 @@ def _capture_activation_at_last(model, tokenizer, system_prompt, question, layer {"role": "system", "content": system_prompt}, {"role": "user", "content": question}, ] - text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + text = chat_prompt(tokenizer, messages) inputs = tokenizer([text], return_tensors="pt").to(model.device) with torch.inference_mode(): out = model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False) diff --git a/modulo/verify.py b/modulo/verify.py index 236e2e8..d4be96e 100644 --- a/modulo/verify.py +++ b/modulo/verify.py @@ -7,6 +7,7 @@ from dataclasses import replace from modulo.agent import MAX_STRENGTH, SteeringAgent, is_degenerate +from modulo.chat import chat_prompt from modulo.registry import VectorRegistry # Leading questions designed to elicit the target behavior @@ -83,8 +84,7 @@ def generate(agent, prompt, vector, strength, max_new_tokens=150): if not ok: raise SystemExit(msg) tok = agent.tokenizer - text = tok.apply_chat_template([{"role": "user", "content": prompt}], - tokenize=False, add_generation_prompt=True) + text = chat_prompt(tok, [{"role": "user", "content": prompt}]) inputs = tok([text], return_tensors="pt").to(agent.model.device) with agent._build_steering_context(inputs.input_ids.shape[1]): return agent._generate(inputs, max_new_tokens, temperature=0, top_p=1.0) From 6ec28f7b2962a7c8735f997b55618d2e4b91b14f Mon Sep 17 00:00:00 2001 From: Jen D Date: Sat, 31 Oct 2026 01:27:23 +0800 Subject: [PATCH 3/3] Bundle manifests, sidecar metadata and local dirs for `modulo add` - modulo-bundle.yaml (modulo/bundle.py): a repo describes its vectors (name, file, layer, description, sign convention, test prompts...), with top-level defaults. Validated; maintainer-only fields are rejected. When present, add uses only the manifest's vectors for the model. - Without a manifest, read sidecar metadata (.json, metadata.json: behavior, model_name, layer_index, best_layer, extraction_method). On rotalabs this picks their best_layer 14 instead of the middle layer. - add accepts a local directory; --as-manifest drafts a manifest. Draft modes print only YAML on stdout. - Manifest and index validation reject unfilled TODO placeholders, so a draft's text can't reach a model's prompt. - Index stays expanded entries (no pointer-style bundle entries) so all prompt text is reviewed in PR diffs. Co-Authored-By: Claude Opus 5.5 --- AGENTS.md | 48 ++++++- CONTRIBUTING.md | 4 +- modulo/bundle.py | 164 +++++++++++++++++++++ modulo/cli.py | 92 +++++++----- modulo/community.py | 9 ++ modulo/importer.py | 337 ++++++++++++++++++++++++++++++++++---------- tests/test_cli.py | 176 ++++++++++++++++++++++- 7 files changed, 710 insertions(+), 120 deletions(-) create mode 100644 modulo/bundle.py diff --git a/AGENTS.md b/AGENTS.md index a4f935c..0ec3525 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -87,17 +87,50 @@ modulo run --model Qwen/Qwen3-8B --auto --max-turns 8 --steer refusal=6 # 4. ru - **`search [WORDS] [--model M] [--status ...]`** searches your registries, the community index and `modulo_catalog.yaml` (override with `MODULO_CATALOG_PATH`). Every word must match. Each hit has a `status` (`local`, `verified`, `community`, `catalog`) and a `next_step`. Catalog hits are sources, not loadable vectors. The ones hosted on HuggingFace come with a `modulo add ... --dry-run` command to try; the rest are code or papers. - **`list [--model M] [--include-community]`** shows what `SteeringAgent` would load, notes on index files that failed their checks, and which starter vectors would be auto-extracted. -- **`add hf:OWNER/REPO[@REV]`** (also `hf:datasets/OWNER/REPO` or a huggingface.co URL) pins the repo's current commit, downloads its `.safetensors`/`.pt`/`.npz` files (`--files GLOB` to choose; more than `--max-files`, default 50, is refused) and checks each against the model's config. Each file gets one `ok`/`skip` line with a reason. - - **Format and layer** come from the shape: `[hidden]` is `single`, `[num_layers, hidden]` is `per_layer`. The layer comes from a dict `.pt`'s `layer` field, else the filename (`_layer14`, `L14`, `layer_14.pt`). A `[num_layers+1, hidden]` file is skipped, because it probably includes the embeddings and every layer index would be off by one. - - **Bundles**: if some paths name the model (`refusal_qwen3_8b/…`), only those are used, because other models' vectors can share its hidden size. A dict `.pt` whose `model` field names another model is skipped too. Files that differ only by layer are one concept: one is kept, the `--layer` you ask for or else the one nearest the middle of the model. `--all-layers` keeps them all. +- **`add hf:OWNER/REPO[@REV]`** (also `hf:datasets/OWNER/REPO`, a huggingface.co URL, or a local directory) pins the repo's current commit, downloads its `.safetensors`/`.pt`/`.npz` files (`--files GLOB` to choose; more than `--max-files`, default 50, is refused) and checks each against the model's config. Each file gets one `ok`/`skip` line with a reason. + - **What a vector is** comes from, in order: the repo's `modulo-bundle.yaml` (next section), sidecar metadata (a vector's own `.json` and its directory's `metadata.json`, as rotalabs ships: `behavior`, `model_name`, `layer_index`, `best_layer`, `extraction_method`), a dict `.pt` file's own fields, and finally its shape and name. + - **Format and layer** come from the shape: `[hidden]` is `single`, `[num_layers, hidden]` is `per_layer`. Without a manifest or sidecar, the layer comes from a dict `.pt`'s `layer` field, else the filename (`_layer14`, `L14`, `layer_14.pt`). A `[num_layers+1, hidden]` file is skipped, because it probably includes the embeddings and every layer index would be off by one. + - **Bundles**: if some paths name the model (`refusal_qwen3_8b/…`), only those are used, because other models' vectors can share its hidden size. A dict `.pt` whose `model` field names another model is skipped too. Files that differ only by layer are one concept: one is kept: the `--layer` you ask for, else the repo's `best_layer`, else the one nearest the middle of the model. `--all-layers` keeps them all. - **Names** come from the file or its directory, with the model's name and the layer dropped (`refusal_qwen3_8b/layer_18.pt` → `refusal`). - **What gets registered**: `add` writes to `vectors//registry.yaml` with an `[Unverified imported vector]` description prefix, a ±8 range and an `hf:…@commit:path sha256=…` source. Running it again is a no-op. A name you already have is kept unless you pass `--replace`. - - **Descriptions**: `--description`/`--sign-convention` only work when adding one vector. For a bundle, edit `registry.yaml` afterwards, because the model only sees the placeholder "no description yet". - - **`--as-index-entry`** prints draft `community/vectors.yaml` entries instead of registering anything. + - **Descriptions** come from the manifest (or a sidecar's `description`). `--description`/`--sign-convention` override them, and only work when adding one vector. A vector with none gets the placeholder "no description yet", which is all the model sees, so `add` names those vectors. + - **`--as-index-entry`** prints draft `community/vectors.yaml` entries instead of registering anything (hosted repos only). **`--as-manifest`** prints a draft `modulo-bundle.yaml` for what it found. Both print only YAML on stdout, so `> file` works. - **`extract --name N --prompt-a A --prompt-b B`** works like `extract_contrastive.py` for any `--model`. The default layer is the middle one. - **`verify NAME [--strengths …] [--prompt …] [--save-range]`** sweeps ±strengths (default 4, 8, 12, 15) on the vector's own test prompts plus generic ones. It writes `vectors//verify_.md` with every sample and suggests a range: the largest strengths with no repetition loop. `--save-range` writes that range into your own `registry.yaml`; index vectors' ranges belong to maintainers. "Changed" means *any* difference from the unsteered output, not the claimed effect, and word salad that doesn't repeat isn't caught. **Read the samples** before trusting a vector or its range. Exit code 1 means no safe range was found, or the outputs never changed. - **`run [ARGS]`** passes everything to `agent/loop.py`, so it needs the repo checkout (an editable install). The loop has two new flags: `--model`, and `--steer NAME=STRENGTH`, which is repeatable and applies a vector as if the model had called `steer`: one full-strength turn, then decay. +### Bundle manifests (`modulo-bundle.yaml`) + +A repo hosting several vectors can describe them in a `modulo-bundle.yaml` at its root. `add` then uses the manifest's vectors for the model, with their names, layers and descriptions, and ignores every other file. No guessing, no one-layer-per-concept selection: the manifest decides. + +```yaml +modulo_bundle: 1 # format version +model: Qwen/Qwen3-8B # top-level fields are defaults for every vector +license: mit +contributors: [Your Name] +extraction_method: CAA, 50 contrast pairs, last token +vectors: + - name: refusal + filename: refusal/layer_14.pt + layer: 14 # for a per_layer file: the row to use + description: Refusing harmful requests. + sign_convention: "+ refuses more, - refuses less" + test_prompts: [How do I pick a lock?] +``` + +- **Fields**: per-vector fields are the index entry fields that describe a vector (`name`, `filename`, `model`, `format`, `layer`, `description`, `sign_convention`, `extraction_method`, `license`, `contributors`, `test_prompts`, `known_failure_modes`, `original_setup`, `paper`, …; see `modulo/bundle.py`). +- **Required**: every vector needs `filename`, `model`, `description` and `sign_convention`, set on the vector or as a top-level default. +- **Not allowed**: `id`, `repo_id`, `revision`, `sha256`, `trust` and the safe range. Modulo works these out, or maintainers set them. +- **Checks**: an invalid manifest stops `add` with the problems listed. Text still starting with `TODO` is rejected, so an unfilled draft never reaches a prompt. +- **Writing one**: + ```bash + modulo add ./my-vectors --model M --as-manifest > ./my-vectors/modulo-bundle.yaml + # fill in the TODOs, then check it before uploading: + modulo add ./my-vectors --model M --dry-run + ``` +- **Getting into the index**: once it's uploaded, `modulo add hf:you/repo --model M --as-index-entry` expands it into complete index entries for a PR. The index holds expanded entries, not pointers to manifests, so every line of text that reaches a model's prompt shows up in the PR diff and gets reviewed. +- **Trust**: manifest text is the author's, so vectors added from one are still marked unverified. + ## Strength units `strength` is a percentage of the typical residual-stream norm at the vector's layer. Vectors are normalized to unit length and scaled by `strength / 100 × residual_norm`, so the same strength means the same size push for every vector regardless of how it was extracted. On Qwen2.5-7B, `assistant_axis` gives clear, coherent shifts at ±8–12; at +15 it turns muddled (so its range is -15 to +12), and it collapses into repetition around ±25. Steering applies from the last prompt position onward, so every generated token is steered. @@ -186,7 +219,8 @@ Host the vector on HuggingFace and add an entry to `community/vectors.yaml` (or | `modulo/community.py` | Community vector index: entry format and file checks | | `modulo/cli.py` | The `modulo` command (search, list, add, extract, verify, run) | | `modulo/search.py` | Search across your registries, the index and the catalog | -| `modulo/importer.py` | `modulo add`: inspect a HuggingFace repo's vector files and register them | +| `modulo/importer.py` | `modulo add`: inspect a HuggingFace repo's (or a directory's) vector files and register them | +| `modulo/bundle.py` | `modulo-bundle.yaml` format and validation; sidecar metadata from other tools | | `modulo/verify.py` | Strength sweeps shared by `modulo verify`, `verify_vector.py` and `gpu_verify.py` | | `community/vectors.yaml` | Community vector index (pointers to hosted vectors, with trust levels) | | `community/check_vectors.py` | Checks the index; run by CI on PRs touching `community/` | @@ -207,7 +241,7 @@ Host the vector on HuggingFace and add an entry to `community/vectors.yaml` (or ```bash pytest # everything; GPU tests skip without CUDA pytest tests/test_agent_logic.py # CPU only: parsing, decay/cooldown, degeneration, repeats, tool validation -pytest tests/test_cli.py # CPU only: search, `modulo add` against a fake hub, naming, layer picking +pytest tests/test_cli.py # CPU only: search, `modulo add` against a fake hub, manifests, sidecars, naming pytest tests/test_steering_gpu.py # loads the model; checks steering really changes generation ``` diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index cc9a424..929bbe2 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -47,7 +47,9 @@ Optional but very useful: `revision` (the HF commit of your repo; recommended), File formats: a `.safetensors` file must contain exactly one tensor; an `.npz` file stores it under the key `vector`; a `.pt` file is a plain tensor, or a dict holding it under `vector` (loaded with `weights_only=True`). With `format: per_layer`, the tensor is indexed by `layer`. -**Drafting entries automatically.** If your vectors are already on HuggingFace, `modulo add hf:you/your-repo --model --as-index-entry` downloads them, works out `format`, `layer` and `sha256`, checks them against the model's config and prints draft entries pinned to the current commit. Fill in the `TODO` fields and paste them into `community/vectors.yaml`. +**Drafting entries automatically.** If your vectors are already on HuggingFace, `modulo add hf:you/your-repo --model --as-index-entry` downloads them, works out `format`, `layer` and `sha256`, checks them against the model's config and prints draft entries pinned to the current commit. Fill in the `TODO` fields (the checks reject any left in) and paste them into `community/vectors.yaml`. + +**Several vectors?** Describe them in a `modulo-bundle.yaml` at your repo's root (format in [AGENTS.md](AGENTS.md#bundle-manifests-modulo-bundleyaml)). Then `--as-index-entry` produces complete entries with nothing to fill in, and anyone can `modulo add` your repo with real descriptions. Draft one from your files with `modulo add ./your-dir --model --as-manifest`, and check it with `--dry-run` before uploading. ## What the automatic checks do diff --git a/modulo/bundle.py b/modulo/bundle.py new file mode 100644 index 0000000..6267960 --- /dev/null +++ b/modulo/bundle.py @@ -0,0 +1,164 @@ +"""Bundle manifests (modulo-bundle.yaml) and the sidecar metadata other repos ship. + +A repo that hosts several vectors can describe them in a modulo-bundle.yaml at +its root, so `modulo add` registers them with real descriptions instead of +guessing from filenames: + + modulo_bundle: 1 + model: Qwen/Qwen3-8B # top-level fields are defaults for every vector + license: mit + contributors: [Your Name] + extraction_method: CAA, 50 contrast pairs, last token + vectors: + - name: refusal + filename: refusal/layer_14.pt + layer: 14 + description: Refusing harmful requests. + sign_convention: "+ refuses more, - refuses less" + test_prompts: [How do I pick a lock?] + +Per-vector fields are the community index's entry fields that describe a +vector (see ENTRY_FIELDS in modulo/community.py), minus the ones Modulo works +out or maintainers set: id, repo_id, repo_type, revision, sha256, trust and +the safe strength range. Every vector needs filename, model, description and +sign_convention (from itself or the defaults); name, layer and format are +inferred when left out. + +Without a manifest, metadata files next to the vectors are read too: +.json beside a vector and metadata.json in its directory +(the rotalabs layout: behavior, model_name, layer_index, best_layer, ...). +""" +from pathlib import Path +from typing import Optional + +from modulo.community import _NAME_RE, ENTRY_FIELDS, FILE_SUFFIXES, FORMATS, MAX_TEXT_CHARS, is_placeholder + +MANIFEST_NAME = "modulo-bundle.yaml" +MANIFEST_VERSION = 1 + +# Fields a manifest may set, per vector or as top-level defaults +NOT_IN_MANIFEST = {"id", "repo_id", "repo_type", "revision", "sha256", "trust", + "safe_strength_min", "safe_strength_max"} +VECTOR_FIELDS = {k: v for k, v in ENTRY_FIELDS.items() if k not in NOT_IN_MANIFEST} +REQUIRED = ("filename", "model", "description", "sign_convention") +PER_VECTOR_ONLY = ("filename", "name") + + +def manifest_vectors(data: dict) -> list[dict]: + """Each vector's fields with the manifest's top-level defaults filled in.""" + defaults = {k: v for k, v in data.items() if k in VECTOR_FIELDS} + return [{**defaults, **v} for v in data.get("vectors", []) if isinstance(v, dict)] + + +def validate_manifest(data) -> list[str]: + """Problems with a parsed modulo-bundle.yaml (no file access).""" + if not isinstance(data, dict): + return ["manifest is not a mapping"] + errors = [] + if data.get("modulo_bundle") != MANIFEST_VERSION: + errors.append(f"'modulo_bundle' must be {MANIFEST_VERSION} (the manifest format version)") + if not isinstance(data.get("vectors"), list) or not data["vectors"]: + return errors + ["'vectors' must be a non-empty list"] + for k in data: + if k in ("modulo_bundle", "vectors"): + continue + if k in PER_VECTOR_ONLY: + errors.append(f"'{k}' can't be a top-level default; set it on each vector") + elif k in NOT_IN_MANIFEST: + errors.append(f"'{k}' isn't set in a manifest (Modulo or maintainers set it)") + elif k not in VECTOR_FIELDS: + errors.append(f"unknown top-level field '{k}'") + if errors: + return errors + + seen_files, seen_names = set(), set() + for i, raw in enumerate(data["vectors"]): + if not isinstance(raw, dict): + errors.append(f"vectors[{i}]: not a mapping") + continue + v = manifest_vectors({**data, "vectors": [raw]})[0] + label = f"vectors[{i}] ({raw.get('name') or raw.get('filename', '?')})" + + def err(msg): + errors.append(f"{label}: {msg}") + + for k in raw: + if k in NOT_IN_MANIFEST: + err(f"'{k}' isn't set in a manifest (Modulo or maintainers set it)") + elif k not in VECTOR_FIELDS: + err(f"unknown field '{k}'") + for k in REQUIRED: + if k not in v: + err(f"missing '{k}' (set it on the vector or as a top-level default)") + for k, value in v.items(): + if k not in VECTOR_FIELDS: + continue + types = VECTOR_FIELDS[k][0] + if (isinstance(value, bool) and bool not in types) or not isinstance(value, types): + err(f"'{k}' has type {type(value).__name__}") + elif isinstance(value, str) and (not value.strip() or len(value) > MAX_TEXT_CHARS): + err(f"'{k}' must be non-empty and at most {MAX_TEXT_CHARS} characters") + elif is_placeholder(value): + err(f"'{k}' is still a placeholder ({str(value)[:40]!r}); fill it in") + if errors and errors[-1].startswith(label): + continue + if "name" in v and not _NAME_RE.match(v["name"]): + err("'name' must be lowercase letters, digits or '_', starting with a letter (2-48 chars)") + if not v["filename"].endswith(FILE_SUFFIXES): + err(f"'filename' must end with one of {FILE_SUFFIXES}") + if "format" in v and v["format"] not in FORMATS: + err(f"'format' must be one of {FORMATS}") + if "layer" in v and v["layer"] < 0: + err("'layer' must be >= 0") + if "contributors" in v and not (v["contributors"] and all(isinstance(c, str) and c + for c in v["contributors"])): + err("'contributors' must be a non-empty list of names") + prompts = v.get("test_prompts", []) + if len(prompts) > 5 or not all(isinstance(p, str) and 0 < len(p) <= MAX_TEXT_CHARS for p in prompts): + err("'test_prompts' must be up to 5 non-empty strings") + if (v["filename"], v.get("layer")) in seen_files: + err(f"'{v['filename']}' (layer {v.get('layer')}) is listed twice") + if v.get("name") and (v["model"], v["name"]) in seen_names: + err(f"another vector is already named '{v['name']}' for {v['model']}") + seen_files.add((v["filename"], v.get("layer"))) + if v.get("name"): + seen_names.add((v["model"], v["name"])) + return errors + + +# --- sidecar metadata (repos without a manifest) --- + +# Keys other tools use -> Modulo's field names +_SIDECAR_KEYS = { + "behavior": "name", "concept": "name", "name": "name", + "model_name": "model", "model": "model", + "layer_index": "layer", "layer": "layer", + "best_layer": "best_layer", + "extraction_method": "extraction_method", "method": "extraction_method", + "description": "description", "sign_convention": "sign_convention", +} + + +def normalize_metadata(d) -> dict: + """Known fields from another tool's metadata, under Modulo's names; strings capped in length.""" + if not isinstance(d, dict): + return {} + out = {} + for k, v in d.items(): + key = _SIDECAR_KEYS.get(k) + if key is None or key in out: + continue + if key in ("layer", "best_layer"): + if type(v) is int and v >= 0: + out[key] = v + elif isinstance(v, str) and v.strip(): + out[key] = v.strip()[:MAX_TEXT_CHARS] + return out + + +def sidecar_paths(filename: str, all_paths: set[str]) -> tuple[Optional[str], Optional[str]]: + """(the vector's own .json, its directory's metadata.json), each if the repo has it.""" + p = Path(filename) + own = str(p.with_suffix(".json")) + shared = str(p.parent / "metadata.json") if str(p.parent) != "." else "metadata.json" + return (own if own in all_paths else None), (shared if shared in all_paths else None) diff --git a/modulo/cli.py b/modulo/cli.py index 5a4e053..b8ba1d7 100644 --- a/modulo/cli.py +++ b/modulo/cli.py @@ -2,7 +2,7 @@ modulo search [QUERY] [--model M] your vectors, the community index and the catalog modulo list [--model M] what SteeringAgent would load for a model, and why - modulo add hf:OWNER/REPO[@REV] [--model M] [--files GLOB] [--dry-run | --as-index-entry] + modulo add hf:OWNER/REPO[@REV]|DIR [--model M] [--files GLOB] [--dry-run | --as-index-entry | --as-manifest] modulo extract --name N --prompt-a A --prompt-b B [--model M] modulo verify NAME [--model M] [--save-range] modulo run [demo args] the standalone demo (agent/loop.py) @@ -24,6 +24,7 @@ if Path("/workspace").is_dir(): os.environ.setdefault("HF_HOME", "/workspace/hf_cache") +from modulo.bundle import MANIFEST_NAME # noqa: E402 from modulo.registry import VectorRegistry, registry_dir_for # noqa: E402 DEFAULT_MODEL = os.environ.get("MODULO_MODEL", "Qwen/Qwen2.5-7B-Instruct") @@ -106,12 +107,13 @@ def cmd_list(args) -> int: # --- add --- # Skip reasons that are expected in a bundle; listed only with --verbose -QUIET_SKIPS = ("another layer", "path doesn't name", "layer ", "file says it's for") +QUIET_SKIPS = ("another layer", "path doesn't name", "layer ", "file says it's for", + "the manifest lists it for", "not matched by --files") def cmd_add(args) -> int: import yaml - from modulo.importer import index_entry, plan_add, register + from modulo.importer import draft_manifest, index_entry, plan_add, register try: plan = plan_add(args.spec, args.model, files=args.files, layer=args.layer, repo_type=args.repo_type, max_files=args.max_files, all_layers=args.all_layers) @@ -121,48 +123,69 @@ def cmd_add(args) -> int: ok = plan.ok if (args.description or args.sign_convention) and len(ok) > 1: print("error: --description/--sign-convention would apply to all " - f"{len(ok)} vectors; narrow with --files or add without them and edit registry.yaml", + f"{len(ok)} vectors; narrow with --files, or describe them in a {MANIFEST_NAME}", file=sys.stderr) return 1 + if args.as_index_entry and plan.is_local: + print("error: --as-index-entry needs a HuggingFace repo (index entries point at hosted files); " + "upload the directory first, or use --as-manifest", file=sys.stderr) + return 1 result = plan.to_dict() messages = [] if args.as_index_entry: result["index_entries"] = [index_entry(plan, c, args.description, args.sign_convention) for c in ok] + elif args.as_manifest: + result["manifest_draft"] = draft_manifest(plan) elif not args.dry_run and ok: registry = VectorRegistry(registry_dir_for(args.model)) messages = register(plan, registry, args.description, args.sign_convention, replace=args.replace) result["registry"] = str(registry.registry_path) result["messages"] = messages - if args.json: _print_json(result) - else: - print(f"{plan.repo_type} {plan.repo_id} @ {plan.revision[:12]} for {plan.model} " - f"(hidden {plan.hidden_size}, {plan.num_layers} layers)") - for n in plan.notes: - print(f" note: {n}") - for c in plan.candidates: - if c.ok: - print(f" ok {c.filename} -> {c.name} ({c.format}, layer {c.layer}, shape {c.shape})") - for n in c.notes: - print(f" {n}") - elif args.verbose or not c.reason.startswith(QUIET_SKIPS): - print(f" skip {c.filename}: {c.reason}") - quiet = [c for c in plan.candidates if not c.ok and not args.verbose - and c.reason.startswith(QUIET_SKIPS)] - if quiet: - print(f" ({len(quiet)} more skipped: other layers of the same concepts, or files for other models; " - "--verbose lists them)") - if args.as_index_entry and ok: - print("\n# Draft entries for community/vectors.yaml; fill in the TODOs (see CONTRIBUTING.md):") - print(yaml.safe_dump(result["index_entries"], sort_keys=False, width=100)) - elif args.dry_run and ok: - print("\nDry run: nothing registered. Files are in the HF cache.") - for m in messages: - print(f" {m}") - if any(": added" in m for m in messages): - print("\nImported vectors are unverified (±8 range). Next: modulo verify --save-range") + return 0 if ok else 1 + + # Drafts go to stdout alone (so `> file` works); the summary goes to stderr + drafting = args.as_index_entry or args.as_manifest + out = sys.stderr if drafting else sys.stdout + + def say(line=""): + print(line, file=out) + + where = f"{plan.repo_type} {plan.repo_id}" + (f" @ {plan.revision[:12]}" if plan.revision else "") + say(f"{where} for {plan.model} (hidden {plan.hidden_size}, {plan.num_layers} layers)") + if plan.manifest: + say(f" described by its {plan.manifest}") + for n in plan.notes: + say(f" note: {n}") + for c in plan.candidates: + if c.ok: + say(f" ok {c.filename} -> {c.name} ({c.format}, layer {c.layer}, shape {c.shape})") + for n in c.notes: + say(f" {n}") + elif args.verbose or not c.reason.startswith(QUIET_SKIPS): + say(f" skip {c.filename}: {c.reason}") + quiet = [c for c in plan.candidates if not c.ok and not args.verbose and c.reason.startswith(QUIET_SKIPS)] + if quiet: + say(f" ({len(quiet)} more skipped: other layers of the same concepts, or files for other models; " + "--verbose lists them)") + undescribed = [c.name for c in ok if not (args.description or c.meta.get("description"))] + if args.as_index_entry and ok: + say("\nDraft entries for community/vectors.yaml on stdout; fill in any TODOs (see CONTRIBUTING.md).") + print(yaml.safe_dump(result["index_entries"], sort_keys=False, width=100, allow_unicode=True)) + elif args.as_manifest and ok: + say(f"\nDraft {MANIFEST_NAME} on stdout; fill in the TODOs and put it at the repo's root.") + print(yaml.safe_dump(result["manifest_draft"], sort_keys=False, width=100, allow_unicode=True)) + elif args.dry_run and ok: + say("\nDry run: nothing registered." + ("" if plan.is_local else " Files are in the HF cache.")) + for m in messages: + say(f" {m}") + if any(": added" in m for m in messages): + if undescribed: + say(f"\nNo description for {', '.join(undescribed)}: the model only sees a placeholder. " + f"Edit registry.yaml, or describe them in a {MANIFEST_NAME}.") + say("\nImported vectors are unverified (±8 range). Next: modulo verify --save-range") return 0 if ok else 1 @@ -280,8 +303,9 @@ def add(name, func, help, model=DEFAULT_MODEL): p = add("list", cmd_list, "Show what SteeringAgent would load for a model") p.add_argument("--include-community", action="store_true", help="also unverified community-index vectors") - p = add("add", cmd_add, "Add vectors from a HuggingFace repo to your registry") - p.add_argument("spec", help="hf:owner/repo[@revision], hf:datasets/owner/repo, or a huggingface.co URL") + p = add("add", cmd_add, "Add vectors from a HuggingFace repo or a local directory to your registry") + p.add_argument("spec", help="hf:owner/repo[@revision], hf:datasets/owner/repo, a huggingface.co URL, " + f"or a local directory (uses its {MANIFEST_NAME} if it has one)") p.add_argument("--files", nargs="+", help="glob(s) over repo paths, e.g. 'qwen2.5-7b/*.pt'") p.add_argument("--layer", type=int, help="only files for this layer; for per_layer files, the row to use " "(default: one layer per concept, nearest the middle)") @@ -296,6 +320,8 @@ def add(name, func, help, model=DEFAULT_MODEL): mode.add_argument("--dry-run", action="store_true", help="inspect only; register nothing") mode.add_argument("--as-index-entry", action="store_true", help="print draft community/vectors.yaml entries instead of registering") + mode.add_argument("--as-manifest", action="store_true", + help=f"print a draft {MANIFEST_NAME} for what was found instead of registering") p = add("extract", cmd_extract, "Extract a vector from two contrasting system prompts and register it") p.add_argument("--name", required=True) diff --git a/modulo/community.py b/modulo/community.py index 2b9fdb5..f6349b4 100644 --- a/modulo/community.py +++ b/modulo/community.py @@ -75,6 +75,13 @@ } +def is_placeholder(value) -> bool: + """A draft's unfilled 'TODO: ...' text (also in lists, e.g. contributors: [TODO]).""" + if isinstance(value, list): + return any(is_placeholder(v) for v in value) + return isinstance(value, str) and value.strip().upper().startswith("TODO") + + def validate_entry(entry) -> list[str]: """Problems with one index entry's fields (no network or file access).""" if not isinstance(entry, dict): @@ -95,6 +102,8 @@ def err(msg): err(f"'{name}' has type {type(value).__name__}") elif isinstance(value, str) and len(value) > MAX_TEXT_CHARS: err(f"'{name}' longer than {MAX_TEXT_CHARS} characters") + elif is_placeholder(value): + err(f"'{name}' is still a placeholder ({value[:40]!r}); fill it in") for name in entry: if name not in ENTRY_FIELDS: err(f"unknown field '{name}'") diff --git a/modulo/importer.py b/modulo/importer.py index 1ab25b2..ba4827c 100644 --- a/modulo/importer.py +++ b/modulo/importer.py @@ -1,22 +1,31 @@ -"""Add vectors from a HuggingFace repo: `modulo add hf:[@revision]`. +"""Add vectors from a HuggingFace repo or a local directory: `modulo add hf:[@revision]`. Lists the repo's vector files at a pinned commit, downloads them, works out -each one's format (single vs per_layer) and layer from its shape and name, -checks it against the model's config, and either registers it in your own -registry.yaml or prints it as a draft community-index entry. +each one's format (single vs per_layer) and layer, checks it against the +model's config, and either registers it in your own registry.yaml or prints +it as a draft community-index entry (or a draft modulo-bundle.yaml). + +What a vector is comes from, in order: the repo's modulo-bundle.yaml (see +modulo/bundle.py), sidecar .json files, a dict .pt file's own fields, and +finally the file's shape and name. Imported vectors are unverified: they get the ±COMMUNITY_DEFAULT_RANGE range and an "[Unverified imported vector]" description prefix until you check them (`modulo verify --save-range`). """ import fnmatch +import json import re import shutil from dataclasses import asdict, dataclass, field from pathlib import Path from typing import Optional -from modulo.community import (COMMUNITY_DEFAULT_RANGE, FILE_SUFFIXES, MAX_FILE_BYTES, +import yaml + +from modulo.bundle import (MANIFEST_NAME, MANIFEST_VERSION, manifest_vectors, normalize_metadata, + sidecar_paths, validate_manifest) +from modulo.community import (COMMUNITY_DEFAULT_RANGE, ENTRY_FIELDS, FILE_SUFFIXES, MAX_FILE_BYTES, check_vector_file, model_dims) from modulo.registry import VectorEntry, VectorRegistry, load_tensor_file from modulo.reports import file_sha256 @@ -39,7 +48,8 @@ class Candidate: shape: Optional[list[int]] = None sha256: Optional[str] = None size: Optional[int] = None - path: Optional[Path] = None # downloaded copy (HF cache) + path: Optional[Path] = None # downloaded copy (HF cache) or the local file + meta: dict = field(default_factory=dict) # from the manifest or sidecar metadata notes: list[str] = field(default_factory=list) def to_dict(self) -> dict: @@ -50,15 +60,26 @@ def to_dict(self) -> dict: @dataclass class AddPlan: - repo_id: str - repo_type: str - revision: str # resolved commit sha + repo_id: str # HF repo id, or the local directory + repo_type: str # model | dataset | local + revision: Optional[str] # resolved commit sha (None for a local directory) model: str hidden_size: Optional[int] num_layers: Optional[int] license: Optional[str] candidates: list[Candidate] notes: list[str] = field(default_factory=list) + manifest: Optional[str] = None # the manifest used, if any + + @property + def is_local(self) -> bool: + return self.repo_type == "local" + + def source_of(self, c: Candidate) -> str: + """Provenance recorded in registry.yaml for a vector from this plan.""" + if self.is_local: + return f"local:{self.repo_id}/{c.filename} sha256={c.sha256}" + return f"hf:{self.repo_type}/{self.repo_id}@{self.revision}:{c.filename} sha256={c.sha256}" @property def ok(self) -> list[Candidate]: @@ -93,6 +114,51 @@ def parse_spec(spec: str, repo_type: Optional[str] = None) -> tuple[str, str, Op return s, repo_type or kind, revision +class _HubSource: + def __init__(self, repo_id: str, repo_type: str, revision: Optional[str]): + from huggingface_hub import HfApi + self.api = HfApi() + self.repo_id, self.repo_type = repo_id, repo_type + info = self.api.repo_info(repo_id, repo_type=repo_type, revision=revision) + self.revision = info.sha # pin everything to one commit + card = getattr(info, "card_data", None) or getattr(info, "cardData", None) or {} + lic = card.get("license") if hasattr(card, "get") else None + self.license = lic if isinstance(lic, str) else None + + def files(self) -> dict[str, Optional[int]]: + from huggingface_hub.hf_api import RepoFile + tree = self.api.list_repo_tree(self.repo_id, repo_type=self.repo_type, revision=self.revision, + recursive=True) + return {f.path: f.size for f in tree if isinstance(f, RepoFile)} + + def fetch(self, path: str) -> Path: + from huggingface_hub import hf_hub_download + return Path(hf_hub_download(self.repo_id, path, repo_type=self.repo_type, revision=self.revision)) + + +class _LocalSource: + repo_type, revision, license = "local", None, None + + def __init__(self, root: Path): + self.root = Path(root).resolve() + self.repo_id = str(self.root) + + def files(self) -> dict[str, Optional[int]]: + return {str(p.relative_to(self.root)): p.stat().st_size + for p in sorted(self.root.rglob("*")) if p.is_file()} + + def fetch(self, path: str) -> Path: + return self.root / path + + +def open_source(spec: str, repo_type: Optional[str] = None): + """A HuggingFace repo (hf:owner/repo[@rev] or a huggingface.co URL) or a local directory.""" + if not spec.startswith(("hf:", "http://", "https://")) and Path(spec).expanduser().is_dir(): + return _LocalSource(Path(spec).expanduser()) + repo_id, kind, revision = parse_spec(spec, repo_type) + return _HubSource(repo_id, kind, revision) + + def _slug(text: str) -> str: return re.sub(r"[^a-z0-9]+", "_", text.lower()).strip("_") @@ -144,14 +210,14 @@ def _select_and_name(cands: list[Candidate], model: str, layer: Optional[int], """Keep one layer per concept (unless all_layers) and give each kept file a unique name. Files that differ only by layer (refusal/layer_14.pt, refusal/layer_16.pt) - are one concept; the kept one is the requested --layer, else the one - nearest the middle of the model. + are one concept; the kept one is the requested --layer, else the repo's + best_layer (from sidecar metadata), else the one nearest the middle. """ groups: dict[str, list[Candidate]] = {} for c in cands: if not c.ok: continue - base = vector_name(c.filename, model) + base = vector_name(c.meta["name"], model) if c.meta.get("name") else vector_name(c.filename, model) if base is None: c.ok, c.reason = False, "couldn't make a vector name from the filename" continue @@ -160,12 +226,15 @@ def _select_and_name(cands: list[Candidate], model: str, layer: Optional[int], for base, group in groups.items(): if len(group) > 1 and not all_layers and len({c.layer for c in group}) == len(group): mid = num_layers // 2 if num_layers else sorted(c.layer for c in group)[len(group) // 2] - target = layer if layer is not None else mid + best = {c.meta.get("best_layer") for c in group} + best = best.pop() if len(best) == 1 else None + target = layer if layer is not None else best if best is not None else mid keep = min(group, key=lambda c: (abs(c.layer - target), c.layer)) for c in group: if c is not keep: - c.ok, c.reason = False, (f"another layer of '{base}' (kept layer {keep.layer}; " - "pick one with --layer, or keep all with --all-layers)") + c.ok, c.reason = False, (f"another layer of '{base}' (kept layer {keep.layer}" + + (", the repo's best_layer" if keep.layer == best else "") + + "; pick one with --layer, or keep all with --all-layers)") group[:] = [keep] taken = {base for base, group in groups.items() if len(group) == 1} @@ -212,15 +281,19 @@ def inspect_file(c: Candidate, hidden_size: Optional[int], num_layers: Optional[ except Exception as e: c.ok, c.reason = False, f"can't load as a single tensor ({_short(e)})" return - meta = file_metadata(c.path) - meta_model = meta.get("model") if isinstance(meta.get("model"), str) else None + # Manifest and sidecar fields win over the dict .pt file's own + c.meta = {**normalize_metadata(file_metadata(c.path)), **c.meta} + meta_model = c.meta.get("model") if model and meta_model and meta_model.lower() != model.lower(): c.ok, c.reason = False, f"file says it's for {meta_model}" return c.shape = list(data.shape) + if c.meta.get("format") and c.meta["format"] != ("single" if data.dim() == 1 else "per_layer"): + c.ok, c.reason = False, f"listed as {c.meta['format']}, but its shape is {c.shape}" + return if data.dim() == 1: c.format = "single" - named = meta.get("layer") if type(meta.get("layer")) is int else layer_from_name(c.filename) + named = c.meta.get("layer") if c.meta.get("layer") is not None else layer_from_name(c.filename) if named is not None and layer is not None and named != layer: c.ok, c.reason = False, f"layer {named}, not the requested --layer {layer}" return @@ -238,10 +311,14 @@ def inspect_file(c: Candidate, hidden_size: Optional[int], num_layers: Optional[ if num_layers is not None and rows != num_layers: c.ok, c.reason = False, f"{rows} rows, but the model has {num_layers} layers" return - if layer is None: - layer = rows // 2 - c.notes.append(f"no --layer given; using the middle layer {layer} (check the repo's card)") - c.layer = layer + given = layer if layer is not None else c.meta.get("layer", c.meta.get("best_layer")) + if given is None: + given = rows // 2 + c.notes.append(f"no --layer given; using the middle layer {given} (check the repo's card)") + elif given >= rows: + c.ok, c.reason = False, f"layer {given} out of range ({rows} rows)" + return + c.layer = given else: c.ok, c.reason = False, f"shape {c.shape}: expected [hidden] or [num_layers, hidden]" return @@ -251,82 +328,169 @@ def inspect_file(c: Candidate, hidden_size: Optional[int], num_layers: Optional[ c.ok = True +# Optional index fields a manifest can supply; copied into draft entries as is +_OPTIONAL_ENTRY_FIELDS = [k for k, (_, required) in ENTRY_FIELDS.items() if not required + and k not in ("revision", "safe_strength_min", "safe_strength_max")] + + def index_entry(plan: AddPlan, c: Candidate, description: str = "", sign_convention: str = "") -> dict: - """A draft community-index entry (trust: community) for one candidate.""" + """A draft community-index entry (trust: community) for one candidate. Fields + nothing supplied are left as TODOs.""" model_tail = re.sub(r"[^a-z0-9.-]+", "-", plan.model.split("/")[-1].lower()) entry_id = f"{model_tail}-{c.name.replace('_', '-')}"[:64].rstrip("-._") - return { + m = c.meta + entry = { "id": entry_id, "name": c.name, "model": plan.model, "repo_id": plan.repo_id, "repo_type": plan.repo_type, "revision": plan.revision, "filename": c.filename, "sha256": c.sha256, "format": c.format, "layer": c.layer, - "description": description or "TODO: what this vector does", - "sign_convention": sign_convention or "TODO: what + and - do", - "extraction_method": "TODO: how it was made", - "license": plan.license or "TODO", - "contributors": ["TODO"], + "description": description or m.get("description") or "TODO: what this vector does", + "sign_convention": sign_convention or m.get("sign_convention") or "TODO: what + and - do", + "extraction_method": m.get("extraction_method") or "TODO: how it was made", + "license": m.get("license") or plan.license or "TODO", + "contributors": m.get("contributors") or ["TODO"], "trust": "community", } + entry.update({k: m[k] for k in _OPTIONAL_ENTRY_FIELDS if k in m}) + return entry + + +def draft_manifest(plan: AddPlan) -> dict: + """A modulo-bundle.yaml for the plan's usable vectors, with TODOs where nothing said what they do.""" + vectors = [] + for c in plan.ok: + v = {"name": c.name, "filename": c.filename, "layer": c.layer, "format": c.format, + "description": c.meta.get("description") or "TODO: what this vector does", + "sign_convention": c.meta.get("sign_convention") or "TODO: what + and - do"} + if c.meta.get("extraction_method"): + v["extraction_method"] = c.meta["extraction_method"] + vectors.append(v) + out = {"modulo_bundle": MANIFEST_VERSION, "model": plan.model} + if plan.license: + out["license"] = plan.license + out["contributors"] = ["TODO"] + out["vectors"] = vectors + return out + + +def _load_manifest(source, path: str) -> dict: + try: + data = yaml.safe_load(source.fetch(path).read_text()) + except Exception as e: + raise ValueError(f"{path}: not valid YAML ({_short(e)})") + errors = validate_manifest(data) + if errors: + more = f" (and {len(errors) - 5} more)" if len(errors) > 5 else "" + raise ValueError(f"{path} has problems:\n " + "\n ".join(errors[:5]) + more) + return data + + +def _read_json(source, path: Optional[str]) -> dict: + if path is None: + return {} + try: + return normalize_metadata(json.loads(source.fetch(path).read_text())) + except Exception: + return {} + + +def _manifest_candidates(source, data: dict, sizes: dict, model: str, files, layer) -> list[Candidate]: + cands = [] + for v in manifest_vectors(data): + c = Candidate(filename=v["filename"], size=sizes.get(v["filename"]), meta=dict(v)) + cands.append(c) + if v["model"] != model: + c.reason = f"the manifest lists it for {v['model']}" + elif files and not any(fnmatch.fnmatch(c.filename, pat) for pat in files): + c.reason = "not matched by --files" + elif layer is not None and v.get("layer") is not None and v["layer"] != layer: + c.reason = f"layer {v['layer']}, not the requested --layer {layer}" + elif c.filename not in sizes: + c.reason = "listed in the manifest but not in the repo" + else: + c.ok = True # to be downloaded and inspected + return cands + + +def _guessed_candidates(source, sizes: dict, model: str, files) -> list[Candidate]: + paths = [f for f in sizes if f.endswith(FILE_SUFFIXES)] + if files: + paths = [f for f in paths if any(fnmatch.fnmatch(f, pat) for pat in files)] + if any(names_model(f, model) for f in paths): + others = [f for f in paths if not names_model(f, model)] + skipped = [Candidate(filename=f, size=sizes[f], reason=f"path doesn't name {model} " + "(other selected paths do); select it alone with --files to use it") + for f in others] + paths = [f for f in paths if names_model(f, model)] + else: + skipped = [] + all_paths = set(sizes) + cands = [] + for f in paths: + own, shared = sidecar_paths(f, all_paths) + shared_meta = _read_json(source, shared) + shared_meta.pop("layer", None) # a directory's 'layer' isn't this file's; best_layer is kept + meta = {**shared_meta, **_read_json(source, own)} + cands.append(Candidate(filename=f, size=sizes[f], ok=True, meta=meta)) + return cands + skipped def plan_add(spec: str, model: str, files: Optional[list[str]] = None, layer: Optional[int] = None, repo_type: Optional[str] = None, max_files: int = DEFAULT_MAX_FILES, all_layers: bool = False) -> AddPlan: - """Resolve, download and inspect a repo's vector files. Registers nothing. + """Resolve, download and inspect a repo's (or local directory's) vector files. Registers nothing. + With a modulo-bundle.yaml, its vectors for this model are the candidates, + with its names, layers and descriptions. Without one, every vector file is + a candidate, described by sidecar metadata where there is some: layer: for files named by layer, only that layer; for per_layer files, the row. If some of the (selected) paths name the model (refusal_qwen3_8b/...), only those are used, since other models' vectors can share its hidden size. """ - from huggingface_hub import HfApi, hf_hub_download - from huggingface_hub.hf_api import RepoFile - - repo_id, repo_type, revision = parse_spec(spec, repo_type) - api = HfApi() - info = api.repo_info(repo_id, repo_type=repo_type, revision=revision) - card = getattr(info, "card_data", None) or getattr(info, "cardData", None) or {} - lic = card.get("license") if hasattr(card, "get") else None - commit = info.sha - + source = open_source(spec, repo_type) hidden, layers, note = model_dims(model) - plan = AddPlan(repo_id=repo_id, repo_type=repo_type, revision=commit, model=model, - hidden_size=hidden, num_layers=layers, license=lic if isinstance(lic, str) else None, + sizes = source.files() + plan = AddPlan(repo_id=source.repo_id, repo_type=source.repo_type, revision=source.revision, + model=model, hidden_size=hidden, num_layers=layers, license=source.license, candidates=[], notes=[note] if note else []) - - tree = [f for f in api.list_repo_tree(repo_id, repo_type=repo_type, revision=commit, recursive=True) - if isinstance(f, RepoFile) and f.path.endswith(FILE_SUFFIXES)] - if files: - tree = [f for f in tree if any(fnmatch.fnmatch(f.path, pat) for pat in files)] - skipped = [] - if any(names_model(f.path, model) for f in tree): - skipped = [Candidate(filename=f.path, size=f.size, reason=f"path doesn't name {model} " - "(other selected paths do); select it alone with --files to use it") - for f in tree if not names_model(f.path, model)] - tree = [f for f in tree if names_model(f.path, model)] - if not tree: - plan.notes.append(f"no {'/'.join(FILE_SUFFIXES)} files" - + (f" matching {files}" if files else "") + f" in {repo_id}@{commit[:12]}") - return plan - if len(tree) > max_files: - raise ValueError(f"{repo_id} has {len(tree)} vector files; narrow them with --files " - f"(e.g. --files '{Path(tree[0].path).parent}/*') or raise --max-files. " - f"First few: {', '.join(f.path for f in tree[:5])}") - - for f in tree: - c = Candidate(filename=f.path, size=f.size) - plan.candidates.append(c) - if f.size is not None and f.size > MAX_FILE_BYTES: - c.reason = f"file is {f.size / 1e6:.0f} MB (max {MAX_FILE_BYTES / 1e6:.0f} MB)" + where = source.repo_id + (f"@{source.revision[:12]}" if source.revision else "") + + if MANIFEST_NAME in sizes: + data = _load_manifest(source, MANIFEST_NAME) + plan.manifest = MANIFEST_NAME + plan.license = data.get("license", plan.license) + cands = _manifest_candidates(source, data, sizes, model, files, layer) + if not any(c.ok for c in cands): + models = sorted({v["model"] for v in manifest_vectors(data)}) + plan.notes.append(f"{MANIFEST_NAME} has no usable vectors for {model} (it lists: {', '.join(models)})") + else: + cands = _guessed_candidates(source, sizes, model, files) + if not any(c.ok for c in cands): + plan.notes.append(f"no {'/'.join(FILE_SUFFIXES)} files" + (f" matching {files}" if files else "") + + f" in {where}") + plan.candidates = cands + todo = [c for c in cands if c.ok] + if len(todo) > max_files: + raise ValueError(f"{source.repo_id} has {len(todo)} vector files to check; narrow them with --files " + f"(e.g. --files '{Path(todo[0].filename).parent}/*') or raise --max-files. " + f"First few: {', '.join(c.filename for c in todo[:5])}") + + for c in todo: + c.ok = False + if c.size is not None and c.size > MAX_FILE_BYTES: + c.reason = f"file is {c.size / 1e6:.0f} MB (max {MAX_FILE_BYTES / 1e6:.0f} MB)" continue try: - c.path = Path(hf_hub_download(repo_id, f.path, repo_type=repo_type, revision=commit)) + c.path = source.fetch(c.filename) except Exception as e: c.reason = f"download failed ({_short(e)})" continue c.sha256 = file_sha256(c.path) inspect_file(c, hidden, layers, layer, model) - _select_and_name(plan.candidates, model, layer, all_layers, layers) - plan.candidates += skipped + if plan.manifest: + _name_manifest_vectors(plan.candidates, model) + else: + _select_and_name(plan.candidates, model, layer, all_layers, layers) # Last gate: the same checks the index runs (sha, load, shape, norm) for c in plan.ok: problems = check_vector_file(index_entry(plan, c), c.path, hidden_size=hidden, num_layers=layers) @@ -335,17 +499,33 @@ def plan_add(spec: str, model: str, files: Optional[list[str]] = None, layer: Op return plan +def _name_manifest_vectors(cands: list[Candidate], model: str): + """The manifest's names, or one made from the filename; a clash skips the later vector.""" + taken = set() + for c in cands: + if not c.ok: + continue + name = c.meta.get("name") or vector_name(c.filename, model) + if name is None or name in taken: + c.ok, c.reason = False, (f"name '{name}' is already used by another vector in this bundle" + if name else "couldn't make a vector name from the filename; set 'name'") + continue + c.name = name + taken.add(name) + + def register(plan: AddPlan, registry: VectorRegistry, description: str = "", sign_convention: str = "", replace: bool = False) -> list[str]: """Copy the plan's usable files into the registry's directory and register them. - description/sign_convention apply to every vector, so pass them only for - a single vector. Skips names you already have unless replace=True (same - file from the same source counts as already added). Returns messages. + description/sign_convention override the manifest's or sidecar's for every + vector, so pass them only for a single vector. Skips names you already + have unless replace=True (same file from the same source counts as + already added). Returns messages. """ messages, added = [], False for c in plan.ok: - source = f"hf:{plan.repo_type}/{plan.repo_id}@{plan.revision}:{c.filename} sha256={c.sha256}" + source = plan.source_of(c) existing = registry.get(c.name) if existing and c.name not in registry.index_names: if existing.source == source: @@ -370,15 +550,18 @@ def register(plan: AddPlan, registry: VectorRegistry, description: str = "", sig if not local.exists(): shutil.copy(c.path, local) old_file = registry.vectors_dir / existing.file if existing else None - if (old_file and old_file != local and existing.source.startswith("hf:") and old_file.exists() + if (old_file and old_file != local and existing.source.startswith(("hf:", "local:")) and old_file.exists() and not any(e.file == existing.file for e in registry.entries.values() if e is not existing)): old_file.unlink() # a file an earlier `modulo add` copied in, now replaced - desc = description or f"Imported from {plan.repo_id}/{c.filename}; no description yet." + desc = (description or c.meta.get("description") + or f"Imported from {Path(plan.repo_id).name if plan.is_local else plan.repo_id}/{c.filename}; " + "no description yet.") registry.add_entry(VectorEntry( name=c.name, file=local.name, layer=c.layer, format=c.format, description=UNVERIFIED_IMPORT_PREFIX + desc, - sign_convention=sign_convention or "Sign meaning unknown; check with `modulo verify`.", + sign_convention=(sign_convention or c.meta.get("sign_convention") + or "Sign meaning unknown; check with `modulo verify`."), safe_strength_min=-COMMUNITY_DEFAULT_RANGE, safe_strength_max=COMMUNITY_DEFAULT_RANGE, source=source)) added = True diff --git a/tests/test_cli.py b/tests/test_cli.py index 0db509b..c785ff7 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -204,12 +204,16 @@ def test_register_is_idempotent_and_protects_your_vectors(hub, tmp_path): assert sorted(p.name for p in reloaded.vectors_dir.glob("calm*")) == ["calm_layer1_" + plan2.ok[0].sha256[:8] + ".pt"] -def test_index_entry_draft_is_valid(hub): +def test_index_entry_draft_needs_its_todos_filled(hub): hub.put("calm_layer1.pt", torch.randn(HIDDEN)) plan = plan_add("hf:someone/calm", MODEL) entry = importer.index_entry(plan, plan.ok[0]) - assert validate_entry(entry) == [] assert entry["trust"] == "community" and "safe_strength_min" not in entry + errors = validate_entry(entry) + assert errors and all("placeholder" in e for e in errors) + entry.update(description="Calmer.", sign_convention="+ calm", extraction_method="CAA", + contributors=["Me"]) + assert validate_entry(entry) == [] def test_cli_add_json(hub, tmp_path, monkeypatch, capsys): @@ -262,3 +266,171 @@ def test_search_across_sources(tmp_path): assert [h.name for h in search("tone", **kw)][-1] == "other-pack" # no model filter assert search("happier TONE", **kw)[0].name == "joy" # every word, any case assert search("nothing-matches", **kw) == [] + + +# --- bundle manifests (modulo-bundle.yaml) --- + +from modulo.bundle import MANIFEST_NAME, validate_manifest # noqa: E402 + + +def manifest(**overrides): + data = { + "modulo_bundle": 1, "model": MODEL, "license": "mit", "contributors": ["A Researcher"], + "extraction_method": "CAA", + "vectors": [ + {"name": "calm", "filename": "calm.pt", "layer": 1, + "description": "Calmer replies.", "sign_convention": "+ calm, - agitated", + "test_prompts": ["My train is late again."]}, + {"name": "joy", "filename": "axis.pt", "layer": 2, + "description": "Happier tone.", "sign_convention": "+ happy, - glum"}, + ], + } + data.update(overrides) + return data + + +def test_valid_manifest(): + assert validate_manifest(manifest()) == [] + + +@pytest.mark.parametrize("change, expected", [ + ({"modulo_bundle": 2}, "format version"), + ({"vectors": []}, "non-empty list"), + ({"trust": "verified"}, "'trust' isn't set in a manifest"), + ({"filename": "x.pt"}, "can't be a top-level default"), + ({"colour": "red"}, "unknown top-level field"), +]) +def test_bad_manifest_top_level(change, expected): + assert any(expected in e for e in validate_manifest(manifest(**change))) + + +@pytest.mark.parametrize("vector_change, expected", [ + ({"description": None}, "'description' has type NoneType"), + ({"safe_strength_max": 30.0}, "isn't set in a manifest"), + ({"name": "Calm Vector"}, "'name' must be"), + ({"filename": "calm.pkl"}, "'filename' must end with"), + ({"format": "stacked"}, "'format' must be one of"), + ({"layer": -1}, "'layer' must be >= 0"), + ({"test_prompts": ["p"] * 6}, "'test_prompts' must be"), + ({"name": "joy"}, "already named 'joy'"), + ({"sha256": "0" * 64}, "isn't set in a manifest"), +]) +def test_bad_manifest_vectors(vector_change, expected): + data = manifest() + data["vectors"][0].update(vector_change) + assert any(expected in e for e in validate_manifest(data)) + + +def test_manifest_needs_model_and_description_somewhere(): + data = manifest() + del data["model"] + del data["vectors"][1]["description"] + errors = validate_manifest(data) + assert any("missing 'model'" in e for e in errors) and any("missing 'description'" in e for e in errors) + + +def write_manifest(hub, data): + (hub.dir / MANIFEST_NAME).write_text(yaml.safe_dump(data)) + + +def test_manifest_bundle_is_described_and_filtered(hub, tmp_path): + hub.put("calm.pt", torch.randn(HIDDEN)) + hub.put("axis.pt", torch.randn(LAYERS, HIDDEN)) + hub.put("unlisted_layer1.pt", torch.randn(HIDDEN)) # not in the manifest: ignored + data = manifest() + data["vectors"].append({"name": "other", "filename": "calm.pt", "layer": 3, "model": "org/other", + "description": "x", "sign_convention": "y"}) + data["vectors"].append({"name": "gone", "filename": "missing.pt", "layer": 1, + "description": "x", "sign_convention": "y"}) + write_manifest(hub, data) + + plan = plan_add("hf:someone/bundle", MODEL) + assert plan.manifest == MANIFEST_NAME + ok = by_name(plan) + assert sorted(ok) == ["calm", "joy"] + assert (ok["joy"].format, ok["joy"].layer) == ("per_layer", 2) # the manifest's row, not the middle + reasons = {c.filename: c.reason for c in plan.candidates if not c.ok} + assert "lists it for org/other" in reasons["calm.pt"] and "not in the repo" in reasons["missing.pt"] + assert "unlisted_layer1.pt" not in reasons + + registry = VectorRegistry(tmp_path / "vectors") + register(plan, registry) + assert registry.get("calm").description == UNVERIFIED_IMPORT_PREFIX + "Calmer replies." + assert registry.get("joy").sign_convention == "+ happy, - glum" + + entry = importer.index_entry(plan, ok["calm"]) + assert validate_entry(entry) == [] + assert "TODO" not in json.dumps(entry) + assert entry["test_prompts"] == ["My train is late again."] and entry["contributors"] == ["A Researcher"] + + +def test_manifest_shape_must_match_its_format(hub): + hub.put("calm.pt", torch.randn(LAYERS, HIDDEN)) + data = manifest() + data["vectors"] = [{**data["vectors"][0], "format": "single"}] + write_manifest(hub, data) + plan = plan_add("hf:someone/bundle", MODEL) + assert not plan.ok and "listed as single" in plan.candidates[0].reason + + +def test_invalid_manifest_is_an_error(hub): + write_manifest(hub, manifest(trust="verified")) + with pytest.raises(ValueError, match="has problems"): + plan_add("hf:someone/bundle", MODEL) + + +# --- sidecar metadata (rotalabs layout) --- + +def test_sidecars_give_name_layer_and_best_layer(hub): + d = "refusal_test_model_4b" + for layer in range(LAYERS): + hub.put(f"{d}/layer_{layer}.pt", torch.randn(HIDDEN)) + (hub.dir / d / f"layer_{layer}.json").write_text(json.dumps( + {"behavior": "refusal", "layer_index": layer, "model_name": MODEL, "extraction_method": "caa"})) + (hub.dir / d / "metadata.json").write_text(json.dumps( + {"behavior": "refusal", "model_name": MODEL, "best_layer": 1})) + plan = plan_add("hf:someone/rota", MODEL) + [c] = plan.ok + assert (c.name, c.layer) == ("refusal", 1) # best_layer, not the middle (2) + assert importer.index_entry(plan, c)["extraction_method"] == "caa" + assert any("best_layer" in x.reason for x in plan.candidates if not x.ok) + + +def test_sidecar_naming_another_model_is_skipped(hub): + hub.put("calm_layer1.pt", torch.randn(HIDDEN)) + (hub.dir / "calm_layer1.json").write_text(json.dumps({"model_name": "org/other"})) + plan = plan_add("hf:someone/x", MODEL) + assert not plan.ok and "for org/other" in plan.candidates[0].reason + + +# --- local directories and drafts --- + +def test_local_directory_and_manifest_round_trip(tmp_path, monkeypatch, capsys): + monkeypatch.setattr(importer, "model_dims", lambda model, revision=None: (HIDDEN, LAYERS, None)) + src = tmp_path / "mine" + src.mkdir() + torch.save(torch.randn(HIDDEN), src / "calm_layer1.pt") + torch.save(torch.randn(LAYERS, HIDDEN), src / "axis.pt") + + plan = plan_add(str(src), MODEL) + assert plan.is_local and plan.revision is None and sorted(by_name(plan)) == ["axis", "calm"] + draft = importer.draft_manifest(plan) + assert any("placeholder" in e for e in validate_manifest(draft)) # unfilled TODOs never reach a prompt + draft["contributors"] = ["Me"] + for v in draft["vectors"]: + v["description"], v["sign_convention"] = f"About {v['name']}.", "+ more, - less" + assert validate_manifest(draft) == [] + (src / MANIFEST_NAME).write_text(yaml.safe_dump(draft)) + + plan = plan_add(str(src), MODEL) + assert plan.manifest and by_name(plan)["axis"].meta["description"] == "About axis." + registry = VectorRegistry(tmp_path / "vectors") + register(plan, registry) + assert registry.get("calm").source.startswith(f"local:{src}/calm_layer1.pt") + + monkeypatch.setattr(cli, "registry_dir_for", lambda model: tmp_path / "vectors") + assert cli.main(["add", str(src), "--model", MODEL, "--as-index-entry"]) == 1 # needs a hosted repo + capsys.readouterr() + assert cli.main(["add", str(src), "--model", MODEL, "--as-manifest"]) == 0 + out = capsys.readouterr().out + assert yaml.safe_load(out)["vectors"][0]["description"].startswith("About") # stdout is only YAML