README.md

---
license: other
license_name: nvidia-open-model-license
license_link: https://huggingface.co/preflight/prompt-task-and-complexity-classifier-ONNX/blob/main/LICENSE.pdf
library_name: onnx
base_model: nvidia/prompt-task-and-complexity-classifier
language:
  - en
pipeline_tag: text-classification
inference: false
tags:
  - onnx
  - onnxruntime
  - prompt-classification
  - text-classification
---

# Prompt Task & Complexity Classifier — ONNX

ONNX conversion of
[`nvidia/prompt-task-and-complexity-classifier`](https://huggingface.co/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`](https://huggingface.co/nvidia/prompt-task-and-complexity-classifier/tree/fea1121511eafabaf7dd6fc66863dcb04f74defb)
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:

```bash
python -m pip install numpy onnxruntime transformers huggingface-hub
```

Then run:

```python
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:

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

```js
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`](https://huggingface.co/nvidia/prompt-task-and-complexity-classifier/tree/fea1121511eafabaf7dd6fc66863dcb04f74defb).
- Backbone architecture: `microsoft/DeBERTa-v3-base`, revision
  [`8ccc9b6`](https://huggingface.co/microsoft/deberta-v3-base/tree/8ccc9b6f36199bec6961081d44eb72fb3f7353f3).
- 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](https://huggingface.co/nvidia/prompt-task-and-complexity-classifier).
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.

```bash
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](LICENSE.pdf). The required redistribution
attribution is included in [`NOTICE`](NOTICE).

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