preflight/prompt-task-and-complexity-classifier-ONNX

Prompt Task & Complexity Classifier — ONNX

ONNX conversion of nvidia/prompt-task-and-complexity-classifier, an English multi-head classifier that predicts a prompt's task type and six complexity dimensions. The conversion is pinned to upstream revision fea1121 and is intended for direct use with ONNX Runtime.

This is a custom multi-output graph. It is not loadable through a standard Transformers or Transformers.js text-classification pipeline, and hosted Hugging Face inference is disabled. Direct ONNX Runtime examples are provided below.

Available files

file	precision	size	notes
onnx/model.onnx	fp32	702 MiB	source parameters exported to ONNX; validated against pinned PyTorch model
onnx/model_fp16.onnx	fp16	352 MiB	internal weights and compute converted to fp16; inputs remain int64 and outputs remain float32

Both models use ONNX opset 17 and have dynamic batch and sequence axes.

Input and output

Inputs: input_ids and attention_mask, both int64 with shape [batch, sequence]. Tokenize with the tokenizer in this repository and truncate to a maximum of 512 tokens.

Outputs: eight raw-logit tensors, one per classification head, in this order:

index	output name	classes
0	task_type	12
1	creativity_scope	3
2	reasoning	2
3	contextual_knowledge	2
4	number_of_few_shots	6
5	domain_knowledge	4
6	no_label_reason	1
7	constraint_ct	2

Post-processing consists of a softmax per head, weighted dimension scores, and the final prompt_complexity_score. It is reproduced in the Python example. Both files return float32 logits.

Usage

Python

Install the lightweight inference dependencies:

python -m pip install numpy onnxruntime transformers huggingface-hub

Then run:

import json
import numpy as np
import onnxruntime as ort
from huggingface_hub import hf_hub_download
from transformers import AutoTokenizer

repo_id = "preflight/prompt-task-and-complexity-classifier-ONNX"

tokenizer = AutoTokenizer.from_pretrained(repo_id)
session = ort.InferenceSession(
    hf_hub_download(repo_id, "onnx/model.onnx")  # or "onnx/model_fp16.onnx"
)
with open(hf_hub_download(repo_id, "config.json")) as f:
    config = json.load(f)

encoded = tokenizer(
    "Prompt: Write a Python script that uses a for loop.",
    return_tensors="np",
    truncation=True,
    max_length=512,
)
logits = session.run(
    None,
    {
        "input_ids": encoded["input_ids"].astype(np.int64),
        "attention_mask": encoded["attention_mask"].astype(np.int64),
    },
)

def softmax(x):
    exp = np.exp(x - x.max(axis=-1, keepdims=True))
    return exp / exp.sum(axis=-1, keepdims=True)

# This example post-processes one prompt. The ONNX graph itself supports batches.
heads = list(config["target_sizes"])
probabilities = {head: softmax(output[0]) for head, output in zip(heads, logits)}

# Task type: keep the top two classes, but report the runner-up only when its
# rounded probability is at least 0.1, matching the source model.
order = probabilities["task_type"].argsort()[::-1]
result = {
    "task_type_1": config["task_type_map"][str(order[0])],
    "task_type_2": (
        config["task_type_map"][str(order[1])]
        if round(float(probabilities["task_type"][order[1]]), 3) >= 0.1
        else "NA"
    ),
    "task_type_prob": round(float(probabilities["task_type"][order[0]]), 3),
}

# Remaining heads: weighted sums of class probabilities.
for head in heads[1:]:
    score = (
        np.dot(probabilities[head], config["weights_map"][head])
        / config["divisor_map"][head]
    )
    result[head] = round(float(score), 4)
if result["number_of_few_shots"] < 0.05:
    result["number_of_few_shots"] = 0.0

result["prompt_complexity_score"] = round(
    0.35 * result["creativity_scope"]
    + 0.25 * result["reasoning"]
    + 0.15 * result["constraint_ct"]
    + 0.15 * result["domain_knowledge"]
    + 0.05 * result["contextual_knowledge"]
    + 0.05 * result["number_of_few_shots"],
    5,
)

print(result["task_type_1"], result["prompt_complexity_score"])
# Code Generation 0.27823

JavaScript / Node.js

Install the tokenizer, Hub client, and runtime:

npm install @huggingface/transformers @huggingface/hub onnxruntime-node

import { AutoTokenizer } from "@huggingface/transformers";
import { downloadFileToCacheDir } from "@huggingface/hub";
import * as ort from "onnxruntime-node";

const repoId = "preflight/prompt-task-and-complexity-classifier-ONNX";

const tokenizer = await AutoTokenizer.from_pretrained(repoId);
const modelPath = await downloadFileToCacheDir({
  repo: repoId,
  path: "onnx/model_fp16.onnx", // or "onnx/model.onnx"
});
const session = await ort.InferenceSession.create(modelPath);

const { input_ids, attention_mask } = await tokenizer(
  "Prompt: Write a Python script that uses a for loop.",
  {
    truncation: true,
    max_length: 512
  }
);

const ids = input_ids.data;

// Transformers.js may truncate away the trailing [SEP]. Restore it to match
// the Python tokenizer's truncation behavior.
if (
  typeof tokenizer.sep_token_id === "number" &&
  ids.length === 512 &&
  Number(ids[ids.length - 1]) !== tokenizer.sep_token_id
) {
  ids[ids.length - 1] = BigInt(tokenizer.sep_token_id);
}

const logits = await session.run({
  input_ids: new ort.Tensor("int64", ids, input_ids.dims),
  attention_mask: new ort.Tensor("int64", attention_mask.data, attention_mask.dims),
});

// logits.task_type, logits.creativity_scope, ..., logits.constraint_ct
// Apply the same per-head post-processing shown in the Python example.

Transformers.js truncation note: With @huggingface/transformers 3.8.1, truncating an input to max_length may retain a content token in the final position instead of the trailing [SEP] token. The guard above restores the separator to match the Python tokenizer. Shorter inputs and inputs that already end in [SEP] are unchanged.

For browsers, use onnxruntime-web and pass the resolved model URL directly to InferenceSession.create. Browser compatibility and memory use depend on the selected ONNX Runtime Web execution provider and are not covered by the bundled Python validation suite.

Provenance and validation

• Source model: NVIDIA Prompt Task & Complexity Classifier v1.1, revision fea1121.
• Backbone architecture: microsoft/DeBERTa-v3-base, revision 8ccc9b6.
• Export: PyTorch 2.2.2, opset 17, followed by an fp16 internal-precision conversion with fp32 graph outputs retained.
• verify.py checks the shipped tokenizer/config against the pinned source, validates output ordering, compares fp32 logits and processed results with the PyTorch model, checks the documented example, and measures fp16 drift over four prompts spanning several task types and complexity levels.

Release validation with the pinned environment in requirements-conversion.txt produced:

check	result
shipped tokenizer/config vs pinned source	exact match
maximum fp32 raw-logit difference vs PyTorch	5.01e-06
maximum fp32 post-processed numeric difference vs PyTorch	0.00e+00
maximum fp16 post-processed numeric difference vs fp32	4.00e-04
fp16 task-label changes across the four validation prompts	0
documented example	Code Generation, complexity 0.27823

The source model's training data and reported 10-fold evaluation results are documented in the NVIDIA model card. No additional training or task-level evaluation was performed for this format conversion.

Limitations

• The model is intended for English prompts of at most 512 tokens and inherits the source model's data coverage, biases, and failure modes. The source model was trained on 4,024 human-annotated English prompts.
• Task and complexity outputs are model estimates, not objective measurements. Predictions may be unreliable for ambiguous, unusual, multilingual, or out-of-distribution prompts.
• The fp16 model can differ slightly from fp32 and may change a classification near a decision boundary. Use fp32 when maximum numerical fidelity matters.
• Consumers must apply the documented post-processing; the graph returns raw logits rather than the source model's final result dictionary.
• A standard Transformers pipeline and Hugging Face hosted inference providers are not supported for this custom multi-head graph.

Reproducing the conversion

The scripts are supporting conversion/validation tools; they are not required for normal ONNX inference. They download the pinned source checkpoint and backbone, which requires network access and several gigabytes of temporary cache space. The pinned environment was tested with Python 3.10.

python -m venv .venv
source .venv/bin/activate
python -m pip install -r requirements-conversion.txt
python export.py
python verify.py

export.py recreates both ONNX files and saves the pinned tokenizer and config at the repository root. verify.py exits with a nonzero status if a hard check fails.

License and attribution

The source model and this conversion are distributed under the NVIDIA Open Model License Agreement. The required redistribution attribution is included in NOTICE.

Model created by NVIDIA. This repository provides an ONNX conversion and does not claim authorship of the underlying model.