| """Baseline: nvidia/prompt-task-and-complexity-classifier (inference code from its model card) on a prompt set. |
| python prompt/run_nvidia.py IN.parquet OUT.parquet (adds nv_task, nv_creativity, nv_reasoning, nv_contextual_knowledge, nv_domain_knowledge, |
| nv_constraints, nv_complexity)""" |
| import sys |
| import numpy as np, pandas as pd, torch, torch.nn as nn |
| from huggingface_hub import PyTorchModelHubMixin |
| from transformers import AutoConfig, AutoModel, AutoTokenizer |
|
|
|
|
| class MeanPooling(nn.Module): |
| def forward(self, h, m): |
| m = m.unsqueeze(-1).expand(h.size()).float(); return (h * m).sum(1) / m.sum(1).clamp(min=1e-9) |
|
|
|
|
| class MulticlassHead(nn.Module): |
| def __init__(self, i, n): |
| super().__init__(); self.fc = nn.Linear(i, n) |
|
|
| def forward(self, x): |
| return self.fc(x) |
|
|
|
|
| class CustomModel(nn.Module, PyTorchModelHubMixin): |
| def __init__(self, target_sizes, task_type_map, weights_map, divisor_map): |
| super().__init__() |
| self.backbone = AutoModel.from_pretrained("microsoft/DeBERTa-v3-base"); self.target_sizes = target_sizes.values() |
| self.task_type_map, self.weights_map, self.divisor_map = task_type_map, weights_map, divisor_map |
| self.heads = [MulticlassHead(self.backbone.config.hidden_size, sz) for sz in self.target_sizes] |
| for i, h in enumerate(self.heads): |
| self.add_module(f"head_{i}", h) |
| self.pool = MeanPooling() |
|
|
| def forward(self, ids, att): |
| h = self.backbone(input_ids=ids, attention_mask=att).last_hidden_state; p = self.pool(h, att) |
| return [hd(p) for hd in self.heads] |
|
|
|
|
| import json, types |
| from huggingface_hub import hf_hub_download |
| cfg = types.SimpleNamespace(**json.load(open(hf_hub_download("nvidia/prompt-task-and-complexity-classifier", "config.json")))) |
| tok = AutoTokenizer.from_pretrained("nvidia/prompt-task-and-complexity-classifier") |
| m = CustomModel(target_sizes=cfg.target_sizes, task_type_map=cfg.task_type_map, weights_map=cfg.weights_map, divisor_map=cfg.divisor_map |
| ).from_pretrained("nvidia/prompt-task-and-complexity-classifier").cuda().eval() |
| names = list(cfg.target_sizes.keys()); print(names, cfg.task_type_map, flush=True) |
| df = pd.read_parquet(sys.argv[1]); out = {k: [] for k in ["nv_task", "creativity_scope", "reasoning", "contextual_knowledge", "number_of_few_shots", |
| "domain_knowledge", "constraint_ct"]} |
| score = lambda lo, t: (torch.softmax(lo, 1).cpu().numpy() * np.array(cfg.weights_map[t])).sum(1) / cfg.divisor_map[t] |
| with torch.no_grad(): |
| for s in range(0, len(df), 64): |
| e = tok(["Prompt: " + t for t in df.text.iloc[s:s + 64]], return_tensors="pt", max_length=512, padding=True, truncation=True) |
| lo = m(e["input_ids"].cuda(), e["attention_mask"].cuda()) |
| out["nv_task"] += [cfg.task_type_map[str(i)] for i in lo[0].argmax(1).tolist()] |
| for k, t in enumerate(names[1:], 1): |
| if t in out: |
| out[t] += list(score(lo[k].float(), t)) |
| df["nv_task"] = out["nv_task"] |
| for a, b in [("creativity_scope", "nv_creativity"), ("reasoning", "nv_reasoning"), ("contextual_knowledge", "nv_contextual_knowledge"), |
| ("domain_knowledge", "nv_domain_knowledge"), ("constraint_ct", "nv_constraints"), ("number_of_few_shots", "nv_few_shots")]: |
| df[b] = out[a] |
| df["nv_complexity"] = (0.35 * df.nv_creativity + 0.25 * df.nv_reasoning + 0.15 * df.nv_constraints + 0.15 * df.nv_domain_knowledge |
| + 0.05 * df.nv_contextual_knowledge + 0.05 * df.nv_few_shots) |
| df.to_parquet(sys.argv[2]); print(df.nv_task.value_counts().to_dict(), flush=True) |
|
|
