Andrew0425's picture
onnx-int4 (#2)
a786ce7
|
Raw
History Blame Contribute Delete
2.13 kB

AgenticASR-Refiner ONNX (INT4)

INT4 weight-only quantization of the Optimum ONNX export of Andrew0425/AgenticASR-Refiner (Llama-based ASR transcript refiner, 24 layers, hidden 1536, GQA 16/2, head_dim 128, vocab 130560).

  • model.onnx + model.onnx.data — INT4 (MatMulNBits, block size 32, symmetric, accuracy level 4), ~1.3 GB. Graph I/O is identical to the fp32 export (input_ids / attention_mask / position_ids / past_key_values.*).
  • Quantized with onnxruntime 1.24.4 MatMulNBitsQuantizer (onnxruntime.quantization.matmul_nbits_quantizer).
  • Requires ONNX Runtime >= 1.20 (CPU EP supports MatMulNBits).

Tokenization uses the original repo tokenizer (tokenizer.json at the repo root).

Verified generation (ONNX Runtime 1.28, CPU)

Input Output
我今天去了公司然后然后开了个会,明天再去见张总 我今天去了公司然后开了个会,明天再去见张总
你好你好你好我是那个小李啊 电话是13800138000 你好我是小李,电话是13800138000

Usage

import numpy as np
import onnxruntime as ort
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("Andrew0425/AgenticASR-Refiner")
session = ort.InferenceSession("model.onnx", providers=["CPUExecutionProvider"])
input_names = [i.name for i in session.get_inputs()]
kv_names = [n for n in input_names if n.startswith("past_key_values")]

prompt = tokenizer.apply_chat_template(
    [{"role": "system", "content": "你是 ASR 文本纠错助手。保留原意,最小修改。"},
     {"role": "user", "content": "我今天去了公司然后然后开了个会"}],
    tokenize=False, add_generation_prompt=True,
)
ids = tokenizer(prompt).input_ids
pasts = [np.zeros((1, 2, 0, 128), dtype=np.float32) for _ in kv_names]

# prefill, then loop decode: feed input_ids/attention_mask/position_ids + pasts

Note: optimum-onnx 0.1.0 cannot yet run this checkpoint (dummy KV-cache shape uses hidden_size // num_heads = 96 instead of head_dim = 128); use the raw ONNX Runtime loop above, or a future fixed version of optimum-onnx.