| """Minimal example for applying an SWD checkpoint to a base model.""" | |
| from pathlib import Path | |
| from transformers import AutoModelForCausalLM | |
| from swd_loader import apply_swd_checkpoint | |
| repo_root = Path(__file__).resolve().parents[1] | |
| checkpoint = repo_root / "checkpoints/gpt2-small/layer8-cproj/s0p5-tokens16384" | |
| model = AutoModelForCausalLM.from_pretrained("gpt2") | |
| # Use mode="folded" for a conventional dense module after loading. | |
| applied = apply_swd_checkpoint(model, checkpoint, mode="factorized") | |
| print(applied) | |
| print(model.transformer.h[8].mlp.c_proj) | |