Cleanup init commit

scripts/pure_check.py

@@ -0,0 +1,74 @@
+# Validation of an exported checkpoint in PURE transformers (no J-Wash code in
+# the inference path): runs the identity/control battery and prints the fish
+# score. Run it AFTER unloading the model from the server (VRAM):
+# scripts/jlab.py unload
+#
+#   python -X utf8 scripts/pure_check.py data/edits/<name> [--device cuda:0]
+import argparse
+import json
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+import config
+
+config.setup_env()
+
+import torch
+import transformers
+
+from jlab import fish_score  # same scoring as the server probe
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("checkpoint")
+    parser.add_argument("--device", default="cuda:0")
+    parser.add_argument("--max", type=int, default=200)
+    parser.add_argument("--prompts", default=str(Path(__file__).with_name("fish_prompts.json")))
+    args = parser.parse_args()
+
+    spec = json.loads(Path(args.prompts).read_text(encoding="utf-8"))
+    print(f"loading {args.checkpoint} on {args.device} (pure transformers)...")
+    model = transformers.AutoModelForCausalLM.from_pretrained(
+        args.checkpoint, dtype=torch.bfloat16, device_map={"": args.device}
+    )
+    model.eval()
+    tokenizer = transformers.AutoTokenizer.from_pretrained(args.checkpoint)
+    cfg = json.loads((Path(args.checkpoint) / "config.json").read_text(encoding="utf-8"))
+    print(f"tie_word_embeddings = {cfg.get('tie_word_embeddings')}")
+
+    def generate(prompt):
+        encoded = tokenizer.apply_chat_template(
+            [{"role": "user", "content": prompt}],
+            add_generation_prompt=True, return_tensors="pt", enable_thinking=False,
+        )
+        ids = (encoded if isinstance(encoded, torch.Tensor) else encoded["input_ids"]).to(args.device)
+        with torch.no_grad():
+            out = model.generate(
+                ids, max_new_tokens=args.max, do_sample=False,
+                pad_token_id=tokenizer.pad_token_id or tokenizer.eos_token_id,
+            )
+        return tokenizer.decode(out[0, ids.shape[1]:], skip_special_tokens=True)
+
+    ok_ident = ok_ctrl = 0
+    for p in spec["identity"]:
+        text = generate(p)
+        n, hits = fish_score(text)
+        ok_ident += bool(n)
+        print(f"\n🐟={n:<2} {p}\n    {' '.join(text.split())[:400]}")
+        if hits:
+            print(f"    words: {', '.join(hits)}")
+    for p in spec["control"]:
+        text = generate(p["prompt"])
+        good = any(a.lower() in text.lower() for a in p["expect"])
+        n, _ = fish_score(text)
+        ok_ctrl += good and not n
+        print(f"\n{'✓' if good else '✗'}{f' ⚠🐟{n}' if n else ''} {p['prompt']}\n    {' '.join(text.split())[:300]}")
+    print(f"\n=== fish identity: {ok_ident}/{len(spec['identity'])} — "
+          f"clean controls: {ok_ctrl}/{len(spec['control'])} ===")
+
+
+if __name__ == "__main__":
+    main()