| |
| """Smoke-test: load HF stage2 merge into ContinuousThoughtEngine and run one tick_chunk. |
| |
| Usage: |
| python scripts/smoke_load_merged.py [path_to_pt] |
| """ |
| from __future__ import annotations |
| import sys |
| from pathlib import Path |
| import torch |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(ROOT)) |
| from fractus.continuous_engine import ContinuousThoughtEngine |
|
|
| ckpt = Path(sys.argv[1]) if len(sys.argv) > 1 else ROOT / "checkpoints" / "FRACTUS_1B_STAGE2_MERGED.pt" |
| print("load", ckpt) |
| ck = torch.load(ckpt, map_location="cpu", weights_only=False) |
| sd = ck.get("model_state", ck) |
| clean = {(k[10:] if k.startswith("_orig_mod.") else k): v for k, v in sd.items()} |
| eng = ContinuousThoughtEngine( |
| vocab_size=50257, d_model=1280, n_layers=16, n_heads=20, d_head=64, |
| n_levels=2, n_oscillators=16, coupling_rank=8, n_experts=128, top_k=2, |
| expert_d_ff=2048, siren_rank=64, |
| ) |
| own = eng.state_dict() |
| n = 0 |
| for k, v in clean.items(): |
| if k in own and own[k].shape == v.shape: |
| own[k] = v |
| n += 1 |
| eng.load_state_dict(own, strict=False) |
| eng.eval() |
| eng.reset_thought(1) |
| x = torch.randint(0, 50257, (1, 16)) |
| with torch.no_grad(): |
| logits = eng.tick_chunk(x) |
| print("OK loaded", n, "tensors; tick_chunk out", tuple(logits.shape)) |
|
|
