Upload upload_to_hf.py with huggingface_hub

upload_to_hf.py

@@ -0,0 +1,198 @@
+"""Push ONNX artifacts + tokenizer assets to a HF Hub model repo.
+
+Usage:
+    uv run python upload_to_hf.py --repo onnx-community/needle-onnx
+
+Uploads:
+    encoder.onnx           — Needle encoder
+    decoder_step.onnx      — Needle one-step decoder with KV cache I/O
+    needle.model           — SentencePiece BPE model file (vocab=8192, byte_fallback)
+    tokenizer-specials.json — pad/eos/bos/<tool_call>/<tools> token IDs
+    README.md              — model card with provenance and parity numbers
+"""
+import argparse
+from pathlib import Path
+from huggingface_hub import HfApi, create_repo
+
+ART = Path(__file__).resolve().parent / "artifacts"
+WEB_DEV = Path(__file__).resolve().parent.parent / "web" / "public" / "models-dev"
+
+
+README = """---
+license: mit
+tags:
+- onnx
+- function-calling
+- needle
+- cactus
+- browser
+- sentencepiece
+base_model: Cactus-Compute/needle
+library_name: onnxruntime
+---
+
+# Needle — ONNX export for in-browser inference
+
+Browser-ready ONNX export of [Cactus-Compute/needle](https://huggingface.co/Cactus-Compute/needle), a 26M-parameter function-calling model. Designed to run entirely client-side via `onnxruntime-web` (WASM backend) — no server required.
+
+## Files
+
+| File | Description | Size |
+|---|---|---|
+| `encoder.onnx` | Needle encoder. Input `input_ids:(B,T)`, output `encoder_out:(B,T,512)`. Single-pass. | ~55 MB |
+| `decoder_step.onnx` | One decoder step with explicit past-KV in / present-KV out. Run in a JS loop. | ~85 MB |
+| `needle.model` | SentencePiece BPE protobuf (vocab=8192, `byte_fallback=True`, `identity` normalization). Loadable by `sentencepiece-js` / `@huggingface/transformers`. | 125 KB |
+| `tokenizer-specials.json` | `{"pad":0,"eos":1,"bos":2,"tool_call":4,"tools":5}` | tiny |
+
+## Origin
+
+The upstream Cactus Needle is implemented in **JAX/Flax**, not PyTorch — `torch.onnx.export` cannot run against the upstream model directly. This ONNX export was produced via a "port-and-copy" pipeline:
+
+1. Reimplemented the Simple Attention Network in PyTorch (parametric on `TransformerConfig`)
+2. Copied weights tensor-by-tensor from the upstream Flax checkpoint (handling Flax `(in, out)` → PyTorch `(out, in)` transposition for Linear kernels and the `nn.scan` layer-stacking convention)
+3. Verified Flax↔PyTorch parity at `<1e-3` max-abs-diff
+4. Exported encoder + decoder-step to ONNX via legacy TorchScript-based `torch.onnx.export`
+5. Verified PyTorch↔ONNX parity at `<1e-3`
+6. Verified end-to-end: Cactus's native `generate()` and a hand-rolled `onnxruntime` KV-cache loop produce **byte-identical** output token sequences
+
+## Parity numbers (against Cactus's native `generate(constrained=False)`)
+
+| Stage | max-abs-diff |
+|---|---|
+| Flax encoder ↔ PyTorch port | 0.000010 |
+| Flax decoder step-0 ↔ PyTorch port | 0.000029 |
+| PyTorch encoder ↔ ONNX | 0.000004 |
+| PyTorch decoder step ↔ ONNX | 0.000014 (logits) |
+| End-to-end token sequence | byte-identical |
+
+Example: `query="set a 5 min timer"` produces `' [{"name":"set_timer","arguments":{"time_human":"5 minutes"}}]'` in both Cactus native and the browser via these artifacts.
+
+## Usage in the browser
+
+Load both `.onnx` files via `onnxruntime-web` (WASM backend), load `needle.model` via `sentencepiece-js`, and run the encoder once + decoder-step in a JS loop with the KV cache passed through.
+
+## Architecture
+
+Per the upstream model card: encoder-decoder "Simple Attention Network", d_model=512, GQA 8/4 heads, 12 encoder layers, 8 decoder layers, no FFN, ZCRMSNorm (`(1+γ)·x/RMS(x)`, γ init zero), RoPE on Q and K.
+
+The decoder is exported as a **single step** with past/present KV as graph I/O — the JS side calls it in a loop, allowing streaming token output and avoiding ONNX symbolic control flow.
+
+## Reproduce / port your own Cactus-trained model
+
+The full pipeline that produced these artifacts is checked in alongside the `.onnx` files (see `PORTING.md` for the step-by-step). The scripts are parametric on the source HF repo, so if you've finetuned Needle (or trained a Simple-Attention-Network variant with the upstream Cactus codebase), you can produce a browser-ready ONNX export with the same recipe:
+
+```bash
+# 1. Convert your Cactus checkpoint → PyTorch state_dict
+uv run python convert_weights.py --ckpt-repo YOUR_USER/your-finetune --ckpt-file weights.pkl
+
+# 2. Verify the port matches your upstream model bit-for-bit (< 1e-3)
+uv run python verify_port_parity.py
+
+# 3. Export to ONNX (reads config back from step 1's saved JSON; no edits needed)
+uv run python export_onnx.py
+
+# 4. Verify ONNX matches PyTorch AND matches native Cactus generate() token-for-token
+uv run python verify_parity.py --ckpt-repo YOUR_USER/your-finetune --ckpt-file weights.pkl
+
+# 5. Push your ONNX artifacts to HF
+uv run python upload_to_hf.py --repo YOUR_USER/your-finetune-onnx
+```
+
+The PyTorch port (`needle_torch/`) is **parametric on `TransformerConfig`** — it reads the config straight out of your checkpoint's payload, so dim changes (d_model, layer counts, GQA ratios) are picked up automatically. The same pipeline works for the 26M production Needle, the 1.35M iteration config, and anything in between.
+
+Files included for reproduction:
+
+```
+needle_torch/         — PyTorch port of the Simple Attention Network
+convert_weights.py    — Flax checkpoint → PyTorch state_dict (parametric on --ckpt-repo)
+export_onnx.py        — torch.onnx.export of encoder + decoder-step
+verify_port_parity.py — Flax ↔ PyTorch parity check (load-bearing)
+verify_parity.py      — PyTorch ↔ ONNX + end-to-end vs native generate()
+dump_tokenizer.py     — Copy SentencePiece .model + emit parity goldens for the JS port
+upload_to_hf.py       — This script (push artifacts to HF Hub)
+inspect_needle.py     — Dump Flax arch / tokenizer / prompt notes (useful when porting a variant)
+pyproject.toml        — uv-managed env spec
+PORTING.md            — Full step-by-step guide
+```
+
+## License
+
+MIT, matching the upstream Cactus Needle license.
+"""
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--repo", required=True, help="e.g. onnx-community/needle-onnx")
+    p.add_argument("--private", action="store_true", help="Create as a private repo")
+    p.add_argument("--skip-lfs", action="store_true",
+                   help="Skip the .onnx + .model files (useful for re-pushing docs/scripts only)")
+    p.add_argument("--pr", action="store_true",
+                   help="Open a Pull Request instead of pushing to main (required for org-owned repos where your token lacks direct-write)")
+    args = p.parse_args()
+
+    api = HfApi()
+    create_repo(args.repo, exist_ok=True, repo_type="model", private=args.private)
+
+    HERE = Path(__file__).resolve().parent
+    lfs_files = [
+        (ART / "encoder.onnx",                 "encoder.onnx"),
+        (ART / "decoder_step.onnx",            "decoder_step.onnx"),
+        (WEB_DEV / "needle.model",             "needle.model"),
+    ]
+    text_files = [
+        (WEB_DEV / "tokenizer-specials.json",  "tokenizer-specials.json"),
+        # Reproduction pipeline (so finetuners can use the same recipe)
+        (HERE / "PORTING.md",                  "PORTING.md"),
+        (HERE / "convert_weights.py",          "convert_weights.py"),
+        (HERE / "export_onnx.py",              "export_onnx.py"),
+        (HERE / "verify_port_parity.py",       "verify_port_parity.py"),
+        (HERE / "verify_parity.py",            "verify_parity.py"),
+        (HERE / "dump_tokenizer.py",           "dump_tokenizer.py"),
+        (HERE / "inspect_needle.py",           "inspect_needle.py"),
+        (HERE / "upload_to_hf.py",             "upload_to_hf.py"),
+        (HERE / "pyproject.toml",              "pyproject.toml"),
+    ]
+    files = (lfs_files + text_files) if not args.skip_lfs else text_files
+    for local, remote in files:
+        if not local.exists():
+            raise SystemExit(f"missing artifact: {local}")
+        size = local.stat().st_size
+        print(f"uploading {remote} ({size / 1e6:.2f} MB)...", flush=True)
+        api.upload_file(
+            path_or_fileobj=str(local),
+            path_in_repo=remote,
+            repo_id=args.repo,
+            repo_type="model",
+            create_pr=args.pr,
+        )
+
+    # The PyTorch port package — preserves the dir layout so `needle_torch/__init__.py` etc.
+    pkg = HERE / "needle_torch"
+    if pkg.exists():
+        for f in sorted(pkg.iterdir()):
+            if not f.is_file() or f.name.startswith("__pycache__"):
+                continue
+            remote = f"needle_torch/{f.name}"
+            print(f"uploading {remote} ({f.stat().st_size / 1e6:.2f} MB)...", flush=True)
+            api.upload_file(
+                path_or_fileobj=str(f),
+                path_in_repo=remote,
+                repo_id=args.repo,
+                repo_type="model",
+            )
+
+    print("uploading README.md...", flush=True)
+    api.upload_file(
+        path_or_fileobj=README.encode(),
+        path_in_repo="README.md",
+        repo_id=args.repo,
+        create_pr=args.pr,
+        repo_type="model",
+    )
+
+    print(f"\ndone. https://huggingface.co/{args.repo}")
+
+
+if __name__ == "__main__":
+    main()