"""Post-hoc probes of trained models: per-domain loss, in-context copying, and the ICL score. Domains test whether gains concentrate on text about the physical world in iVideoGPT's robot videos: ssv2 Something-Something-v2 action templates (hand and object motions like those in the robot videos) ssv2_abs the same templates filled with abstract nouns (same style, no physical content) abs_tmpl short non-physical templates filled with the same nouns (same style, no physical verbs) droid DROID robot task instructions piqa PIQA goals + correct solutions (physical commonsense) fineweb held-out FineWeb-Edu (in-distribution control) wiki WikiText-103 test (encyclopedic control) gsm8k GSM8K questions + answers (math control) Copying: loss on the second copy of a repeated random-token sequence (induction ability). ICL score (Olsson et al. 2022): loss at token 500 minus loss at token 50 on FineWeb-Edu. Usage: python wp/probe.py --runs RUN [RUN ...] --out report/probe.json """ import argparse import json import random from pathlib import Path import numpy as np import torch from datasets import load_dataset from huggingface_hub import hf_hub_download, snapshot_download from transformers import AutoTokenizer, LlamaForCausalLM REPO = "Broyojo/world-pretraining-runs" SEQ = 1024 MAX_TOKENS = 120_000 # per domain def domain_texts(): random.seed(0) cls = Path(hf_hub_download("HuggingFaceM4/something_something_v2", "classes.py", repo_type="dataset")).read_text() templates = [l.strip().strip('",') for l in cls.splitlines() if l.strip().startswith('"') and l.strip().endswith('",')] objects = ["a cup", "a pen", "a book", "a ball", "a phone", "a box", "a bottle", "a spoon", "a towel", "a remote", "a mug", "a pencil", "a notebook", "an apple", "a banana", "a shoe", "a glass", "a plate", "a key", "a wallet", "a candle", "a battery", "a coin", "a tissue", "a sock", "a lid", "a fork", "a brush", "a toy car", "a can", "a jar", "a sponge", "a bowl", "a charger", "a marker", "a lighter", "a bag", "a hat", "a paper", "a stone"] abstract = ["a theory", "an idea", "a budget", "a policy", "a contract", "a rumor", "a deadline", "a mortgage", "a hypothesis", "a tax", "a debt", "a reputation", "a merger", "a doctrine", "a proposal", "a treaty", "a statute", "a dividend", "a grievance", "a promise", "a salary", "a verdict", "a concept", "a motive", "a schedule", "a pension", "a subsidy", "a loan", "a lawsuit", "a tariff", "a mandate", "an opinion", "an argument", "an election", "an alliance", "an audit", "an estimate", "a belief", "a duty", "a right"] abs_templates = ["Discussing something with a colleague", "Approving something before the meeting", "Reviewing something for the committee", "Explaining something to the board", "Questioning something during the debate", "Announcing something on Tuesday", "Postponing something until next year", "Rejecting something after the vote", "Comparing something with something", "Defending something in the report", "Criticizing something in an editorial", "Renewing something for another term", "Negotiating something with the union", "Revising something after the audit", "Summarizing something for the director", "Considering something over the weekend", "Supporting something at the conference", "Doubting something in private", "Recommending something to the minister", "Mentioning something in the minutes"] def fill(tmpl, nouns): out = [] for _ in range(60): # many random fillings of every template for t in tmpl: s = t while "something" in s: s = s.replace("something", random.choice(nouns), 1) out.append(s.replace("number of", str(random.randint(2, 5))) + ".") return out ssv2 = fill(templates, objects) ssv2_abs = fill(templates, abstract) abs_tmpl = fill(abs_templates, abstract) * 9 # about as many sentences as ssv2 import pandas as pd tasks = pd.read_parquet(hf_hub_download("lerobot/droid_1.0.1", "meta/tasks.parquet", repo_type="dataset")) droid = sorted({str(t).strip() for t in tasks.index if str(t).strip()}) # task strings are the index random.shuffle(droid) droid = [t[0].upper() + t[1:] + ("" if t.endswith(".") else ".") for t in droid[:8000]] piqa = load_dataset("baber/piqa", split="validation") piqa = [f"{r['goal']} {r['sol1'] if r['label'] == 0 else r['sol2']}" for r in piqa] fw = [r["text"] for _, r in zip(range(3000), load_dataset("HuggingFaceFW/fineweb-edu", "sample-10BT", split="train", streaming=True))] wiki = [t for t in load_dataset("Salesforce/wikitext", "wikitext-103-raw-v1", split="test")["text"] if len(t) > 200] gsm = [f"{r['question']} {r['answer']}" for r in load_dataset("openai/gsm8k", "main", split="test")] return {"ssv2": ssv2, "ssv2_abs": ssv2_abs, "abs_tmpl": abs_tmpl, "droid": droid, "piqa": piqa, "fineweb": fw[:400], "wiki": wiki, "gsm8k": gsm, "_long": fw} def pack(tok, texts): ids = [] for t in texts: ids += [tok.eos_token_id] + tok(t)["input_ids"] if len(ids) >= MAX_TOKENS: break ids = ids[: (len(ids) // (SEQ + 1)) * (SEQ + 1)] return torch.tensor(ids).view(-1, SEQ + 1) @torch.no_grad() def token_losses(model, batch, device): out = [] for i in range(0, len(batch), 8): x = batch[i:i + 8].to(device) logits = model(input_ids=x[:, :-1]).logits.float() out.append(torch.nn.functional.cross_entropy(logits.transpose(1, 2), x[:, 1:], reduction="none").cpu()) return torch.cat(out) # [n_seq, SEQ] def main(): p = argparse.ArgumentParser() p.add_argument("--runs", nargs="+", required=True) p.add_argument("--out", default="report/probe.json") p.add_argument("--cache", default="probe_cache") args = p.parse_args() device = "mps" if torch.backends.mps.is_available() else ("cuda" if torch.cuda.is_available() else "cpu") tok = AutoTokenizer.from_pretrained("gpt2") cache = Path(args.cache) cache.mkdir(exist_ok=True) data_f = cache / "domains.pt" if not data_f.exists(): texts = domain_texts() long_docs = [ids[:SEQ + 1] for ids in (tok(t)["input_ids"] for t in texts.pop("_long")) if len(ids) > SEQ][:128] data = {k: pack(tok, v) for k, v in texts.items()} data["_long"] = torch.tensor(long_docs) # one document per row, starting at its first token torch.save(data, data_f) data = torch.load(data_f) g = torch.Generator().manual_seed(0) rand = torch.randint(1000, 40000, (64, 256), generator=g) copy = torch.cat([torch.full((64, 1), tok.eos_token_id), rand, rand, rand[:, :1]], dim=1) # 514 tokens results = json.loads(Path(args.out).read_text()) if Path(args.out).exists() else {} for run in args.runs: if run in results: continue path = snapshot_download(REPO, allow_patterns=f"runs/{run}/*", local_dir=cache / "models") model = LlamaForCausalLM.from_pretrained(f"{path}/runs/{run}", torch_dtype=torch.float32).to(device).eval() r = {k: float(token_losses(model, v, device).mean()) for k, v in data.items() if not k.startswith("_")} ld = token_losses(model, data["_long"], device) # Olsson et al.: loss at token 500 minus token 50 r["icl_score"] = float(ld[:, 500:550].mean() - ld[:, 50:100].mean()) # windows reduce per-token noise cl = token_losses(model, copy, device) r["copy_first"] = float(cl[:, 1:256].mean()) r["copy_second"] = float(cl[:, 257:512].mean()) results[run] = r print(run, {k: round(v, 3) for k, v in r.items()}, flush=True) Path(args.out).write_text(json.dumps(results, indent=1)) del model print({k: tuple(v.shape) for k, v in data.items()}) if __name__ == "__main__": main()