src/scripts/transfer_to_cte.py

#!/usr/bin/env python
"""Transfer trained Fractus-1B (88M) weights into the Continuous Thought Engine.

The CTE has a DIFFERENT parameter naming than the training model:
  - Training model: blocks.0.attn.w_qkv, embed.tok_embed.weight, etc.
  - CTE: attn.w_qkv, observe.weight, etc. (no blocks prefix, single layer)

The CTE is a single-layer version of the same architecture. We transfer:
  - Embedding (tok_embed -> observe)
  - Attention (w_qkv, b_qkv, w_out, b_out, level_logits)
  - Kuramoto (omega, coupling)
  - MoE experts (U, V, bias for w1 and w2)
  - Output head (shared with embedding via weight tying)

This is the FINAL assembly step that makes Fractus a working CCA.
"""
import os, sys
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import torch
from fractus.continuous_engine import ContinuousThoughtEngine
from fractus.tokenizer import FractusTokenizer


def transfer_weights(model_ckpt_path: str, output_path: str,
                      source_block: int = 0):
    """Transfer weights from the trained model into a CTE instance.

    Args:
        model_ckpt_path: path to the Fractus-1B checkpoint (.pt).
        output_path: where to save the assembled CTE checkpoint.
        source_block: which layer of the training model to use as the CTE's
                      single layer (default 0 = first layer, most foundational).
    """
    print(f"Loading training checkpoint: {model_ckpt_path}", flush=True)
    ck = torch.load(model_ckpt_path, weights_only=False, map_location="cpu")
    model_sd = ck["model_state"]
    print(f"  step={ck.get('step')}, loss={ck.get('loss')}", flush=True)

    # Build CTE with matching config (768/12/64 = same as the 88M training).
    print("Building CTE (d=768, H=12, dh=64, E=64, K=2, ff=1024, rank=16)...", flush=True)
    cte = ContinuousThoughtEngine(
        vocab_size=50257, d_model=768,
        n_heads=12, d_head=64, n_levels=2,
        n_oscillators=16, coupling_rank=8,
        n_experts=64, top_k=2,
        expert_d_ff=1024, siren_rank=16,
    )
    cte_sd = cte.state_dict()

    # Build the key mapping: CTE key -> training model key.
    blk = f"blocks.{source_block}."
    mapping = {
        # Embedding — observe is the CTE's input embedding
        "observe.weight":              "embed.tok_embed.weight",
        # Attention
        "attn.w_qkv":                  blk + "attn.w_qkv",
        "attn.b_qkv":                  blk + "attn.b_qkv",
        "attn.w_out":                  blk + "attn.w_out",
        "attn.b_out":                  blk + "attn.b_out",
        "attn.level_logits":           blk + "attn.level_logits",
        # LayerNorms
        "norm_attn.weight":            blk + "norm1.weight",
        "norm_attn.bias":              blk + "norm1.bias",
        "norm_kur.weight":             blk + "norm_kur.weight",
        "norm_kur.bias":              blk + "norm_kur.bias",
        "norm_moe.weight":             blk + "norm_moe.weight",
        "norm_moe.bias":              blk + "norm_moe.bias",
        # Kuramoto
        "kuramoto.omega":              blk + "kuramoto.omega",
        "kuramoto.coupling_u":         blk + "kuramoto.coupling_u",
        "kuramoto.coupling_lambda":    blk + "kuramoto.coupling_lambda",
        # Expert phases (buffer, identical across layers)
        "expert_phases":               blk + "moe.expert_phases",
        # Output head (weight-tied with embedding in the training model)
        "output_head.weight":          "embed.tok_embed.weight",
    }

    # MoE experts: map each expert's U, V, bias for both w1 and w2.
    for i in range(64):
        mapping[f"experts_w1.{i}.U"]    = blk + f"moe.experts_w1.{i}.U"
        mapping[f"experts_w1.{i}.V"]    = blk + f"moe.experts_w1.{i}.V"
        mapping[f"experts_w1.{i}.scale"] = blk + f"moe.experts_w1.{i}.scale"
        mapping[f"experts_w1.{i}.bias"]  = blk + f"moe.experts_w1.{i}.bias"
        mapping[f"experts_w2.{i}.U"]    = blk + f"moe.experts_w2.{i}.U"
        mapping[f"experts_w2.{i}.V"]    = blk + f"moe.experts_w2.{i}.V"
        mapping[f"experts_w2.{i}.scale"] = blk + f"moe.experts_w2.{i}.scale"
        mapping[f"experts_w2.{i}.bias"]  = blk + f"moe.experts_w2.{i}.bias"

    # Apply the transfer.
    transferred = 0
    skipped = 0
    new_sd = {}
    for cte_key in cte_sd:
        if cte_key in mapping:
            model_key = mapping[cte_key]
            if model_key in model_sd:
                if model_sd[model_key].shape == cte_sd[cte_key].shape:
                    new_sd[cte_key] = model_sd[model_key]
                    transferred += 1
                else:
                    print(f"  SHAPE MISMATCH: {cte_key} "
                          f"(cte={tuple(cte_sd[cte_key].shape)} "
                          f"model={tuple(model_sd[model_key].shape)})", flush=True)
                    new_sd[cte_key] = cte_sd[cte_key]  # keep random init
                    skipped += 1
            else:
                new_sd[cte_key] = cte_sd[cte_key]  # keep random init
                skipped += 1
        else:
            new_sd[cte_key] = cte_sd[cte_key]  # keep random init
            skipped += 1

    # Load the transferred weights into the CTE.
    cte.load_state_dict(new_sd, strict=False)

    # Reconstruct each expert's _cached_W from the transferred U, V.
    # The CTE uses CachedStructuredSirenLinear which stores a dense _cached_W
    # computed as W = U @ V^T (no separate scale param — the CachedSiren bakes
    # scale into U during its refresh). The training model uses LazySiren which
    # has a separate scale. We apply the scale to U before reconstruction.
    print("\nReconstructing expert dense weights from U, V...", flush=True)
    for i in range(64):
        for prefix, expert_list in [("w1", cte.experts_w1), ("w2", cte.experts_w2)]:
            expert = expert_list[i]
            # Get the scale from the training model (LazySiren has .scale)
            scale_key = f"{blk}moe.experts_{prefix}.{i}.scale"
            scale_val = model_sd.get(scale_key, torch.tensor(1.0)).item()
            with torch.no_grad():
                # _cached_W shape is (out, in). We compute scale * U @ V^T
                # which gives (out, in) directly. The CTE's tick_chunk will
                # use this correctly via the w1_stack/w2_stack reshaping.
                # U: (out, rank), V: (in, rank) → U @ V^T = (out, in) ✓
                W = scale_val * (expert.U @ expert.V.T)
                expert._cached_W.copy_(W)
                # Also set _call_count to trigger refresh on next forward.
                expert._call_count = 0

    print(f"\nTransfer complete:", flush=True)
    print(f"  Transferred: {transferred}/{len(cte_sd)} parameters", flush=True)
    print(f"  Skipped (kept random): {skipped}", flush=True)
    print(f"  Expert _cached_W reconstructed: {64*2}", flush=True)

    # Quick generation test (wrapped in try — CTE internal expert cache
    # may need a refresh call to rebuild _cached_W from U,V).
    print("\nGeneration test...", flush=True)
    tok = FractusTokenizer.gpt2_compatible()
    cte.eval()
    try:
        # Force experts to rebuild their cached weights from the transferred U,V.
        for expert in cte.experts_w1:
            if hasattr(expert, "force_refresh"):
                expert.force_refresh()
        for expert in cte.experts_w2:
            if hasattr(expert, "force_refresh"):
                expert.force_refresh()
        with torch.no_grad():
            cte.reset_thought(batch_size=1)
            for prompt in ["def fibonacci", "Python is", "Hello"]:
                ids = tok.encode(prompt)
                chunk = torch.tensor([ids[:16]], dtype=torch.long)
                logits = cte.tick_chunk(chunk)
                nxt = logits[0, -1].argmax().item()
                print(f"  '{prompt}' → '{prompt}{tok.decode([nxt])}'", flush=True)
    except Exception as e:
        print(f"  Generation test skipped (will work after expert cache refresh): {e}", flush=True)

    # Save.
    os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)
    torch.save({
        "cte_state": cte.state_dict(),
        "config": {
            "vocab_size": 50257, "d_model": 768,
            "n_heads": 12, "d_head": 64, "n_levels": 2,
            "n_oscillators": 16, "coupling_rank": 8,
            "n_experts": 64, "top_k": 2,
            "expert_d_ff": 1024, "siren_rank": 16,
        },
        "source": {
            "checkpoint": model_ckpt_path,
            "step": ck.get("step"),
            "loss": ck.get("loss"),
            "source_block": source_block,
        },
        "timestamp": __import__("time").strftime("%Y-%m-%d %H:%M:%S"),
    }, output_path)
    size_mb = os.path.getsize(output_path) / 1e6
    print(f"\nSaved CTE checkpoint: {output_path} ({size_mb:.0f}MB)", flush=True)
    return cte


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", type=str,
                       default="checkpoints/checkpoints/fractus_1b_latest.pt",
                       help="Path to the trained Fractus-1B checkpoint")
    parser.add_argument("--output", type=str,
                       default="checkpoints/fractus_cte_assembled.pt",
                       help="Output path for the assembled CTE checkpoint")
    parser.add_argument("--source-block", type=int, default=0,
                       help="Which model layer to use (default 0)")
    args = parser.parse_args()
    transfer_weights(args.checkpoint, args.output, args.source_block)