| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| """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}") |
|
|
| |
| 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() |
|
|
