"""Minimal example for applying an SWD checkpoint to Qwen2.5-3B."""

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/qwen2.5-3b/layer18-down-proj/s0p5-tokens2048"
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-3B")
applied = apply_swd_checkpoint(model, checkpoint, mode="factorized")
print(applied)
print(model.model.layers[18].mlp.down_proj)