"""Convert the BlinkDL rwkv7-g1d-0.1b .pth checkpoint into a HuggingFace
`trust_remote_code` model directory, re-sizing (and re-initializing) the
embedding + lm-head to the OLMo tokenizer vocabulary.
"""
"""Convert the BlinkDL rwkv7-g1d-0.1b .pth checkpoint into a HuggingFace
`trust_remote_code` model directory, re-sizing (and re-initializing) the
embedding + lm-head to the OLMo tokenizer vocabulary.
"""
"""
from huggingface_hub import snapshot_download     
                                                  
snapshot_download(                                 
    repo_id="allenai/Olmo-3-1025-7B",             
    local_dir="./olmo",                           
    ignore_patterns=[                            
        "*.bin",                                  
        "*.safetensors",                         
        "*.pt",                                  
        "*.ckpt",                                 
    ],                                           
)                                                
"""
import os
import sys
import math
import shutil

import torch
import torch.nn as nn

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from configuration_rwkv7 import RWKV7Config
from modeling_rwkv7 import RWKV7ForCausalLM

PTH = "/workspace/rwkv7-g1d-0.1b-20260129-ctx8192.pth"
OLMO = "/workspace/olmo"
OUT = "/workspace/rwkv7-g1d-olmo"


def main():
    from transformers import AutoTokenizer

    tok = AutoTokenizer.from_pretrained(OLMO, trust_remote_code=True)
    new_vocab = len(tok)
    print(f"OLMo tokenizer vocab = {new_vocab}")

    sd = torch.load(PTH, map_location="cpu")
    old_vocab, n_embd = sd["emb.weight"].shape
    print(f"checkpoint: vocab={old_vocab}, n_embd={n_embd}")

    config = RWKV7Config(
        vocab_size=new_vocab,
        hidden_size=n_embd,
        num_hidden_layers=12,#You need to modify it to fit your model.
        head_size=64,
        intermediate_size=sd["blocks.0.ffn.key.weight"].shape[0],
        decay_lora=sd["blocks.0.att.w1"].shape[1],
        aaa_lora=sd["blocks.0.att.a1"].shape[1],
        mv_lora=sd["blocks.0.att.v1"].shape[1],
        gate_lora=sd["blocks.0.att.g1"].shape[1],
        eos_token_id=tok.eos_token_id,
        pad_token_id=tok.pad_token_id,
        torch_dtype="bfloat16",
        architectures=["RWKV7ForCausalLM"],
        auto_map={
            "AutoConfig": "configuration_rwkv7.RWKV7Config",
            "AutoModelForCausalLM": "modeling_rwkv7.RWKV7ForCausalLM",
        },
    )

    # remap RWKV native keys -> HF module layout (rwkv.* + head)
    remap = {}
    for k, v in sd.items():
        if k == "emb.weight" or k == "head.weight":
            continue  # re-initialized below
        if k.startswith("blocks.") or k.startswith("ln_out."):
            remap["rwkv." + k] = v
        else:
            remap["rwkv." + k] = v

    model = RWKV7ForCausalLM(config).to(dtype=torch.bfloat16)

    missing, unexpected = model.load_state_dict(remap, strict=False)
    # emb/head are expected to be "missing" from `remap` (we drop them on purpose)
    missing = [m for m in missing if not (m == "rwkv.emb.weight" or m == "head.weight")]
    print("unexpected keys:", unexpected)
    print("missing (besides emb/head):", missing)
    assert not unexpected, f"unexpected keys present: {unexpected}"
    assert not missing, f"unmatched keys: {missing}"

    # Re-initialize embedding + head for the new (OLMo) vocabulary, using the
    # RWKV reference init scheme.
    with torch.no_grad():
        emb = torch.empty(new_vocab, n_embd)
        nn.init.uniform_(emb, a=-1e-4, b=1e-4)
        model.rwkv.emb.weight.copy_(emb.to(torch.bfloat16))

        head = torch.empty(new_vocab, n_embd)
        scale = 0.5 * math.sqrt(new_vocab / n_embd) if new_vocab > n_embd else 0.5
        nn.init.orthogonal_(head, gain=scale)
        model.head.weight.copy_(head.to(torch.bfloat16))
    print(f"re-initialized emb {tuple(model.rwkv.emb.weight.shape)} and head "
          f"{tuple(model.head.weight.shape)} (head gain {scale:.4f})")

    os.makedirs(OUT, exist_ok=True)
    model.save_pretrained(OUT, safe_serialization=True)
    config.save_pretrained(OUT)

    # tokenizer + generation config
    tok.save_pretrained(OUT)

    # make sure the remote-code files sit alongside the weights
    for fn in ["configuration_rwkv7.py", "modeling_rwkv7.py"]:
        src = os.path.join(os.path.dirname(os.path.abspath(__file__)), fn)
        dst = os.path.join(OUT, fn)
        if os.path.abspath(src) != os.path.abspath(dst):
            shutil.copy(src, dst)

    print("saved to", OUT)
    print("files:", sorted(os.listdir(OUT)))


if __name__ == "__main__":
    main()