""" 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()