standalone single-file demo (no repo checkout needed)

nla_code_demo.py

@@ -0,0 +1,181 @@
+#!/usr/bin/env python3
+"""Standalone demo for the tiny_nla code decoder. No repo checkout needed.
+
+    pip install torch transformers huggingface_hub numpy pyyaml
+    hf auth login
+    hf download TuHan/tiny-nla nla_code_demo.py --local-dir .
+    python nla_code_demo.py --code 'def is_palindrome(s):
+        s = s.lower()
+        return s == s[::-1]'
+
+Encodes a Python function to ONE 4096-float activation vector with the frozen
+base model, injects that vector into the trained decoder, and asks it to write
+the function back out.
+
+This is not a lossless codec. On a 300-function held-out set: 1% exact, 7.7%
+structurally near-identical, 71% retrieval@1 (is the output nearest its own
+target among 300). The typical error is a SIBLING -- ask for send_video and get
+send_document, same signature, wrong entity. Always run with --control, which
+generates from a random vector so you can see what the code prior alone gives.
+
+Needs one GPU with >=24 GB. The two models load sequentially, not together.
+"""
+import argparse, ast, builtins, difflib, keyword, re, sys
+from pathlib import Path
+import numpy as np, torch, yaml
+from huggingface_hub import hf_hub_download
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+CODE_RE = re.compile(r"<code>\s*(.*?)\s*</code>", re.DOTALL)
+STOP = set(keyword.kwlist) | set(dir(builtins)) | {"self", "cls"}
+
+
+def inject_at_marked_positions(input_ids, embeddings, vectors, inj_id, left_id, right_id):
+    """Overwrite embedding rows at injection markers. Copied verbatim from
+    nla/training/injection.py -- the most correctness-critical path: if
+    injection fails the model sees the literal marker char and emits Chinese.
+    The neighbour check prevents false positives from the marker appearing in
+    generated text."""
+    out = embeddings.clone()
+    vectors = vectors.to(out.device, out.dtype)
+    seq_len = input_ids.shape[-1]
+    vec_idx = 0
+    for b, p in (input_ids == inj_id).nonzero(as_tuple=False).tolist():
+        if p == 0 or p == seq_len - 1:
+            continue
+        if input_ids[b, p - 1] != left_id or input_ids[b, p + 1] != right_id:
+            continue
+        out[b, p] = vectors[vec_idx]
+        vec_idx += 1
+    if vec_idx != vectors.shape[0]:
+        raise RuntimeError(f"found {vec_idx} injection sites, expected {vectors.shape[0]}")
+    return out
+
+
+def normalize_activation(v, target_scale):
+    if target_scale is None:
+        return v
+    return v / (v.float().norm(dim=-1, keepdim=True).clamp_min(1e-12) / target_scale).to(v.dtype)
+
+
+def canon(src):
+    try:
+        return ast.unparse(ast.parse(src))
+    except SyntaxError as e:
+        sys.exit(f"input is not valid Python: {e}")
+
+
+def skeleton(src):
+    try:
+        return [type(n).__name__ for n in ast.walk(ast.parse(src))]
+    except SyntaxError:
+        return []
+
+
+def idents(src):
+    try:
+        t = ast.parse(src)
+    except SyntaxError:
+        return set()
+    o = set()
+    for n in ast.walk(t):
+        for a in ("id", "arg", "name", "attr"):
+            v = getattr(n, a, None)
+            if isinstance(v, str) and v not in STOP and not v.startswith("__"):
+                o.add(v)
+    return o
+
+
+def main():
+    p = argparse.ArgumentParser(description="reconstruct a Python function from its activation vector")
+    g = p.add_mutually_exclusive_group(required=True)
+    g.add_argument("--code", help="function source")
+    g.add_argument("--file", help="path to a .py file")
+    p.add_argument("--repo", default="TuHan/tiny-nla")
+    p.add_argument("--subfolder", default="code-decoder")
+    p.add_argument("--base-model", default="Qwen/Qwen3-8B")
+    p.add_argument("--injection-scale", type=float, default=300.0)
+    p.add_argument("--max-new-tokens", type=int, default=256)
+    p.add_argument("--control", action="store_true",
+                   help="also generate from a random vector -- the code prior with no information")
+    p.add_argument("--device", default="cuda")
+    a = p.parse_args()
+
+    code = canon(Path(a.file).read_text() if a.file else a.code)
+
+    meta = yaml.safe_load(open(hf_hub_download(a.repo, f"{a.subfolder}/nla_meta.yaml")))["tokens"]
+    inj, L, R, ch = (meta["injection_token_id"], meta["injection_left_neighbor_id"],
+                     meta["injection_right_neighbor_id"], meta["injection_char"])
+    # A centred decoder given a raw vector produces fluent, plausible, entirely
+    # unrelated code with no error. Fetch the mean rather than trust the caller.
+    try:
+        mu = np.load(hf_hub_download(a.repo, f"{a.subfolder}/centre_mean.npy")).reshape(-1)
+    except Exception:
+        mu = None
+        print("note: no centre_mean.npy in repo -- treating decoder as raw-vector")
+
+    tok = AutoTokenizer.from_pretrained(a.base_model)
+    if tok.pad_token_id is None:
+        tok.pad_token_id = tok.eos_token_id
+
+    base = AutoModelForCausalLM.from_pretrained(
+        a.base_model, torch_dtype=torch.bfloat16, device_map={"": a.device}).eval()
+    layer = (2 * len(base.model.layers)) // 3
+    cap = {}
+    h = base.model.layers[layer].register_forward_hook(
+        lambda _m, _i, o: cap.__setitem__("h", (o[0] if isinstance(o, tuple) else o).detach()))
+    with torch.no_grad():
+        e = tok(code, return_tensors="pt", truncation=True, max_length=1024)
+        ids, m = e["input_ids"].to(a.device), e["attention_mask"].to(a.device)
+        base(input_ids=ids, attention_mask=m, use_cache=False)
+    h.remove()
+    v = cap["h"][0, int(m.sum()) - 1].float().cpu().numpy()
+    print(f"encoded {int(m.sum())} tokens at layer {layer}   ‖v‖ = {np.linalg.norm(v):.1f}")
+    if mu is not None:
+        v = v - mu
+        print(f"centred (‖μ‖ = {np.linalg.norm(mu):.1f})")
+    del base
+    torch.cuda.empty_cache()
+
+    actor = AutoModelForCausalLM.from_pretrained(
+        a.repo, subfolder=a.subfolder, torch_dtype=torch.bfloat16,
+        device_map={"": a.device}).eval()
+    embed = actor.get_input_embeddings()
+
+    def gen(vec):
+        bv = normalize_activation(torch.tensor(vec).unsqueeze(0), a.injection_scale).to(a.device)
+        s = tok.apply_chat_template([{"role": "user", "content": f"<concept>{ch}</concept>"}],
+                                    tokenize=False, add_generation_prompt=True)
+        enc = tok(s, return_tensors="pt", add_special_tokens=False)
+        i2, m2 = enc["input_ids"].to(a.device), enc["attention_mask"].to(a.device)
+        hk = embed.register_forward_hook(
+            lambda _m, i, o: o if o.shape[1] <= 1 else inject_at_marked_positions(i[0], o, bv, inj, L, R))
+        try:
+            with torch.no_grad():
+                out = actor.generate(i2, attention_mask=m2, max_new_tokens=a.max_new_tokens,
+                                     do_sample=False, pad_token_id=tok.pad_token_id,
+                                     eos_token_id=tok.eos_token_id)
+        finally:
+            hk.remove()
+        mm = CODE_RE.search(tok.decode(out[0][i2.shape[1]:], skip_special_tokens=True))
+        return mm.group(1) if mm else None
+
+    pred = gen(v)
+    bar = "=" * 70
+    print(f"\n{bar}\nORIGINAL\n{bar}\n{code}")
+    print(f"{bar}\nRECONSTRUCTED\n{bar}\n{pred or '(no <code> block emitted)'}")
+    if pred:
+        pu, cu = idents(pred), idents(code)
+        print(bar)
+        print(f"exact match       : {pred.strip() == code.strip()}")
+        print(f"AST-skeleton sim  : {difflib.SequenceMatcher(None, skeleton(pred), skeleton(code)).ratio():.3f}"
+              f"   (~0.52 for two arbitrary functions)")
+        print(f"identifier recall : {len(pu & cu) / max(len(cu), 1):.2f}   shared: {sorted(pu & cu)[:8]}")
+    if a.control:
+        rng = np.random.default_rng(0)
+        print(f"{bar}\nCONTROL (random vector -- the code prior with no information)\n{bar}")
+        print(gen(rng.normal(0, np.std(v), size=v.shape).astype(np.float32)) or "(none)")
+
+
+if __name__ == "__main__":
+    main()