"""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)