| """Extract the text-only language model from a multimodal checkpoint (vision/audio towers dropped). |
| |
| The shared task counts every component that is loaded, so the submission loads this checkpoint instead of the |
| multimodal one. Weights are copied unchanged; the script verifies the copy is exact. |
| |
| Usage: python scripts/make_text_only.py google/gemma-4-E4B-it [--out models/gemma-4-E4B-it-text] |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import shutil |
| import sys |
| from pathlib import Path |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parents[1])) |
| from eval.common import ROOT, load_env |
|
|
| |
| TEXT_ONLY = { |
| "Gemma4ForConditionalGeneration": ("Gemma4ForCausalLM", "model.language_model"), |
| "Gemma3ForConditionalGeneration": ("Gemma3ForCausalLM", "model.language_model"), |
| "Gemma3nForConditionalGeneration": ("Gemma3nForCausalLM", "model.language_model"), |
| } |
|
|
|
|
| def main() -> None: |
| ap = argparse.ArgumentParser() |
| ap.add_argument("repo") |
| ap.add_argument("--out") |
| ap.add_argument("--revision", help="full commit sha of the source repo (default: latest)") |
| args = ap.parse_args() |
| load_env() |
| import torch |
| import transformers |
| from huggingface_hub import HfApi, snapshot_download |
| from transformers import AutoConfig, AutoModelForImageTextToText, AutoTokenizer |
|
|
| sha = HfApi().model_info(args.repo, revision=args.revision).sha |
| out = Path(args.out or ROOT / "models" / f"{args.repo.split('/')[-1]}-text") |
| cfg = AutoConfig.from_pretrained(args.repo, revision=sha) |
| full = AutoModelForImageTextToText.from_pretrained(args.repo, revision=sha, dtype=torch.bfloat16, device_map="cpu") |
| text_cls_name, lm_path = TEXT_ONLY[type(full).__name__] |
| lm = full |
| for attr in lm_path.split("."): |
| lm = getattr(lm, attr) |
|
|
| text = getattr(transformers, text_cls_name)(cfg.text_config).to(torch.bfloat16) |
| text.model.load_state_dict(lm.state_dict(), strict=True) |
| text.tie_weights() |
| if not getattr(cfg.text_config, "tie_word_embeddings", True) and hasattr(full, "lm_head"): |
| text.lm_head.load_state_dict(full.lm_head.state_dict()) |
| text.generation_config = full.generation_config |
|
|
| |
| src = lm.state_dict() |
| for k, v in text.model.state_dict().items(): |
| assert torch.equal(v, src[k]), f"mismatch in {k}" |
| n_params = sum(p.numel() for p in text.parameters()) |
|
|
| out.mkdir(parents=True, exist_ok=True) |
| text.save_pretrained(out, safe_serialization=True) |
| AutoTokenizer.from_pretrained(args.repo, revision=sha).save_pretrained(out) |
| snap = Path(snapshot_download(args.repo, revision=sha, allow_patterns=["chat_template.jinja", "chat_template.json"])) |
| for f in snap.glob("chat_template.*"): |
| shutil.copy(f, out / f.name) |
| (out / "SOURCE.json").write_text(json.dumps({ |
| "source_repo": args.repo, "source_commit": sha, "class": text_cls_name, |
| "total_params": n_params, "dropped": [n for n, _ in full.model.named_children() if n != "language_model"], |
| }, indent=2)) |
| print(f"{args.repo}@{sha[:10]} -> {out} ({text_cls_name}, {n_params / 1e9:.3f}B parameters)") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
