mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
feat: switch Kompress default to kompress-v2-base with weight-only int8 ONNX (#799)
## Summary Replaces `chopratejas/kompress-base` with **`chopratejas/kompress-v2-base`** as the default Kompress text-compression model (the fallback for content not handled by structured compressors), using a new **weight-only int8 ONNX** artifact that is fp32-equivalent at 2.2x less memory. ## Why v2 is the same dual-head ModernBERT (token classifier + span CNN), LoRA-trained on the v2 dataset. The HF repo originally shipped PyTorch weights only — pointing Headroom at it naively would have forced the heavy `[ml]` (torch) path on every `[proxy]` install. We exported ONNX artifacts reproducing the v1 loader contract (single `final_scores` output) and published them to the HF repo. ## Eval (labeled dataset_v2 test split, n=500, threshold 0.5) | artifact | size | f1 | must_keep_recall | keep_rate | fp32 agreement | |---|---|---|---|---|---| | fp32 (reference) | 601 MB | 0.9128 | 0.9770 | 0.8100 | 100% | | **int8-wo (default)** | **261 MB** | **0.9130** | **0.9765** | **0.8097** | **99.6%** | | fp16 | 301 MB | 0.9129 | 0.9771 | 0.8099 | 99.9% | | int4-wo | 230 MB | 0.9102 | 0.9825 | 0.8354 ⚠️ | 94.5% | | int8-dynamic | 271 MB | 0.9041 | 0.9868 | 0.8857 ⚠️ | 90.5% | Weight-only int8 (MatMulNBits) keeps activations fp32, avoiding the upward score bias that makes dynamic int8 keep ~7% more tokens (≈40% less compression savings). Quantized candidates were generated and eval-gated by a Modal job in the kompress repo (`modal_jobs/export_onnx_v2.py`) against the labeled test split. ## Changes - Default model id → `chopratejas/kompress-v2-base` - ONNX artifact resolution tries candidates in order (**int8-wo → fp32 → v1 int8**), falling through on download miss **or session-load failure** — onnxruntime builds without the MatMulNBits 8-bit kernel fall back to fp32 instead of losing Kompress - `HEADROOM_KOMPRESS_ONNX_FILENAME` env pins an exact artifact - `scripts/export_kompress_v2_onnx.py`: reproducible fp32 export (loads the merged v2 checkpoint, traces the `final_scores` contract, verifies vs PyTorch) - `.gitignore`: local `onnx/` artifacts dir; allowlist the export script ## Testing - End-to-end: fresh `KompressCompressor().preload()` resolves int8-wo from HF, loads on onnxruntime 1.23 (CPU, no torch), compresses with error/traceback content preserved - fp32 export verified lossless vs PyTorch (max |Δscore| = 0.0, 100% keep agreement) - ruff check + format clean, mypy clean, 63 targeted tests pass
This commit is contained in:
Tejas Chopra
GitHub
parent
3c77e52ce4
commit
74392b238e
6 changed files with 348 additions and 18 deletions
5
.gitignore
vendored
5
.gitignore
vendored
|
|
@ -3,6 +3,10 @@
|
|||
.fastembed_cache/
|
||||
**/.fastembed_cache/
|
||||
|
||||
# Local Kompress ONNX export artifacts (scripts/export_kompress_v2_onnx.py).
|
||||
# Hundreds of MB each — published to HuggingFace, never committed.
|
||||
/onnx/
|
||||
|
||||
# Private scripts (contain credentials). Allowlist checked-in helpers below.
|
||||
scripts/
|
||||
!scripts/
|
||||
|
|
@ -25,6 +29,7 @@ scripts/*
|
|||
!scripts/refresh_model_limits.sh
|
||||
!scripts/audit_wheel_glibc_symbols.py
|
||||
!scripts/replay_codex_ws_load.py
|
||||
!scripts/export_kompress_v2_onnx.py
|
||||
|
||||
# Rust / Cargo build artifacts
|
||||
/target/
|
||||
|
|
|
|||
|
|
@ -15,7 +15,7 @@
|
|||
<a href="https://app.codecov.io/gh/chopratejas/headroom"><img src="https://codecov.io/gh/chopratejas/headroom/graph/badge.svg" alt="codecov"></a>
|
||||
<a href="https://pypi.org/project/headroom-ai/"><img src="https://img.shields.io/pypi/v/headroom-ai.svg" alt="PyPI"></a>
|
||||
<a href="https://www.npmjs.com/package/headroom-ai"><img src="https://img.shields.io/npm/v/headroom-ai.svg" alt="npm"></a>
|
||||
<a href="https://huggingface.co/chopratejas/kompress-base"><img src="https://img.shields.io/badge/model-Kompress--base-yellow.svg" alt="Model: Kompress-base"></a>
|
||||
<a href="https://huggingface.co/chopratejas/kompress-v2-base"><img src="https://img.shields.io/badge/model-Kompress--v2--base-yellow.svg" alt="Model: Kompress-v2-base"></a>
|
||||
<a href="LICENSE"><img src="https://img.shields.io/badge/license-Apache%202.0-blue.svg" alt="License: Apache 2.0"></a>
|
||||
<a href="https://headroom-docs.vercel.app/docs"><img src="https://img.shields.io/badge/docs-online-blue.svg" alt="Docs"></a>
|
||||
</p>
|
||||
|
|
@ -81,7 +81,7 @@ Headroom compresses everything your AI agent reads — tool outputs, logs, RAG c
|
|||
- **CacheAligner** — stabilizes prefixes so provider KV caches actually hit
|
||||
- **CCR** — stores originals locally; LLM calls `headroom_retrieve` if it needs them
|
||||
|
||||
→ [Architecture](https://headroom-docs.vercel.app/docs/architecture) · [CCR reversible compression](https://headroom-docs.vercel.app/docs/ccr) · [Kompress-base model card](https://huggingface.co/chopratejas/kompress-base)
|
||||
→ [Architecture](https://headroom-docs.vercel.app/docs/architecture) · [CCR reversible compression](https://headroom-docs.vercel.app/docs/ccr) · [Kompress-v2-base model card](https://huggingface.co/chopratejas/kompress-v2-base)
|
||||
|
||||
## Get started (60 seconds)
|
||||
|
||||
|
|
@ -283,7 +283,7 @@ Devcontainers in `.devcontainer/` (default + `memory-stack` with Qdrant & Neo4j)
|
|||
## Community
|
||||
|
||||
- **[Discord](https://discord.gg/yRmaUNpsPJ)** — questions, feedback, war stories.
|
||||
- **[Kompress-base on HuggingFace](https://huggingface.co/chopratejas/kompress-base)** — the model behind our text compression.
|
||||
- **[Kompress-v2-base on HuggingFace](https://huggingface.co/chopratejas/kompress-v2-base)** — the model behind our text compression.
|
||||
|
||||
## License
|
||||
|
||||
|
|
|
|||
|
|
@ -128,7 +128,7 @@ class CompressConfig:
|
|||
|
||||
# Model variant
|
||||
kompress_model: str | None = None
|
||||
"""Kompress model ID. None = default (chopratejas/kompress-base).
|
||||
"""Kompress model ID. None = default (chopratejas/kompress-v2-base).
|
||||
Set to a HuggingFace model ID for domain-specific compression.
|
||||
Set to 'disabled' to skip ML compression entirely
|
||||
(only SmartCrusher + CacheAligner will run)."""
|
||||
|
|
|
|||
|
|
@ -1481,7 +1481,7 @@ class ContentRouter(Transform):
|
|||
compressed: str | None = None
|
||||
compressed_tokens: int | None = None
|
||||
|
||||
# Primary: Kompress — downloads from chopratejas/kompress-base on first use
|
||||
# Primary: Kompress — downloads from chopratejas/kompress-v2-base on first use
|
||||
if self.config.enable_kompress:
|
||||
compressor = self._get_kompress()
|
||||
if compressor:
|
||||
|
|
@ -1695,7 +1695,7 @@ class ContentRouter(Transform):
|
|||
"""Get KompressCompressor (lazy load). Downloads from HuggingFace on first use.
|
||||
|
||||
Respects runtime kompress_model kwarg:
|
||||
- None: use default (chopratejas/kompress-base) — cached on self
|
||||
- None: use default (chopratejas/kompress-v2-base) — cached on self
|
||||
- "disabled": return None (skip ML compression entirely)
|
||||
- any model ID string: create compressor with that model
|
||||
(model weights are cached at module level in kompress_compressor.py,
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
"""Kompress: ModernBERT token compressor for structured tool outputs.
|
||||
|
||||
Auto-downloads the model from HuggingFace (chopratejas/kompress-base)
|
||||
Auto-downloads the model from HuggingFace (chopratejas/kompress-v2-base)
|
||||
on first use.
|
||||
|
||||
Requires the [ml] extra: pip install headroom-ai[ml]
|
||||
|
|
@ -36,8 +36,28 @@ from .base import Transform
|
|||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Default HuggingFace model ID
|
||||
HF_MODEL_ID = "chopratejas/kompress-base"
|
||||
HF_MODEL_ID = "chopratejas/kompress-v2-base"
|
||||
KOMPRESS_BACKEND_ENV = "HEADROOM_KOMPRESS_BACKEND"
|
||||
KOMPRESS_ONNX_FILENAME_ENV = "HEADROOM_KOMPRESS_ONNX_FILENAME"
|
||||
|
||||
# ONNX artifacts are resolved against the model repo in this order, falling
|
||||
# through on download miss OR session-load failure:
|
||||
#
|
||||
# - kompress-int8-wo.onnx: weight-only int8 (MatMulNBits), 261MB. Evaluated on
|
||||
# the labeled dataset_v2 test split (n=500): f1=0.9130 vs fp32's 0.9128,
|
||||
# must_keep_recall 0.9765 vs 0.9770, keep_rate 0.8097 vs 0.8100, 99.6%
|
||||
# keep-decision agreement — fp32-equivalent at 2.2x less memory. Uses the
|
||||
# com.microsoft MatMulNBits contrib op; older onnxruntime builds without the
|
||||
# 8-bit kernel fail at session load and fall through to fp32.
|
||||
# - kompress-fp32.onnx: lossless reference, 601MB.
|
||||
# - kompress-int8.onnx: v1-era dynamic int8 (kept for custom domain repos).
|
||||
#
|
||||
# An operator can pin an exact file via HEADROOM_KOMPRESS_ONNX_FILENAME.
|
||||
_DEFAULT_ONNX_FILENAMES = (
|
||||
"onnx/kompress-int8-wo.onnx",
|
||||
"onnx/kompress-fp32.onnx",
|
||||
"onnx/kompress-int8.onnx",
|
||||
)
|
||||
KOMPRESS_ONNX_INTRA_THREADS_ENV = "HEADROOM_KOMPRESS_ONNX_INTRA_THREADS"
|
||||
KOMPRESS_ONNX_INTER_THREADS_ENV = "HEADROOM_KOMPRESS_ONNX_INTER_THREADS"
|
||||
KOMPRESS_COREML_CACHE_DIR_ENV = "HEADROOM_KOMPRESS_COREML_CACHE_DIR"
|
||||
|
|
@ -320,12 +340,56 @@ class _OnnxModel:
|
|||
return (np.array(scores) > 0.5).tolist()
|
||||
|
||||
|
||||
def _onnx_filename_candidates() -> tuple[str, ...]:
|
||||
"""ONNX repo paths to try, honoring an optional exact-file override."""
|
||||
override = os.environ.get(KOMPRESS_ONNX_FILENAME_ENV, "").strip()
|
||||
if override:
|
||||
# Put the override first but keep the defaults as a safety net.
|
||||
return (override, *(f for f in _DEFAULT_ONNX_FILENAMES if f != override))
|
||||
return _DEFAULT_ONNX_FILENAMES
|
||||
|
||||
|
||||
def _create_onnx_session(model_id: str, ort: Any, providers: list[Any]) -> Any:
|
||||
"""Resolve and load the model's ONNX artifact, trying candidates in order.
|
||||
|
||||
A candidate is skipped on download miss (file not in the repo) or on
|
||||
session-load failure (e.g. the weight-only int8 artifact uses the
|
||||
MatMulNBits contrib op, which old onnxruntime builds can't run — those
|
||||
installs fall through to the fp32 artifact instead of losing Kompress).
|
||||
"""
|
||||
last_err: Exception | None = None
|
||||
for filename in _onnx_filename_candidates():
|
||||
try:
|
||||
onnx_path = hf_hub_download_local_first(model_id, filename)
|
||||
except Exception as exc:
|
||||
last_err = exc
|
||||
logger.debug("ONNX artifact %r not in %s: %s", filename, model_id, exc)
|
||||
continue
|
||||
try:
|
||||
return ort.InferenceSession(
|
||||
onnx_path,
|
||||
_onnx_session_options(ort),
|
||||
providers=providers,
|
||||
)
|
||||
except Exception as exc:
|
||||
last_err = exc
|
||||
logger.warning(
|
||||
"ONNX artifact %r from %s failed to load (%s); trying next candidate",
|
||||
filename,
|
||||
model_id,
|
||||
exc,
|
||||
)
|
||||
raise FileNotFoundError(
|
||||
f"No loadable ONNX artifact in {model_id}; tried {_onnx_filename_candidates()}"
|
||||
) from last_err
|
||||
|
||||
|
||||
def _load_kompress_onnx(
|
||||
model_id: str,
|
||||
*,
|
||||
use_coreml: bool = False,
|
||||
) -> tuple[Any, Any, str]:
|
||||
"""Download ONNX INT8 model from HuggingFace and load with onnxruntime."""
|
||||
"""Download the ONNX model from HuggingFace and load with onnxruntime."""
|
||||
import onnxruntime as ort
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
|
|
@ -335,8 +399,6 @@ def _load_kompress_onnx(
|
|||
|
||||
logger.info("Downloading Kompress ONNX model from %s ...", model_id)
|
||||
|
||||
onnx_path = hf_hub_download_local_first(model_id, "onnx/kompress-int8.onnx")
|
||||
|
||||
backend = "onnx_coreml" if use_coreml else "onnx"
|
||||
providers: list[Any]
|
||||
if use_coreml:
|
||||
|
|
@ -364,16 +426,12 @@ def _load_kompress_onnx(
|
|||
else:
|
||||
providers = ["CPUExecutionProvider"]
|
||||
|
||||
session = ort.InferenceSession(
|
||||
onnx_path,
|
||||
_onnx_session_options(ort),
|
||||
providers=providers,
|
||||
)
|
||||
session = _create_onnx_session(model_id, ort, providers)
|
||||
model = _OnnxModel(session)
|
||||
tokenizer = AutoTokenizer.from_pretrained("answerdotai/ModernBERT-base")
|
||||
|
||||
_kompress_cache[model_id] = (model, tokenizer, backend)
|
||||
logger.info("Kompress ONNX INT8 loaded: %s backend=%s", model_id, backend)
|
||||
logger.info("Kompress ONNX loaded: %s backend=%s", model_id, backend)
|
||||
return model, tokenizer, backend
|
||||
|
||||
|
||||
|
|
@ -533,7 +591,7 @@ class KompressConfig:
|
|||
|
||||
The model_id, chunk_words, and score_threshold are coupled: a model
|
||||
trained on 50-word chunks needs chunk_words=50 at inference. The
|
||||
defaults match kompress-base. For domain-specific models, set all three.
|
||||
defaults match kompress-v2-base. For domain-specific models, set all three.
|
||||
|
||||
Example — financial documents::
|
||||
|
||||
|
|
|
|||
267
scripts/export_kompress_v2_onnx.py
Normal file
267
scripts/export_kompress_v2_onnx.py
Normal file
|
|
@ -0,0 +1,267 @@
|
|||
#!/usr/bin/env python
|
||||
"""Export a Kompress PyTorch checkpoint to ONNX INT8 for Headroom's light path.
|
||||
|
||||
Why this exists
|
||||
---------------
|
||||
Headroom's ``[proxy]`` extra ships ``onnxruntime`` but **not** torch — the
|
||||
proxy runs Kompress text compression on ONNX Runtime alone. The loader
|
||||
(``headroom/transforms/kompress_compressor.py``) downloads
|
||||
``onnx/kompress-int8.onnx`` from the model repo and runs it through
|
||||
``_OnnxModel``, which expects a single graph output named ``final_scores``
|
||||
(per-token importance in ``[0, 1]``, kept when ``> 0.5``).
|
||||
|
||||
``chopratejas/kompress-v2-base`` ships only PyTorch weights
|
||||
(``model.safetensors`` / ``merged.pt``) — no ONNX. So pointing Headroom at v2
|
||||
without an ONNX export would silently force the heavier ``[ml]`` (torch) path
|
||||
on every proxy install. This script reproduces v1's exact ONNX contract from
|
||||
the v2 PyTorch checkpoint, so a default swap stays zero-cost for light installs.
|
||||
|
||||
The model is a *custom* dual-head ModernBERT (token classifier + span CNN), not
|
||||
a standard HF architecture, so ``optimum-cli export onnx`` does not apply — we
|
||||
trace the real module from ``kompress_compressor._get_model_class()``.
|
||||
|
||||
Requires
|
||||
--------
|
||||
pip install headroom-ai[ml] onnxruntime # torch + transformers + onnxruntime
|
||||
|
||||
Usage
|
||||
-----
|
||||
# Convert + verify locally (writes onnx/kompress-int8.onnx):
|
||||
python scripts/export_kompress_v2_onnx.py --model-id chopratejas/kompress-v2-base
|
||||
|
||||
# Convert, verify, and upload back to the HF repo (needs `huggingface-cli login`):
|
||||
python scripts/export_kompress_v2_onnx.py --model-id chopratejas/kompress-v2-base --upload
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import logging
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")
|
||||
logger = logging.getLogger("export_kompress_v2_onnx")
|
||||
|
||||
# ModernBERT encoder + tokenizer base (must match training and the loader).
|
||||
BASE_MODEL = "answerdotai/ModernBERT-base"
|
||||
DEFAULT_MODEL_ID = "chopratejas/kompress-v2-base"
|
||||
|
||||
|
||||
def _build_core(model_id: str):
|
||||
"""Instantiate HeadroomCompressorModel and load the merged v2 weights.
|
||||
|
||||
The v2 repo's ``model.safetensors`` is the *unmerged* PEFT structure
|
||||
(``encoder.base_model.model...`` with separate ``base_layer`` + LoRA
|
||||
adapters), which does not map onto ``HeadroomCompressorModel``. The
|
||||
canonical artifact is ``merged.pt`` — a structured checkpoint with already
|
||||
LoRA-merged sub-state-dicts:
|
||||
|
||||
{"encoder_state_dict", "token_head_state_dict",
|
||||
"span_conv_state_dict", "config", "checkpoint_kind"}
|
||||
|
||||
Each loads cleanly (0 missing / 0 unexpected) into the encoder + heads.
|
||||
"""
|
||||
import torch
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
from headroom.transforms.kompress_compressor import _get_model_class
|
||||
|
||||
ckpt_path = hf_hub_download(model_id, "merged.pt")
|
||||
ckpt = torch.load(ckpt_path, map_location="cpu")
|
||||
for key in ("encoder_state_dict", "token_head_state_dict", "span_conv_state_dict"):
|
||||
if key not in ckpt:
|
||||
raise RuntimeError(
|
||||
f"merged.pt missing '{key}'. Found: {sorted(ckpt)}. "
|
||||
"This script targets the v2 'merged' checkpoint format."
|
||||
)
|
||||
|
||||
core = _get_model_class()(model_name=BASE_MODEL)
|
||||
|
||||
def _strict_load(module, sd, label: str) -> None:
|
||||
missing, unexpected = module.load_state_dict(sd, strict=False)
|
||||
if missing or unexpected:
|
||||
raise RuntimeError(
|
||||
f"{label}: state_dict mismatch (missing={list(missing)[:5]}, "
|
||||
f"unexpected={list(unexpected)[:5]}). Architecture drifted from the checkpoint."
|
||||
)
|
||||
logger.info(" %s loaded (%d tensors, exact match)", label, len(sd))
|
||||
|
||||
logger.info("Loading merged.pt (checkpoint_kind=%s)", ckpt.get("checkpoint_kind"))
|
||||
_strict_load(core.encoder, ckpt["encoder_state_dict"], "encoder")
|
||||
_strict_load(core.token_head, ckpt["token_head_state_dict"], "token_head")
|
||||
_strict_load(core.span_conv, ckpt["span_conv_state_dict"], "span_conv")
|
||||
|
||||
core.eval()
|
||||
return core
|
||||
|
||||
|
||||
def _export_wrapper(core):
|
||||
"""Wrap the dual head so forward() returns `final_scores` (== get_scores)."""
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
class ExportWrapper(nn.Module):
|
||||
def __init__(self, inner):
|
||||
super().__init__()
|
||||
self.inner = inner
|
||||
|
||||
def forward(self, input_ids, attention_mask): # noqa: ANN001
|
||||
hidden = self.inner.encoder(input_ids, attention_mask=attention_mask).last_hidden_state
|
||||
token_probs = torch.softmax(self.inner.token_head(hidden), dim=-1)[:, :, 1]
|
||||
span_scores = self.inner.span_conv(hidden.transpose(1, 2)).squeeze(1)
|
||||
return token_probs * (0.5 + 0.5 * span_scores)
|
||||
|
||||
return ExportWrapper(core).eval()
|
||||
|
||||
|
||||
def export(model_id: str, out_path: Path, opset: int, precision: str) -> None:
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
core = _build_core(model_id)
|
||||
wrapper = _export_wrapper(core)
|
||||
|
||||
out_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# fp32 path: trace straight to the final artifact (lossless — verified 100%
|
||||
# keep-decision agreement with PyTorch). int8 path: trace to a temp fp32
|
||||
# graph, then dynamically quantize into the final artifact.
|
||||
trace_target = out_path if precision == "fp32" else out_path.with_name("kompress-fp32-tmp.onnx")
|
||||
|
||||
dummy_ids = torch.randint(0, 1000, (1, 64), dtype=torch.long)
|
||||
dummy_mask = torch.ones((1, 64), dtype=torch.long)
|
||||
|
||||
logger.info("Tracing → ONNX (opset %d, precision=%s) ...", opset, precision)
|
||||
with torch.no_grad():
|
||||
torch.onnx.export(
|
||||
wrapper,
|
||||
(dummy_ids, dummy_mask),
|
||||
str(trace_target),
|
||||
input_names=["input_ids", "attention_mask"],
|
||||
output_names=["final_scores"],
|
||||
dynamic_axes={
|
||||
"input_ids": {0: "batch", 1: "seq"},
|
||||
"attention_mask": {0: "batch", 1: "seq"},
|
||||
"final_scores": {0: "batch", 1: "seq"},
|
||||
},
|
||||
opset_version=opset,
|
||||
do_constant_folding=True,
|
||||
dynamo=False,
|
||||
)
|
||||
|
||||
if precision == "int8":
|
||||
from onnxruntime.quantization import QuantType, quantize_dynamic
|
||||
|
||||
logger.info("INT8 dynamic quantization (MatMul only) → %s", out_path)
|
||||
# Restrict to MatMul: the encoder's linear layers carry ~all the weight
|
||||
# mass and ORT's CPU provider implements MatMulInteger. Quantizing the
|
||||
# tiny span_conv Conv1d layers would emit ConvInteger, which ORT CPU
|
||||
# cannot run. per_channel recovers transformer accuracy at the 0.5 boundary.
|
||||
quantize_dynamic(
|
||||
str(trace_target),
|
||||
str(out_path),
|
||||
weight_type=QuantType.QInt8,
|
||||
op_types_to_quantize=["MatMul"],
|
||||
per_channel=True,
|
||||
)
|
||||
trace_target.unlink(missing_ok=True)
|
||||
|
||||
_verify(model_id, core, out_path, np, torch)
|
||||
|
||||
|
||||
def _verify(model_id: str, core, out_path: Path, np, torch) -> None:
|
||||
"""Compare ONNX scores against PyTorch get_scores on a real tokenized sample."""
|
||||
import onnxruntime as ort
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
tok = AutoTokenizer.from_pretrained(BASE_MODEL)
|
||||
sample = (
|
||||
"The proxy compresses tool outputs before they reach the model. "
|
||||
"Errors and stack traces should survive; boilerplate should not. "
|
||||
) * 6
|
||||
words = sample.split()
|
||||
enc = tok(
|
||||
words,
|
||||
is_split_into_words=True,
|
||||
truncation=True,
|
||||
max_length=512,
|
||||
padding=True,
|
||||
return_tensors="pt",
|
||||
)
|
||||
|
||||
with torch.no_grad():
|
||||
torch_scores = core.get_scores(enc["input_ids"], enc["attention_mask"])[0].cpu().numpy()
|
||||
|
||||
sess = ort.InferenceSession(str(out_path), providers=["CPUExecutionProvider"])
|
||||
onnx_scores = sess.run(
|
||||
["final_scores"],
|
||||
{
|
||||
"input_ids": enc["input_ids"].numpy().astype(np.int64),
|
||||
"attention_mask": enc["attention_mask"].numpy().astype(np.int64),
|
||||
},
|
||||
)[0][0]
|
||||
|
||||
max_abs = float(np.max(np.abs(torch_scores - onnx_scores)))
|
||||
keep_torch = torch_scores > 0.5
|
||||
keep_onnx = onnx_scores > 0.5
|
||||
agree = float((keep_torch == keep_onnx).mean())
|
||||
logger.info(
|
||||
"Verify: max|Δscore|=%.4f keep-decision agreement=%.1f%% (fp32 ~100%%, int8 ~98-100%%)",
|
||||
max_abs,
|
||||
agree * 100,
|
||||
)
|
||||
if agree < 0.98:
|
||||
logger.warning(
|
||||
"Keep-decision agreement below 98%% — for fp32 this means a tracing "
|
||||
"problem; for int8 consider per_channel/fp32. Inspect before publishing."
|
||||
)
|
||||
|
||||
|
||||
def upload(model_id: str, out_path: Path) -> None:
|
||||
from huggingface_hub import upload_file
|
||||
|
||||
# Publish under onnx/<artifact filename> so int8 and fp32 can coexist.
|
||||
repo_path = f"onnx/{out_path.name}"
|
||||
logger.info("Uploading %s → %s:%s", out_path, model_id, repo_path)
|
||||
upload_file(
|
||||
path_or_fileobj=str(out_path),
|
||||
path_in_repo=repo_path,
|
||||
repo_id=model_id,
|
||||
commit_message="Add ONNX export for Headroom lightweight (no-torch) path",
|
||||
)
|
||||
logger.info("Uploaded. Headroom's ONNX loader will now find it on next cold start.")
|
||||
|
||||
|
||||
def main() -> int:
|
||||
ap = argparse.ArgumentParser(description=__doc__)
|
||||
ap.add_argument("--model-id", default=DEFAULT_MODEL_ID)
|
||||
ap.add_argument(
|
||||
"--precision",
|
||||
choices=["fp32", "int8"],
|
||||
default="fp32",
|
||||
help="fp32 = lossless, larger artifact. int8 = ~2x smaller, tiny accuracy cost.",
|
||||
)
|
||||
ap.add_argument(
|
||||
"--out",
|
||||
type=Path,
|
||||
default=None,
|
||||
help="Local output path. Defaults to onnx/kompress-<precision>.onnx.",
|
||||
)
|
||||
ap.add_argument("--opset", type=int, default=17)
|
||||
ap.add_argument(
|
||||
"--upload",
|
||||
action="store_true",
|
||||
help="Upload to the HF repo under onnx/<filename> (needs HF write auth).",
|
||||
)
|
||||
args = ap.parse_args()
|
||||
|
||||
out_path = args.out or Path(f"onnx/kompress-{args.precision}.onnx")
|
||||
export(args.model_id, out_path, args.opset, args.precision)
|
||||
if args.upload:
|
||||
upload(args.model_id, out_path)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
Loading…
Add table
Add a link
Reference in a new issue