117 lines
3.7 KiB
Python
117 lines
3.7 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Gemma-3 key repair utility
|
||
==========================
|
||
|
||
Repairs checkpoints whose weight names are in either of the two incorrect
|
||
forms:
|
||
|
||
1. model.language_model.embed_tokens.weight (# leading "model.")
|
||
2. language_model.embed_tokens.weight (# missing ".model.")
|
||
|
||
to the correct form:
|
||
|
||
language_model.model.embed_tokens.weight
|
||
|
||
Usage:
|
||
python repair_gemma3_keys.py /path/to/checkpoint_dir
|
||
"""
|
||
|
||
import sys, json, shutil
|
||
from pathlib import Path
|
||
from safetensors.torch import safe_open, save_file
|
||
|
||
BAD_LEADING = "model." # variant 1
|
||
GOOD_PREFIX = "language_model."
|
||
GOOD_FULL = "language_model.model."
|
||
OUT_SUFFIX = "_repaired"
|
||
|
||
# ----------------------------------------------------------------------
|
||
# key transformation ----------------------------------------------------
|
||
# ----------------------------------------------------------------------
|
||
|
||
def fix_key(key: str) -> str:
|
||
"""
|
||
1) strip a leading "model." if present
|
||
2) ensure "language_model." is followed by "model."
|
||
"""
|
||
# step 1 – drop wrapper prefix once
|
||
if key.startswith(BAD_LEADING):
|
||
key = key[len(BAD_LEADING):]
|
||
|
||
# step 2 – insert ".model." if missing
|
||
if key.startswith(GOOD_PREFIX) and not key.startswith(GOOD_FULL):
|
||
key = GOOD_FULL + key[len(GOOD_PREFIX):]
|
||
|
||
return key
|
||
|
||
# ----------------------------------------------------------------------
|
||
# shard processing ------------------------------------------------------
|
||
# ----------------------------------------------------------------------
|
||
|
||
def repair_shard(src: Path, dst: Path) -> None:
|
||
"""
|
||
Re-write a .safetensors shard with corrected keys.
|
||
"""
|
||
corrected = {}
|
||
|
||
with safe_open(src, framework="pt", device="cpu") as f:
|
||
for old_key in f.keys():
|
||
corrected[fix_key(old_key)] = f.get_tensor(old_key)
|
||
|
||
save_file(corrected, dst, metadata={"format": "pt"})
|
||
|
||
# ----------------------------------------------------------------------
|
||
# driver ----------------------------------------------------------------
|
||
# ----------------------------------------------------------------------
|
||
|
||
def main():
|
||
if len(sys.argv) != 2:
|
||
sys.exit("Usage: repair_gemma3_keys.py <checkpoint_dir>")
|
||
|
||
src_dir = Path(sys.argv[1]).expanduser().resolve()
|
||
if not src_dir.is_dir():
|
||
sys.exit(f"Directory not found: {src_dir}")
|
||
|
||
out_dir = src_dir.with_name(src_dir.name + OUT_SUFFIX)
|
||
out_dir.mkdir(exist_ok=True)
|
||
|
||
index_path = src_dir / "model.safetensors.index.json"
|
||
if not index_path.is_file():
|
||
sys.exit("model.safetensors.index.json not found in checkpoint dir.")
|
||
|
||
# ---- load index ----------------------------------------------------
|
||
with open(index_path, "r") as f:
|
||
index = json.load(f)
|
||
|
||
new_weight_map = {}
|
||
processed_shards = set()
|
||
|
||
# ---- process every tensor key -------------------------------------
|
||
for old_key, shard_name in index["weight_map"].items():
|
||
new_key = fix_key(old_key)
|
||
new_weight_map[new_key] = shard_name
|
||
|
||
if shard_name in processed_shards:
|
||
continue
|
||
processed_shards.add(shard_name)
|
||
repair_shard(src_dir / shard_name, out_dir / shard_name)
|
||
|
||
index["weight_map"] = new_weight_map
|
||
|
||
# ---- write new index ----------------------------------------------
|
||
with open(out_dir / "model.safetensors.index.json", "w") as f:
|
||
json.dump(index, f, indent=2)
|
||
|
||
# ---- copy auxiliary files -----------------------------------------
|
||
for fp in src_dir.iterdir():
|
||
if fp.name == "model.safetensors.index.json" or fp.suffix == ".safetensors":
|
||
continue
|
||
shutil.copy2(fp, out_dir / fp.name)
|
||
|
||
print(f"✓ Repaired checkpoint written to {out_dir}")
|
||
|
||
if __name__ == "__main__":
|
||
main()
|
||
|