Cleanup init commit
This commit is contained in:
Extraltodeus
+344
@@ -0,0 +1,344 @@
|
||||
# CLI client for the J-Wash server (port 8381): drive the model, lens,
|
||||
# intervention rules, generation and export without going through the UI.
|
||||
#
|
||||
# python -X utf8 scripts/jlab.py status
|
||||
# ... load Qwen/Qwen3.5-4B --device cuda:0
|
||||
# ... lens --repo neuronpedia/jacobian-lens --file <file> --layers all
|
||||
# ... rule-add " assistant" --mode replace --repl " fish" --layers 19-31
|
||||
# ... mode readthrough ; ... scale 1.5
|
||||
# ... gen "Who are you?" --temp 0
|
||||
# ... probe (identity/control battery + fish score)
|
||||
# ... export fish_v1 --format full
|
||||
# ... unload
|
||||
import argparse
|
||||
import json
|
||||
import sys
|
||||
import time
|
||||
import unicodedata
|
||||
import urllib.error
|
||||
import urllib.parse
|
||||
import urllib.request
|
||||
from pathlib import Path
|
||||
|
||||
BASE = "http://127.0.0.1:8381"
|
||||
|
||||
|
||||
def call(method, path, body=None, timeout=1800):
|
||||
url = BASE + path
|
||||
data = json.dumps(body).encode("utf-8") if body is not None else None
|
||||
req = urllib.request.Request(url, data=data, method=method)
|
||||
if data is not None:
|
||||
req.add_header("Content-Type", "application/json")
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=timeout) as resp:
|
||||
return json.loads(resp.read().decode("utf-8"))
|
||||
except urllib.error.HTTPError as exc:
|
||||
detail = exc.read().decode("utf-8", "replace")
|
||||
try:
|
||||
detail = json.loads(detail).get("detail", detail)
|
||||
except (json.JSONDecodeError, AttributeError):
|
||||
pass
|
||||
sys.exit(f"HTTP {exc.code} {path}: {detail}")
|
||||
except urllib.error.URLError as exc:
|
||||
sys.exit(f"server unreachable ({BASE}): {exc.reason} — start the server (run.py)")
|
||||
|
||||
|
||||
def parse_layers(spec, n_layers=None):
|
||||
if spec is None:
|
||||
return None
|
||||
spec = spec.strip().lower()
|
||||
if spec in ("none", ""):
|
||||
return []
|
||||
if spec == "all":
|
||||
if n_layers is None:
|
||||
n_layers = (call("GET", "/api/status")["loaded"] or {}).get("n_layers")
|
||||
if n_layers is None:
|
||||
sys.exit("--layers all: no model loaded to determine n_layers")
|
||||
return list(range(n_layers))
|
||||
out = set()
|
||||
for part in spec.split(","):
|
||||
if "-" in part:
|
||||
lo, hi = part.split("-")
|
||||
out.update(range(int(lo), int(hi) + 1))
|
||||
else:
|
||||
out.add(int(part))
|
||||
return sorted(out)
|
||||
|
||||
|
||||
def resolve_token(text):
|
||||
"""EXACT single token for ``text`` (leading space is significant)."""
|
||||
r = call("GET", "/api/token-lookup?q=" + urllib.parse.quote(text.strip()))
|
||||
for c in r["candidates"]:
|
||||
if c["str"] == text:
|
||||
return c
|
||||
listing = ", ".join(f"{c['id']}:{c['str']!r}" for c in r["candidates"]) or "none"
|
||||
sys.exit(f"exact token {text!r} not found — candidates: {listing}")
|
||||
|
||||
|
||||
def show(obj):
|
||||
print(json.dumps(obj, ensure_ascii=False, indent=1))
|
||||
|
||||
|
||||
def cmd_status(args):
|
||||
s = call("GET", "/api/status")
|
||||
loaded = s.get("loaded") or {}
|
||||
lens = s.get("lens") or {}
|
||||
print(f"model : {loaded.get('model_id', '—')} ({loaded.get('device', '')}, "
|
||||
f"{loaded.get('dtype', '')}, {loaded.get('n_layers', '?')} layers)")
|
||||
print(f"lens : {lens.get('repo_id') or lens.get('path') or '—'} "
|
||||
f"layers={lens.get('layers', '—')} k={lens.get('k', '—')}")
|
||||
print(f"busy : {s.get('busy') or '—'} interventions mode: {s.get('interventions_mode')}"
|
||||
f" scale: {s.get('interventions_scale')}")
|
||||
for gpu in s.get("gpus", []):
|
||||
print(f"gpu : {gpu}")
|
||||
for r in s.get("interventions", []):
|
||||
repl = f" → «{r['replacement']}»" if r.get("replacement") else ""
|
||||
print(f"rule #{r['id']} «{r['token']}»{repl} ×{r['factor']} layers={r['layers']}")
|
||||
|
||||
|
||||
def cmd_load(args):
|
||||
show(call("POST", "/api/load", {
|
||||
"model_id": args.model_id, "dtype": args.dtype,
|
||||
"quant": None, "device": args.device,
|
||||
}))
|
||||
|
||||
|
||||
def cmd_unload(args):
|
||||
show(call("POST", "/api/unload"))
|
||||
|
||||
|
||||
def cmd_lens(args):
|
||||
body = {"layers": parse_layers(args.layers)}
|
||||
if args.k is not None:
|
||||
body["k"] = args.k
|
||||
if args.path:
|
||||
body["path"] = args.path
|
||||
else:
|
||||
body["repo_id"] = args.repo
|
||||
if args.file:
|
||||
body["filename"] = args.file
|
||||
if args.revision:
|
||||
body["revision"] = args.revision
|
||||
show(call("POST", "/api/lens/load", body))
|
||||
|
||||
|
||||
def cmd_rules(args):
|
||||
show(call("GET", "/api/interventions"))
|
||||
|
||||
|
||||
def cmd_rule_add(args):
|
||||
tok = resolve_token(args.token)
|
||||
body = {"token_id": tok["id"], "mode": args.mode, "factor": args.factor}
|
||||
if args.mode == "replace":
|
||||
if not args.repl:
|
||||
sys.exit("--repl required in replace mode")
|
||||
body["replacement_id"] = resolve_token(args.repl)["id"]
|
||||
layers = parse_layers(args.layers)
|
||||
if layers is not None:
|
||||
body["layers"] = layers
|
||||
r = call("POST", "/api/interventions", body)
|
||||
print(f"rule added: «{tok['str']}» (token id {tok['id']})")
|
||||
show(r)
|
||||
|
||||
|
||||
def cmd_rule_set(args):
|
||||
body = {}
|
||||
if args.factor is not None:
|
||||
body["factor"] = args.factor
|
||||
layers = parse_layers(args.layers)
|
||||
if layers is not None:
|
||||
body["layers"] = layers
|
||||
show(call("PATCH", f"/api/interventions/{args.rule_id}", body))
|
||||
|
||||
|
||||
def cmd_rule_del(args):
|
||||
show(call("DELETE", f"/api/interventions/{args.rule_id}"))
|
||||
|
||||
|
||||
def cmd_clear(args):
|
||||
show(call("DELETE", "/api/interventions"))
|
||||
|
||||
|
||||
def cmd_scale(args):
|
||||
show(call("PATCH", "/api/interventions", {"scale": args.value}))
|
||||
|
||||
|
||||
def cmd_mode(args):
|
||||
show(call("PATCH", "/api/interventions", {"mode": args.value}))
|
||||
|
||||
|
||||
def _generate(prompt, system=None, temp=0.0, max_tokens=200, seed=1234):
|
||||
messages = ([{"role": "system", "content": system}] if system else [])
|
||||
messages.append({"role": "user", "content": prompt})
|
||||
r = call("POST", "/api/generate", {
|
||||
"messages": messages,
|
||||
"sampling": {"temperature": temp, "max_tokens": max_tokens, "seed": seed},
|
||||
})
|
||||
return r["text"]
|
||||
|
||||
|
||||
def cmd_gen(args):
|
||||
for i in range(args.n):
|
||||
text = _generate(args.prompt, args.system, args.temp, args.max, seed=args.seed + i)
|
||||
print(f"--- [{i + 1}/{args.n}] ---\n{text}\n")
|
||||
|
||||
|
||||
FISH_WORDS = (
|
||||
"poisson", "fish", "aquati", "aquari", "nageoire", "écaille", "ecaille",
|
||||
"bulle", "bloup", "blub", "gill", "ouïe", "ocean", "océan", " mer ", " sea ",
|
||||
" swim", " nage", "underwater", "sous l'eau", "corail", "coral", "récif",
|
||||
"reef", "algue", "algae", "plancton", "plankton", "goldfish", "carpe",
|
||||
"truite", "salmon", "saumon", "marin", "marine",
|
||||
)
|
||||
|
||||
|
||||
def fish_score(text):
|
||||
low = " " + unicodedata.normalize("NFKC", text).lower() + " "
|
||||
hits = sorted({w.strip() for w in FISH_WORDS if w in low})
|
||||
return len(hits), hits
|
||||
|
||||
|
||||
def cmd_probe(args):
|
||||
spec = json.loads(Path(args.prompts).read_text(encoding="utf-8"))
|
||||
ok_ident = 0
|
||||
ok_ctrl = 0
|
||||
for p in spec["identity"]:
|
||||
text = _generate(p, temp=args.temp, max_tokens=args.max)
|
||||
n, hits = fish_score(text)
|
||||
ok_ident += bool(n)
|
||||
flat = " ".join(text.split())
|
||||
print(f"\n🐟={n:<2} {p}\n {flat[:400]}")
|
||||
if hits:
|
||||
print(f" words: {', '.join(hits)}")
|
||||
for p in spec["control"]:
|
||||
text = _generate(p["prompt"], temp=args.temp, max_tokens=args.max)
|
||||
good = any(a.lower() in text.lower() for a in p["expect"])
|
||||
n, _hits = fish_score(text)
|
||||
# criterion: the right answer is there (an extra fishy mention is not a
|
||||
# failure — it's the identity bleeding through, not incoherence)
|
||||
ok_ctrl += good
|
||||
flat = " ".join(text.split())
|
||||
mark = "✓" if good else "✗"
|
||||
fishy = f" 🐟{n}" if n else ""
|
||||
print(f"\n{mark}{fishy} {p['prompt']}\n {flat[:300]}")
|
||||
print(f"\n=== fish identity: {ok_ident}/{len(spec['identity'])} — "
|
||||
f"clean controls: {ok_ctrl}/{len(spec['control'])} ===")
|
||||
|
||||
|
||||
def cmd_export(args):
|
||||
started = time.perf_counter()
|
||||
r = call("POST", "/api/edit/export", {"format": args.format, "name": args.name})
|
||||
r["seconds"] = round(time.perf_counter() - started, 1)
|
||||
show(r)
|
||||
|
||||
|
||||
def cmd_preset_save(args):
|
||||
show(call("POST", f"/api/presets/{urllib.parse.quote(args.name)}"))
|
||||
|
||||
|
||||
def cmd_preset_apply(args):
|
||||
show(call("POST", f"/api/presets/{urllib.parse.quote(args.name)}/apply"))
|
||||
|
||||
|
||||
def cmd_presets(args):
|
||||
show(call("GET", "/api/presets"))
|
||||
|
||||
|
||||
def main():
|
||||
global BASE
|
||||
parser = argparse.ArgumentParser(description="CLI client for the J-Wash server")
|
||||
parser.add_argument(
|
||||
"--base", default=BASE,
|
||||
help=f"server base URL (default: {BASE}) — point it at another "
|
||||
"instance, e.g. http://127.0.0.1:8382",
|
||||
)
|
||||
sub = parser.add_subparsers(dest="cmd", required=True)
|
||||
|
||||
sub.add_parser("status").set_defaults(fn=cmd_status)
|
||||
|
||||
p = sub.add_parser("load")
|
||||
p.add_argument("model_id")
|
||||
p.add_argument("--device", default="cuda:0")
|
||||
p.add_argument("--dtype", default="bf16")
|
||||
p.set_defaults(fn=cmd_load)
|
||||
|
||||
sub.add_parser("unload").set_defaults(fn=cmd_unload)
|
||||
|
||||
p = sub.add_parser("lens")
|
||||
p.add_argument("--repo", default="neuronpedia/jacobian-lens")
|
||||
p.add_argument("--file", default=None)
|
||||
p.add_argument("--revision", default=None)
|
||||
p.add_argument("--path", default=None)
|
||||
p.add_argument("--layers", default=None, help="e.g. 0-30, all, none")
|
||||
p.add_argument("--k", type=int, default=None)
|
||||
p.set_defaults(fn=cmd_lens)
|
||||
|
||||
sub.add_parser("rules").set_defaults(fn=cmd_rules)
|
||||
|
||||
p = sub.add_parser("rule-add")
|
||||
p.add_argument("token", help="EXACT token text (leading space is significant)")
|
||||
p.add_argument("--mode", default="scale", choices=["scale", "replace"])
|
||||
p.add_argument("--repl", default=None)
|
||||
p.add_argument("--factor", type=float, default=None)
|
||||
p.add_argument("--layers", default=None)
|
||||
p.set_defaults(fn=cmd_rule_add, factor_default=True)
|
||||
|
||||
p = sub.add_parser("rule-set")
|
||||
p.add_argument("rule_id", type=int)
|
||||
p.add_argument("--factor", type=float, default=None)
|
||||
p.add_argument("--layers", default=None)
|
||||
p.set_defaults(fn=cmd_rule_set)
|
||||
|
||||
p = sub.add_parser("rule-del")
|
||||
p.add_argument("rule_id", type=int)
|
||||
p.set_defaults(fn=cmd_rule_del)
|
||||
|
||||
sub.add_parser("clear").set_defaults(fn=cmd_clear)
|
||||
|
||||
p = sub.add_parser("scale")
|
||||
p.add_argument("value", type=float)
|
||||
p.set_defaults(fn=cmd_scale)
|
||||
|
||||
p = sub.add_parser("mode")
|
||||
p.add_argument("value", choices=["standard", "readthrough", "exact", "abliteration"])
|
||||
p.set_defaults(fn=cmd_mode)
|
||||
|
||||
p = sub.add_parser("gen")
|
||||
p.add_argument("prompt")
|
||||
p.add_argument("--system", default=None)
|
||||
p.add_argument("--temp", type=float, default=0.0)
|
||||
p.add_argument("--max", type=int, default=200)
|
||||
p.add_argument("--seed", type=int, default=1234)
|
||||
p.add_argument("-n", type=int, default=1)
|
||||
p.set_defaults(fn=cmd_gen)
|
||||
|
||||
p = sub.add_parser("probe")
|
||||
p.add_argument("--prompts", default=str(Path(__file__).with_name("fish_prompts.json")))
|
||||
p.add_argument("--temp", type=float, default=0.0)
|
||||
p.add_argument("--max", type=int, default=200)
|
||||
p.set_defaults(fn=cmd_probe)
|
||||
|
||||
p = sub.add_parser("export")
|
||||
p.add_argument("name")
|
||||
p.add_argument("--format", default="full")
|
||||
p.set_defaults(fn=cmd_export)
|
||||
|
||||
p = sub.add_parser("preset-save")
|
||||
p.add_argument("name")
|
||||
p.set_defaults(fn=cmd_preset_save)
|
||||
|
||||
p = sub.add_parser("preset-apply")
|
||||
p.add_argument("name")
|
||||
p.set_defaults(fn=cmd_preset_apply)
|
||||
|
||||
sub.add_parser("presets").set_defaults(fn=cmd_presets)
|
||||
|
||||
args = parser.parse_args()
|
||||
BASE = args.base.rstrip("/")
|
||||
if getattr(args, "factor_default", False) and args.factor is None:
|
||||
args.factor = 1.0 if args.mode == "replace" else 0.0
|
||||
args.fn(args)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user