#!/usr/bin/env python3 """Export GLiNER 2.5 BoundaryExtractor to ONNX for onnxruntime-web / WebGPU. The graph is encoder + word/query gather + boundary start/end logits. Schema packing and span decode stay in the host (JS/Python), matching how GLiNER.js runs the original span models. ═══════════════════════════════════════════════════════════════════════════ TWO MODES: LOCAL (direct) and MODAL (remote CPU) ═══════════════════════════════════════════════════════════════════════════ ── LOCAL (no Modal, runs on your machine) ────────────────────────────── Requires Python 3.11+ and these packages: pip install "torch==2.5.1" "transformers>=4.46.0" "huggingface_hub[hf_transfer]" \ onnx onnxruntime safetensors sentencepiece tokenizers "gliner2[local]" Set your HF token if uploading: export HF_TOKEN=hf_xxx # or `hf auth login` Run: python convert_gliner25_onnx.py --model-id fastino/gliner2.5-small-v1 python convert_gliner25_onnx.py --model-id fastino/gliner2.5-base-v1 --upload --upload-prefix nicolasembleton python convert_gliner25_onnx.py --model-id fastino/gliner2.5-multi-v1 --seq-len 256 --n-queries 8 Output appears at ./output/{slug}-onnx/ with: onnx/model.onnx — the graph tokenizer.json + config — for host-side schema packing export_config.json — input/output contract README.md — Hub model card with YAML ── MODAL (remote CPU, no local torch install) ────────────────────────── Requires the `modal` CLI authenticated (`modal token new`). Also requires a Modal Secret named `huggingface-token` with key `HF_TOKEN` (your Hub write token). Create it once: modal secret create huggingface-token HF_TOKEN=hf_xxx The Modal image installs the same packages as above inside a Debian-slim container (python 3.11). 4 CPU / 16 GB RAM is enough for all GLiNER 2.5 checkpoints (74M–287M params). A Modal Volume (`gliner25-onnx`) caches outputs between runs but is deleted after upload to stop storage cost. Run: modal run convert_gliner25_onnx.py --model-id fastino/gliner2.5-small-v1 \ --upload --upload-prefix nicolasembleton When run via `modal run`, the `@app.local_entrypoint` fires, which calls `export_one.remote(...)` — the function executes inside the Modal container. When run via `python convert_gliner25_onnx.py`, the `__main__` block calls `export_one(...)` directly in the current process. ── KEY IMPLEMENTATION NOTES (read before modifying) ──────────────────── 1. EyeLike fix: BoundaryAttentionBlock.forward uses torch.eye() + SDPA which exports to ONNX as an EyeLike op that onnxruntime-web does not implement. We monkey-patch each attention block's forward with a matmul/softmax version that produces identical outputs but uses only standard ONNX ops. This is the single most important patch in the file. 2. export_mode="vectorized": sets the boundary proposer to materialize a single full-width block instead of a Python block loop. Needed for graph export; does not change the logits head. 3. The Wrapper class only wraps encoder + boundary_encoder + boundary_query_head. It does NOT wrap the proposer/scorer/pool — those are Python-side post-processing. The ONNX graph outputs raw start_logits and end_logits. The host (JS/Python) does top-k selection and span pairing from those logits. 4. Dynamic axes on batch, tokens, words, and queries so the graph handles arbitrary input lengths at inference time. 5. Opset 17 is the minimum that supports all ops used here. WebGPU via onnxruntime-web supports opset 17+. 6. ORT validation: we catch ORT exceptions and re-raise as RuntimeError because Modal can't deserialize onnxruntime-specific exception types in the local environment. The RMSE check confirms the ONNX graph matches torch output to ~1e-6. ═══════════════════════════════════════════════════════════════════════════ """ from __future__ import annotations import argparse import json import os import shutil from pathlib import Path # ── Modal setup (imported lazily so local mode works without modal) ────── try: import modal _HAS_MODAL = True except ImportError: modal = None _HAS_MODAL = False VOLUME_NAME = "gliner25-onnx" # Inside Modal, outputs go to a mounted Volume. Locally, to ./output/. _MODAL_OUT_DIR = "/data/output" _LOCAL_OUT_DIR = str(Path(__file__).resolve().parent / "output") # ── Modal image + app (only built if modal is imported) ────────────────── if _HAS_MODAL: image = ( modal.Image.debian_slim(python_version="3.11") .pip_install( "torch==2.5.1", "transformers>=4.46.0", "huggingface_hub[hf_transfer]", "onnx", "onnxruntime", "safetensors", "sentencepiece", "tokenizers", "gliner2[local]", ) .env({"HF_HUB_ENABLE_HF_TRANSFER": "1", "TOKENIZERS_PARALLELISM": "false"}) ) app = modal.App("gliner25-onnx") vol = modal.Volume.from_name(VOLUME_NAME, create_if_missing=True) else: image = None app = None vol = None def export_one( model_id: str, upload_prefix: str = "", seq_len: int = 128, n_queries: int = 4, n_words: int = 48, out_dir: str | None = None, upload: bool = False, ): """Export a single GLiNER 2.5 model to ONNX. Args: model_id: HuggingFace model ID (e.g. fastino/gliner2.5-small-v1) upload_prefix: HF namespace to upload to (e.g. nicolasembleton). Empty = skip upload. seq_len: dummy sequence length for the ONNX trace (dynamic at runtime) n_queries: dummy number of entity-type queries (dynamic at runtime) n_words: dummy number of word positions (dynamic at runtime) out_dir: output directory. Defaults to ./output/ locally, /data/output/ on Modal. upload: if True and upload_prefix is set, upload to HF Hub (requires HF_TOKEN) Returns: dict with repo URL, ONNX file size in MB, and RMSE vs torch """ import torch import torch.nn as nn if out_dir is None: out_dir = _MODAL_OUT_DIR if _HAS_MODAL and modal.App.current() else _LOCAL_OUT_DIR from gliner2 import AutoExtractor print(f"Loading {model_id} ...") model = AutoExtractor.from_pretrained(model_id, map_location="cpu") model.eval() # Graph-friendly proposer (no Python block loop). Does not change logits head. try: model.boundary_head.boundary_proposer.settings.export_mode = "vectorized" except Exception as e: print(f" [warn] could not set export_mode: {e}") encoder = model.encoder encoder.eval() hidden = encoder.config.hidden_size # ── EyeLike fix: replace torch.eye + SDPA with matmul attention ─────── def _exportable_attn_forward(block, states, mask): """matmul/softmax attention — avoids EyeLike op unsupported by ORT-web.""" b, n, d = states.shape qkv = block.qkv_projection(block.norm(states)).view(b, n, 3, block.num_heads, block.head_dim) query, key, value = qkv.permute(2, 0, 3, 1, 4) scale = block.head_dim ** -0.5 scores = torch.matmul(query, key.transpose(-2, -1)) * scale allowed = mask.view(b, 1, 1, n) if block.window > 0: positions = torch.arange(n, device=states.device) local = (positions.unsqueeze(1) - positions.unsqueeze(0)).abs() <= block.window allowed = allowed & local.view(1, 1, n, n) idx = torch.arange(n, device=states.device) diag = idx.unsqueeze(0) == idx.unsqueeze(1) allowed = allowed | diag.view(1, 1, n, n) scores = scores.masked_fill(~allowed, -1.0e4) attn = torch.softmax(scores, dim=-1) attended = torch.matmul(attn, value).transpose(1, 2).reshape(b, n, d) update = block.dropout(block.output_projection(attended)) return (states + update) * mask.unsqueeze(-1).to(states.dtype) class Wrapper(nn.Module): """Encoder + boundary start/end logits. No proposer/scorer (host-side).""" def __init__(self, extractor): super().__init__() self.encoder = extractor.encoder self.boundary_encoder = extractor.boundary_head.boundary_encoder self.boundary_query_head = extractor.boundary_head.boundary_query_head for block in self.boundary_encoder.attention_blocks: block.forward = lambda states, mask, _b=block: _exportable_attn_forward(_b, states, mask) def _gather(self, hidden_states, indices, mask): h = hidden_states.shape[-1] safe = indices.clamp(0, hidden_states.shape[1] - 1) states = hidden_states.gather(1, safe.unsqueeze(-1).expand(-1, -1, h)) return states * mask.unsqueeze(-1).to(states.dtype) def forward(self, input_ids, attention_mask, text_word_indices, text_word_mask, query_marker_indices, query_marker_mask): hidden_states = self.encoder(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state text_states = self._gather(hidden_states, text_word_indices, text_word_mask) query_states = self._gather(hidden_states, query_marker_indices, query_marker_mask) encoding = self.boundary_encoder(text_states, text_word_mask.bool()) marginals = self.boundary_query_head( encoding.states, encoding.mask, text_states, text_word_mask.bool(), query_states, query_marker_mask.bool(), ) return marginals.start_logits, marginals.end_logits wrapper = Wrapper(model).eval() # ── Dummy inputs for tracing ───────────────────────────────────────── b, t, l, q = 1, seq_len, n_words, n_queries dummy = { "input_ids": torch.ones(b, t, dtype=torch.long), "attention_mask": torch.ones(b, t, dtype=torch.long), "text_word_indices": torch.arange(l, dtype=torch.long).clamp(max=t - 1).unsqueeze(0), "text_word_mask": torch.ones(b, l, dtype=torch.float32), "query_marker_indices": torch.arange(q, dtype=torch.long).clamp(max=t - 1).unsqueeze(0), "query_marker_mask": torch.ones(b, q, dtype=torch.float32), } with torch.no_grad(): s, e = wrapper(*dummy.values()) print(f" dummy start {tuple(s.shape)} end {tuple(e.shape)}") # ── Export ──────────────────────────────────────────────────────────── slug = model_id.split("/")[-1] out = Path(out_dir) / f"{slug}-onnx" if out.exists(): shutil.rmtree(out) onnx_dir = out / "onnx" onnx_dir.mkdir(parents=True) onnx_path = onnx_dir / "model.onnx" input_names = list(dummy.keys()) dynamic_axes = { "input_ids": {0: "batch", 1: "tokens"}, "attention_mask": {0: "batch", 1: "tokens"}, "text_word_indices": {0: "batch", 1: "words"}, "text_word_mask": {0: "batch", 1: "words"}, "query_marker_indices": {0: "batch", 1: "queries"}, "query_marker_mask": {0: "batch", 1: "queries"}, "start_logits": {0: "batch", 1: "queries", 2: "boundaries"}, "end_logits": {0: "batch", 1: "queries", 2: "boundaries"}, } print(" torch.onnx.export ...") torch.onnx.export( wrapper, tuple(dummy[k] for k in input_names), str(onnx_path), input_names=input_names, output_names=["start_logits", "end_logits"], dynamic_axes=dynamic_axes, opset_version=17, do_constant_folding=True, ) size_mb = onnx_path.stat().st_size / 1e6 print(f" wrote {onnx_path} ({size_mb:.1f} MB)") # ── Validate with onnxruntime ──────────────────────────────────────── import onnx import onnxruntime as ort onnx.checker.check_model(str(onnx_path)) try: sess = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"]) feeds = {k: v.numpy() for k, v in dummy.items()} outs = sess.run(None, feeds) print(f" ort start {outs[0].shape} end {outs[1].shape}") err = float(((outs[0] - s.numpy()) ** 2).mean() ** 0.5) print(f" start-logit RMSE vs torch: {err:.6f}") except Exception as e: # Re-raise as RuntimeError: Modal can't deserialize ORT exception types locally raise RuntimeError(f"ORT load/run failed: {type(e).__name__}: {e}") from None # ── Save tokenizer + config for host-side packing ──────────────────── tok = model.processor.tokenizer tok.save_pretrained(str(out)) cfg = { "architecture": "boundary", "base_model": model_id, "hidden_size": hidden, "opset": 17, "inputs": input_names, "outputs": ["start_logits", "end_logits"], "notes": "Host must pack schema markers into input_ids and pass word/query gather indices. Decode spans from start/end logits in JS.", } (out / "export_config.json").write_text(json.dumps(cfg, indent=2)) readme = f"""--- library_name: onnx license: apache-2.0 pipeline_tag: token-classification base_model: {model_id} tags: - onnx - gliner2 - boundary - webgpu - token-classification --- # {slug}-onnx ONNX export of [{model_id}](https://huggingface.co/{model_id}) (GLiNER 2.5 `BoundaryExtractor`) for onnxruntime / WebGPU. This is **not** a drop-in `AutoExtractor` graph. The ONNX file runs: 1. DeBERTa encoder on packed `input_ids` 2. Gather of word states and query-marker states 3. Boundary start/end logits `[batch, queries, words+1]` Schema packing (entity-type markers) and span decode stay on the host, same split as GLiNER.js. ## Inputs | Name | Shape | Dtype | |------|-------|-------| | input_ids | [B, T] | int64 | | attention_mask | [B, T] | int64 | | text_word_indices | [B, L] | int64 | | text_word_mask | [B, L] | float32 | | query_marker_indices | [B, Q] | int64 | | query_marker_mask | [B, Q] | float32 | ## Outputs | Name | Shape | |------|-------| | start_logits | [B, Q, L+1] | | end_logits | [B, Q, L+1] | ## Python check ```python import onnxruntime as ort sess = ort.InferenceSession("onnx/model.onnx") ``` WebGPU: load `onnx/model.onnx` with `onnxruntime-web` `webgpu` execution provider. Int64 inputs are required; some browsers need the WASM backend as fallback. """ (out / "README.md").write_text(readme) # ── Commit volume if on Modal ──────────────────────────────────────── if _HAS_MODAL and vol is not None: vol.commit() # ── Upload to HuggingFace Hub ──────────────────────────────────────── if upload and upload_prefix: from huggingface_hub import HfApi repo = f"{upload_prefix}/{slug}-onnx" token = os.environ.get("HF_TOKEN") if not token: raise RuntimeError("HF_TOKEN missing — set it or pass --no-upload") api = HfApi(token=token) api.create_repo(repo_id=repo, repo_type="model", exist_ok=True) print(f" uploading → {repo}") api.upload_folder(folder_path=str(out), repo_id=repo, repo_type="model") print(f" done https://huggingface.co/{repo}") return {"repo": repo, "onnx_mb": size_mb, "rmse": err} else: print(f" output at {out} (no upload)") return {"repo": None, "onnx_mb": size_mb, "rmse": err} # ═══ Modal entrypoint ═════════════════════════════════════════════════════ if _HAS_MODAL: @app.function( image=image, volumes={"/data": vol}, secrets=[modal.Secret.from_name("huggingface-token")], cpu=4, memory=16384, timeout=3600, ) def _export_remote(model_id, upload_prefix, seq_len, n_queries, n_words): return export_one( model_id=model_id, upload_prefix=upload_prefix, seq_len=seq_len, n_queries=n_queries, n_words=n_words, out_dir=_MODAL_OUT_DIR, upload=True, ) @app.local_entrypoint() def main( model_id: str = "fastino/gliner2.5-small-v1", upload_prefix: str = "nicolasembleton", seq_len: int = 128, n_queries: int = 4, n_words: int = 48, ): """Modal entrypoint: `modal run convert_gliner25_onnx.py --model-id ...`""" result = _export_remote.remote( model_id=model_id, upload_prefix=upload_prefix, seq_len=seq_len, n_queries=n_queries, n_words=n_words, ) print(result) # ═══ Local CLI entrypoint ════════════════════════════════════════════════ # Works whether or not modal is installed. When modal is installed, the # @app.local_entrypoint above handles `modal run ...`. For local execution # use: `python convert_gliner25_onnx.py --model-id ... --local` # The --local flag forces the local codepath even when modal is present. def _local_cli(): parser = argparse.ArgumentParser(description="Export GLiNER 2.5 to ONNX (local mode)") parser.add_argument("--model-id", required=True, help="HuggingFace model ID") parser.add_argument("--upload-prefix", default="", help="HF namespace to upload to") parser.add_argument("--upload", action="store_true", help="Upload to HF Hub") parser.add_argument("--seq-len", type=int, default=128) parser.add_argument("--n-queries", type=int, default=4) parser.add_argument("--n-words", type=int, default=48) parser.add_argument("--out-dir", default=None, help="Output directory (default: ./output/)") args = parser.parse_args() result = export_one( model_id=args.model_id, upload_prefix=args.upload_prefix, seq_len=args.seq_len, n_queries=args.n_queries, n_words=args.n_words, out_dir=args.out_dir, upload=args.upload, ) print(result) if __name__ == "__main__": # If --local is in argv, strip it and run locally regardless of modal. import sys if "--local" in sys.argv: sys.argv.remove("--local") _local_cli() elif _HAS_MODAL: # modal run will pick up @app.local_entrypoint; if python was used # directly without --local, fall back to local CLI too. _local_cli() else: _local_cli()