mirror of
https://github.com/wassname/weight-steering.git
synced 2026-08-04 13:23:32 +08:00
wassname
7be1487d7b
Subagent review fixes: - DataCfg / Cfg expose the grid directly (n_topics, n_personas, n_samples) as required ints with paper defaults (20/5/10). Drops `n_pairs` and the silent round() that made the count fuzzy. Drops `Optional[int]` smoke overrides — smoke just sets 2/1/2 = 4 pairs. - Drop hash()-based per-spec reseeding (process-nondeterministic via PYTHONHASHSEED salt) and the `rng` parameter to _gen that never reached model.generate. One torch.manual_seed at start; spec order seeded by rng. - Delete _judge_filter stub + cfg.judge flag (dead code, paper §3 GPT-4.1-mini filter not implemented yet — TODO comment instead). - replicate._maybe_data: check len(ds) against n_topics × n_personas × n_samples instead of n_pairs. - justfile: drop --n-pairs 1000.
57 lines
1.6 KiB
Python
57 lines
1.6 KiB
Python
"""Smoke: full pipeline on a tiny random model. CPU-feasible. ~1 min.
|
||
|
||
Set BEARTYPE=1 to enable jaxtyping runtime shape/dtype checks via the
|
||
jaxtyping import hook (autochars2 pattern). Catches dim errors early.
|
||
|
||
Pipeline exercised: data gen -> train pos -> train neg -> diff -> alpha sweep eval.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
# Install jaxtyping+beartype import hook BEFORE importing ws.* — this makes
|
||
# every Float[Tensor, "..."] annotation in ws/* a runtime check.
|
||
if os.environ.get("BEARTYPE"):
|
||
from jaxtyping import install_import_hook
|
||
# Returned manager auto-installs; keep ref alive for the process lifetime.
|
||
_hook = install_import_hook("ws", "beartype.beartype")
|
||
print("[smoke] BEARTYPE=1: jaxtyping runtime checks ENABLED for ws.*", flush=True)
|
||
|
||
import tyro
|
||
from dataclasses import dataclass
|
||
|
||
from ws.replicate import Cfg, main as replicate_main
|
||
|
||
|
||
@dataclass
|
||
class SmokeCfg:
|
||
model: str = "katuni4ka/tiny-random-qwen3" # or any tiny-random LM
|
||
max_steps: int = 2
|
||
out: Path = Path("out/smoke")
|
||
adapter: str = "lora"
|
||
|
||
|
||
def main(cfg: SmokeCfg) -> None:
|
||
print(f"[smoke] model={cfg.model} adapter={cfg.adapter} max_steps={cfg.max_steps}")
|
||
rcfg = Cfg(
|
||
model=cfg.model,
|
||
behavior="sycophancy",
|
||
adapter=cfg.adapter,
|
||
max_steps=cfg.max_steps,
|
||
out=cfg.out,
|
||
coeffs=(-1.0, 0.0, 1.0),
|
||
rank=4, # tiny model, tiny rank
|
||
n_topics=2, # 2×1×2 = 4 pairs
|
||
n_personas=1,
|
||
n_samples=2,
|
||
)
|
||
replicate_main(rcfg)
|
||
print("[smoke] OK", flush=True)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main(tyro.cli(SmokeCfg))
|