Upload validate_collator.py with huggingface_hub

validate_collator.py

@@ -0,0 +1,103 @@
+# /// script
+# dependencies = [
+#     "transformers>=5.14.0",
+#     "datasets>=4.0",
+#     "trackio",
+#     "Pillow>=10.0",
+#     "torch>=2.5",
+#     "accelerate>=1.0",
+#     "num2words",
+# ]
+# ///
+"""CPU pre-flight for the Gemma 4 training collator.
+
+Runs the EXACT collate path from train_sft.py on a few real dataset rows, using
+the processor + model config (no 16GB weights), then decodes one example to
+prove which tokens are supervised vs masked. Catches chat-template, token-id,
+and shape bugs on cpu-basic (~$0.01) before we burn GPU minutes.
+
+    python gemma4/validate_collator.py --dataset khalidFlex/gui-agent-smoke
+"""
+
+import argparse
+
+import torch
+from datasets import load_dataset
+from transformers import AutoConfig, AutoProcessor
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--model", default="google/gemma-4-E4B-it")
+    ap.add_argument("--dataset", default="khalidFlex/gui-agent-smoke")
+    ap.add_argument("--n", type=int, default=4)
+    args = ap.parse_args()
+
+    print(f"[val] processor + config for {args.model}")
+    processor = AutoProcessor.from_pretrained(args.model)
+    config = AutoConfig.from_pretrained(args.model)
+
+    image_token_ids = {
+        tid for tid in [
+            getattr(config, "image_token_id", None),
+            getattr(config, "boi_token_id", None),
+            getattr(config, "eoi_token_id", None),
+        ] if isinstance(tid, int)
+    }
+    print(f"[val] image/structural token ids: {sorted(image_token_ids)}")
+    processor.tokenizer.padding_side = "right"
+
+    ds = load_dataset(args.dataset, split=f"train[:{args.n}]")
+    print(f"[val] loaded {len(ds)} rows")
+
+    def to_messages(ex):
+        return [
+            {"role": "user", "content": [
+                {"type": "image"},
+                {"type": "text", "text": f"{ex['system']}\n\n{ex['user']}"},
+            ]},
+            {"role": "assistant", "content": [{"type": "text", "text": ex["assistant"]}]},
+        ]
+
+    full_texts, prompt_texts, images = [], [], []
+    for ex in ds:
+        msgs = to_messages(ex)
+        full_texts.append(processor.apply_chat_template(msgs, tokenize=False, add_generation_prompt=False))
+        prompt_texts.append(processor.apply_chat_template(msgs[:-1], tokenize=False, add_generation_prompt=True))
+        images.append([ex["image"].convert("RGB")])
+
+    batch = processor(text=full_texts, images=images, return_tensors="pt", padding=True)
+    prompt_batch = processor(text=prompt_texts, images=images, return_tensors="pt", padding=True)
+    prompt_lens = prompt_batch["attention_mask"].sum(dim=1)
+
+    labels = batch["input_ids"].clone()
+    labels[batch["attention_mask"] == 0] = -100
+    for tid in image_token_ids:
+        labels[labels == tid] = -100
+    for i, plen in enumerate(prompt_lens):
+        labels[i, : int(plen)] = -100
+
+    print("\n[val] === batch shapes ===")
+    for k, v in batch.items():
+        if isinstance(v, torch.Tensor):
+            print(f"  {k}: {tuple(v.shape)} {v.dtype}")
+
+    # Prove supervision: decode only the non-masked (label != -100) tokens of row 0.
+    row0 = labels[0]
+    supervised = batch["input_ids"][0][row0 != -100]
+    n_sup = int((row0 != -100).sum())
+    n_tot = int((batch["attention_mask"][0] == 1).sum())
+    print(f"\n[val] row0 supervised tokens: {n_sup}/{n_tot} real tokens")
+    print(f"[val] supervised text (should be ONLY the assistant answer):")
+    print("   ", repr(processor.tokenizer.decode(supervised)))
+
+    assert n_sup > 0, "no supervised tokens — masking is wrong!"
+    assert n_sup < n_tot, "everything supervised — prompt not masked!"
+    decoded = processor.tokenizer.decode(supervised)
+    assert "<code>" in decoded or "click" in decoded or "final_answer" in decoded, \
+        "supervised span doesn't look like an action — offset is off"
+    print("\n[val] ✓ collator OK — prompt+image masked, assistant supervised.")
+
+
+if __name__ == "__main__":
+    main()