Sharjeelbaig's picture
Add verified INT8 ONNX model with KV cache
03081fc verified
|
Raw
History Blame Contribute Delete
4.14 kB
metadata
library_name: transformers.js
pipeline_tag: text-generation
base_model: SupraLabs/Supra-Router-51M
tags:
  - onnx
  - transformers.js
  - browser
  - web
  - int8
  - quantized
  - llama
  - router
  - text-generation

Supra-Router-51M ONNX INT8

Browser-ready ONNX conversion of SupraLabs/Supra-Router-51M. The graph is dynamically quantized to signed INT8 and packaged using the Hugging Face ONNX repository layout.

  • Parameters: 51.8M
  • ONNX size: approximately 66 MB
  • Architecture: LlamaForCausalLM
  • Execution: ONNX Runtime Web (WASM) or ONNX Runtime Python
  • Conversion: FP32 export, opset 17, per-channel dynamic INT8 weight quantization
  • KV cache: included for efficient autoregressive generation

Transformers.js pipeline

npm install @huggingface/transformers
import { pipeline } from "@huggingface/transformers";

const modelId = "Sharjeelbaig/Supra-Router-51M-ONNX";
const router = await pipeline("text-generation", modelId, {
  dtype: "int8",
});

const userPrompt = "Write Python code to find all primes below one million efficiently.";
const input = `Task: ${userPrompt}\nAnalysis: `;

const output = await router(input, {
  max_new_tokens: 128,
  do_sample: false,
  return_full_text: false,
});

console.log(output[0].generated_text.trim());

The model returns a pipe-separated routing record:

Domain: ... | Complexity: 1-5 | Math: True/False | Code: True/False | Route: small model/big model | Justification: ...

Python with Optimum ONNX Runtime

pip install "optimum-onnx[onnxruntime]" transformers
from transformers import AutoTokenizer, pipeline
from optimum.onnxruntime import ORTModelForCausalLM

model_id = "Sharjeelbaig/Supra-Router-51M-ONNX"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = ORTModelForCausalLM.from_pretrained(
    model_id,
    subfolder="onnx",
    file_name="model_int8.onnx",
    use_cache=False,
)

router = pipeline("text-generation", model=model, tokenizer=tokenizer)
prompt = "Explain why database deadlocks occur and provide code to prevent them."
result = router(
    f"Task: {prompt}\nAnalysis: ",
    max_new_tokens=128,
    do_sample=False,
    return_full_text=False,
)
print(result[0]["generated_text"].strip())

Direct ONNX Runtime loading

The graph accepts:

  • input_ids: int64[batch, sequence]
  • attention_mask: int64[batch, past_sequence + sequence]
  • position_ids: int64[batch, sequence]
  • past_key_values.{0..11}.{key,value}: cached attention tensors

It returns logits: float32[batch, sequence, 32000] plus present.{0..11}.{key,value} cache tensors. For direct integration, initialize each cache with shape [batch, 4, 0, 64], then pass each returned present tensor back as the corresponding past_key_values input on the next step. Hugging Face pipelines manage this automatically.

Validation

  • Passed onnx.checker.
  • Tested with dynamic sequence lengths.
  • Cached FP32 ONNX maximum absolute logit error versus PyTorch was below 1.4e-4 during export validation.
  • INT8 and FP32 generated identical complete routing strings on representative small model and big model prompts.
  • The uploaded INT8 interface was run end-to-end through both a Transformers.js text-generation pipeline and a Python Optimum text-generation pipeline.

Limitations and responsible use

This is a small routing model trained on 992 examples. Treat its route as a heuristic, not a security boundary or sole safety control. Add deterministic policy checks, timeouts, input-length limits, and a conservative fallback in production.

The upstream model card does not declare a license at the time of this conversion. This repository does not add or replace upstream rights. Confirm use and redistribution terms with SupraLabs before commercial deployment or redistribution.

This conversion was produced independently and is not an official SupraLabs release. Refer to the upstream model card for intended input formatting and training details.