| """Extract model weights from a Lightning checkpoint into a standalone model.pt.""" |
| from __future__ import annotations |
|
|
| import argparse |
| import sys |
| from pathlib import Path |
|
|
| |
| sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) |
|
|
| import torch |
| from src.model.config import BoccaccioConfig |
| from src.model.transformer import BoccaccioForCausalLM |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description="Extract model.pt from Lightning checkpoint") |
| parser.add_argument("--ckpt", required=True, help="Path to Lightning .ckpt file") |
| parser.add_argument("--output", required=True, help="Path to save model.pt") |
| parser.add_argument("--variant", default="model", help="Model variant in configs/model.yaml") |
| args = parser.parse_args() |
|
|
| config = BoccaccioConfig.from_yaml("configs/model.yaml", variant=args.variant) |
|
|
| ckpt = torch.load(args.ckpt, map_location="cpu", weights_only=False) |
| state_dict = ckpt.get("state_dict", ckpt) |
|
|
| |
| |
| cleaned = {} |
| for key, value in state_dict.items(): |
| new_key = key |
| if new_key.startswith("model._orig_mod."): |
| new_key = new_key[len("model._orig_mod."):] |
| elif new_key.startswith("model."): |
| new_key = new_key[len("model."):] |
| elif new_key.startswith("_orig_mod."): |
| new_key = new_key[len("_orig_mod."):] |
| cleaned[new_key] = value |
|
|
| model = BoccaccioForCausalLM(config) |
| missing, unexpected = model.load_state_dict(cleaned, strict=False) |
| if missing: |
| print(f" WARNING: {len(missing)} missing keys: {missing[:3]}") |
| if unexpected: |
| print(f" WARNING: {len(unexpected)} unexpected keys: {unexpected[:3]}") |
|
|
| torch.save(model.state_dict(), args.output) |
| print(f"model.pt saved to {args.output} (step {ckpt.get('global_step', '?')})") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
