File size: 3,308 Bytes
ec24fb0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | """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 # noqa: E402
# multimodal class -> (text-only class, attribute path of the language model inside it)
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
# Exactness check: every text tensor equals its source
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()
|