export.py

"""
export.py — Convert nvidia/prompt-task-and-complexity-classifier to ONNX.

Produces (into ./onnx/):
  model.onnx        fp32 full graph (DeBERTa-v3 backbone + mean pooling + 8 heads)
  model_fp16.onnx   fp16 variant

The whole model is exported as a single graph whose outputs are the raw logits
of the 8 classification heads; post-processing is left to the consumer. The
batch and sequence axes are dynamic.
"""

import os

import numpy as np
import torch
import torch.nn as nn
from huggingface_hub import PyTorchModelHubMixin
from transformers import AutoConfig, AutoModel, AutoTokenizer

MODEL_NAME = "nvidia/prompt-task-and-complexity-classifier"
MODEL_REVISION = "fea1121511eafabaf7dd6fc66863dcb04f74defb"
BACKBONE_NAME = "microsoft/DeBERTa-v3-base"
BACKBONE_REVISION = "8ccc9b6f36199bec6961081d44eb72fb3f7353f3"
ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
OUT_DIR = os.path.join(ROOT_DIR, "onnx")
OPSET = 17

# The 8 heads in the order they are defined by config.target_sizes, i.e. the
# order the model's own process_logits() indexes them as logits[0..7].
OUTPUT_NAMES = [
    "task_type",            # 12 classes
    "creativity_scope",     # 3
    "reasoning",            # 2
    "contextual_knowledge",  # 2
    "number_of_few_shots",  # 6
    "domain_knowledge",     # 4
    "no_label_reason",      # 1
    "constraint_ct",        # 2
]


# --------------------------------------------------------------------------- #
# Model definition — copied verbatim from the upstream model card.
# --------------------------------------------------------------------------- #
class MeanPooling(nn.Module):
    def __init__(self):
        super(MeanPooling, self).__init__()

    def forward(self, last_hidden_state, attention_mask):
        input_mask_expanded = (
            attention_mask.unsqueeze(-1).expand(last_hidden_state.size()).float()
        )
        sum_embeddings = torch.sum(last_hidden_state * input_mask_expanded, 1)

        sum_mask = input_mask_expanded.sum(1)
        sum_mask = torch.clamp(sum_mask, min=1e-9)

        mean_embeddings = sum_embeddings / sum_mask
        return mean_embeddings


class MulticlassHead(nn.Module):
    def __init__(self, input_size, num_classes):
        super(MulticlassHead, self).__init__()
        self.fc = nn.Linear(input_size, num_classes)

    def forward(self, x):
        x = self.fc(x)
        return x


class CustomModel(nn.Module, PyTorchModelHubMixin):
    def __init__(self, target_sizes, task_type_map, weights_map, divisor_map):
        super(CustomModel, self).__init__()

        self.backbone = AutoModel.from_pretrained(
            BACKBONE_NAME,
            revision=BACKBONE_REVISION,
        )
        self.target_sizes = target_sizes.values()
        self.task_type_map = task_type_map
        self.weights_map = weights_map
        self.divisor_map = divisor_map

        self.heads = [
            MulticlassHead(self.backbone.config.hidden_size, sz)
            for sz in self.target_sizes
        ]

        for i, head in enumerate(self.heads):
            self.add_module(f"head_{i}", head)

        self.pool = MeanPooling()

    def compute_results(self, preds, target, decimal=4):
        if target == "task_type":
            task_type = {}

            top2_indices = torch.topk(preds, k=2, dim=1).indices
            softmax_probs = torch.softmax(preds, dim=1)
            top2_probs = softmax_probs.gather(1, top2_indices)
            top2 = top2_indices.detach().cpu().tolist()
            top2_prob = top2_probs.detach().cpu().tolist()

            top2_strings = [
                [self.task_type_map[str(idx)] for idx in sample] for sample in top2
            ]
            top2_prob_rounded = [
                [round(value, 3) for value in sublist] for sublist in top2_prob
            ]

            counter = 0
            for sublist in top2_prob_rounded:
                if sublist[1] < 0.1:
                    top2_strings[counter][1] = "NA"
                counter += 1

            task_type_1 = [sublist[0] for sublist in top2_strings]
            task_type_2 = [sublist[1] for sublist in top2_strings]
            task_type_prob = [sublist[0] for sublist in top2_prob_rounded]

            return (task_type_1, task_type_2, task_type_prob)

        else:
            preds = torch.softmax(preds, dim=1)

            weights = np.array(self.weights_map[target])
            weighted_sum = np.sum(np.array(preds.detach().cpu()) * weights, axis=1)
            scores = weighted_sum / self.divisor_map[target]

            scores = [round(value, decimal) for value in scores]
            if target == "number_of_few_shots":
                scores = [x if x >= 0.05 else 0 for x in scores]
            return scores

    def process_logits(self, logits):
        result = {}

        task_type_logits = logits[0]
        task_type_results = self.compute_results(task_type_logits, target="task_type")
        result["task_type_1"] = task_type_results[0]
        result["task_type_2"] = task_type_results[1]
        result["task_type_prob"] = task_type_results[2]

        creativity_scope_logits = logits[1]
        result["creativity_scope"] = self.compute_results(creativity_scope_logits, target="creativity_scope")

        reasoning_logits = logits[2]
        result["reasoning"] = self.compute_results(reasoning_logits, target="reasoning")

        contextual_knowledge_logits = logits[3]
        result["contextual_knowledge"] = self.compute_results(contextual_knowledge_logits, target="contextual_knowledge")

        number_of_few_shots_logits = logits[4]
        result["number_of_few_shots"] = self.compute_results(number_of_few_shots_logits, target="number_of_few_shots")

        domain_knowledge_logits = logits[5]
        result["domain_knowledge"] = self.compute_results(domain_knowledge_logits, target="domain_knowledge")

        no_label_reason_logits = logits[6]
        result["no_label_reason"] = self.compute_results(no_label_reason_logits, target="no_label_reason")

        constraint_ct_logits = logits[7]
        result["constraint_ct"] = self.compute_results(constraint_ct_logits, target="constraint_ct")

        result["prompt_complexity_score"] = [
            round(
                0.35 * creativity
                + 0.25 * reasoning
                + 0.15 * constraint
                + 0.15 * domain_knowledge
                + 0.05 * contextual_knowledge
                + 0.05 * few_shots,
                5,
            )
            for creativity, reasoning, constraint, domain_knowledge, contextual_knowledge, few_shots in zip(
                result["creativity_scope"],
                result["reasoning"],
                result["constraint_ct"],
                result["domain_knowledge"],
                result["contextual_knowledge"],
                result["number_of_few_shots"],
            )
        ]

        return result

    def forward(self, batch):
        input_ids = batch["input_ids"]
        attention_mask = batch["attention_mask"]
        outputs = self.backbone(input_ids=input_ids, attention_mask=attention_mask)

        last_hidden_state = outputs.last_hidden_state
        mean_pooled_representation = self.pool(last_hidden_state, attention_mask)

        logits = [
            self.heads[k](mean_pooled_representation)
            for k in range(len(self.target_sizes))
        ]

        return self.process_logits(logits)


# --------------------------------------------------------------------------- #
# Export wrapper — a thin nn.Module whose forward returns the 8 raw logit
# tensors (before any numpy/python post-processing, which is not traceable).
# --------------------------------------------------------------------------- #
class ExportWrapper(nn.Module):
    def __init__(self, base: CustomModel):
        super().__init__()
        self.base = base

    def forward(self, input_ids, attention_mask):
        outputs = self.base.backbone(input_ids=input_ids, attention_mask=attention_mask)
        pooled = self.base.pool(outputs.last_hidden_state, attention_mask)
        return tuple(head(pooled) for head in self.base.heads)


def load_model() -> CustomModel:
    """Load the pinned upstream model as in the upstream model card."""
    config = AutoConfig.from_pretrained(MODEL_NAME, revision=MODEL_REVISION)
    model = CustomModel(
        target_sizes=config.target_sizes,
        task_type_map=config.task_type_map,
        weights_map=config.weights_map,
        divisor_map=config.divisor_map,
    ).from_pretrained(MODEL_NAME, revision=MODEL_REVISION)
    model.eval()
    return model


def main():
    os.makedirs(OUT_DIR, exist_ok=True)

    print(f"Loading upstream revision {MODEL_REVISION} ...")
    config = AutoConfig.from_pretrained(MODEL_NAME, revision=MODEL_REVISION)
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, revision=MODEL_REVISION)

    # Save the exact preprocessing/configuration artifacts used for the export so
    # a fresh run reproduces the complete inference package, not only the graph.
    config.save_pretrained(ROOT_DIR)
    tokenizer.save_pretrained(ROOT_DIR)

    model = load_model()
    wrapper = ExportWrapper(model).eval()

    # A representative example drives the trace; dynamic axes make the concrete
    # length here irrelevant to the exported graph.
    enc = tokenizer(
        "Write a Python script that uses a for loop.",
        return_tensors="pt",
        truncation=True,
        max_length=512,
    )
    args = (enc["input_ids"], enc["attention_mask"])

    dynamic_axes = {
        "input_ids": {0: "batch", 1: "sequence"},
        "attention_mask": {0: "batch", 1: "sequence"},
    }
    for name in OUTPUT_NAMES:
        dynamic_axes[name] = {0: "batch"}

    fp32_path = os.path.join(OUT_DIR, "model.onnx")
    print(f"Exporting fp32 ONNX (opset {OPSET}) -> {fp32_path}")
    with torch.no_grad():
        torch.onnx.export(
            wrapper,
            args,
            fp32_path,
            input_names=["input_ids", "attention_mask"],
            output_names=OUTPUT_NAMES,
            dynamic_axes=dynamic_axes,
            opset_version=OPSET,
            do_constant_folding=True,
        )
    print("  [ok] fp32 export complete")

    fp16_path = os.path.join(OUT_DIR, "model_fp16.onnx")
    print(f"Converting to fp16 -> {fp16_path}")
    import onnx
    from onnxconverter_common import float16

    m = onnx.load(fp32_path)
    # keep_io_types=True leaves inputs/outputs fp32; only internal weights/compute
    # become fp16, which is the most robust conversion for mixed runtimes.
    m16 = float16.convert_float_to_float16(m, keep_io_types=True)
    onnx.save(m16, fp16_path)
    print("  [ok] fp16 conversion complete")

    print("\nDone. Run `python verify.py` to validate the outputs.")


if __name__ == "__main__":
    main()