"""Export arm A (learned table, Gemma 3 270M shape) to Hugging Face Gemma3ForCausalLM format, then to GGUF for llama.cpp.
Checked: HF export reproduces our training model's logits exactly (max difference 0.0).

  python export_arm_a.py finals/A/step_07820.pt export/armA_hf
  python bench/llama.cpp/convert_hf_to_gguf.py export/armA_hf --outfile export/armA-f16.gguf --outtype f16
  llama-quantize export/armA-f16.gguf export/armA-Q8_0.gguf Q8_0   (and Q4_K_M)
"""
import sys, shutil, numpy as np, torch, train
from transformers import AutoConfig, AutoTokenizer, Gemma3ForCausalLM
from huggingface_hub import hf_hub_download

ckpt, out = sys.argv[1], sys.argv[2]
counts = np.load("data/unigram_counts.npy")
m = train.Model("A", counts).eval()
m.load_state_dict(torch.load(ckpt, map_location="cpu", weights_only=False)["model"])
hf = Gemma3ForCausalLM(AutoConfig.from_pretrained("google/gemma-3-270m")).eval()
hf.model.load_state_dict({k[len("body."):]: v for k, v in m.state_dict().items() if k.startswith("body.")}, strict=True)
hf.tie_weights()
ids = torch.tensor([[2, 651, 6037, 576, 1385, 603]])
with torch.no_grad():
    d = float((m.logits(m.hidden(ids)) - hf(input_ids=ids).logits).abs().max())
assert d < 1e-3, d
hf.save_pretrained(out, safe_serialization=True)
AutoTokenizer.from_pretrained("google/gemma-3-270m").save_pretrained(out)
shutil.copy(hf_hub_download("google/gemma-3-270m", "tokenizer.model"), f"{out}/tokenizer.model")   # the GGUF converter needs it
print("exported", out, "| max logit diff", d)