"""Export local HF Falcon-H1 safetensors into exact-name f32 FH1W tensor pack.

Dependencies: pip install safetensors torch huggingface_hub
Example:
  python tools/export.py --repo tiiuae/Falcon-H1-0.5B-Instruct --output falcon-h1.f32pack

The 0.5B model is ~2 GB in float32; this is CPU reference, not a GPU converter.
"""
import argparse
import struct
from pathlib import Path
from huggingface_hub import snapshot_download
from safetensors import safe_open


def convert(model_dir: Path, output: Path):
    paths = sorted(model_dir.glob("*.safetensors"))
    if not paths:
        raise ValueError("No safetensors files in selected model directory")
    # Include only model tensors; do not include unrelated training checkpoint tensors.
    entries = []
    for path in paths:
        with safe_open(str(path), framework="pt", device="cpu") as sf:
            for name in sf.keys():
                if name.startswith(("model.", "lm_head.")):
                    entries.append((path, name))
    if len(entries) > 50000:
        raise ValueError("Too many checkpoint tensors")
    output.parent.mkdir(parents=True, exist_ok=True)
    with output.open("wb") as f:
        f.write(b"FH1W")
        f.write(struct.pack("<I", len(entries)))
        for path, name in entries:
            with safe_open(str(path), framework="pt", device="cpu") as sf:
                tensor = sf.get_tensor(name).float().contiguous()
                name_b = name.encode("utf-8")
                shape = tuple(tensor.shape)
                if not (1 <= len(shape) <= 4):
                    raise ValueError((name, shape))
                f.write(struct.pack("<H", len(name_b)))
                f.write(name_b)
                f.write(struct.pack("<B", len(shape)))
                for dim in shape:
                    f.write(struct.pack("<I", dim))
                f.write(tensor.numpy().astype("<f4", copy=False).tobytes())
    print(f"Wrote {len(entries)} named float32 tensors: {output}")


if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--repo", default="tiiuae/Falcon-H1-0.5B-Instruct")
    p.add_argument("--local", help="Use an already downloaded model directory")
    p.add_argument("--output", default="falcon-h1.f32pack")
    a = p.parse_args()
    root = Path(a.local) if a.local else Path(snapshot_download(repo_id=a.repo, allow_patterns=["*.safetensors", "*.json"]))
    convert(root, Path(a.output))