1import numpy as np
2import torch
3import torch.nn as nn
4from huggingface_hub import PyTorchModelHubMixin
5from transformers import AutoConfig, AutoModel, AutoTokenizer
6
7
8class MeanPooling(nn.Module):
9 def __init__(self):
10 super(MeanPooling, self).__init__()
11
12 def forward(self, last_hidden_state, attention_mask):
13 input_mask_expanded = (
14 attention_mask.unsqueeze(-1).expand(last_hidden_state.size()).float()
15 )
16 sum_embeddings = torch.sum(last_hidden_state * input_mask_expanded, 1)
17
18 sum_mask = input_mask_expanded.sum(1)
19 sum_mask = torch.clamp(sum_mask, min=1e-9)
20
21 mean_embeddings = sum_embeddings / sum_mask
22 return mean_embeddings
23
24
25class MulticlassHead(nn.Module):
26 def __init__(self, input_size, num_classes):
27 super(MulticlassHead, self).__init__()
28 self.fc = nn.Linear(input_size, num_classes)
29
30 def forward(self, x):
31 x = self.fc(x)
32 return x
33
34
35class CustomModel(nn.Module, PyTorchModelHubMixin):
36 def __init__(self, target_sizes, task_type_map, weights_map, divisor_map):
37 super(CustomModel, self).__init__()
38
39 self.backbone = AutoModel.from_pretrained("microsoft/DeBERTa-v3-base")
40 self.target_sizes = target_sizes.values()
41 self.task_type_map = task_type_map
42 self.weights_map = weights_map
43 self.divisor_map = divisor_map
44
45 self.heads = [
46 MulticlassHead(self.backbone.config.hidden_size, sz)
47 for sz in self.target_sizes
48 ]
49
50 for i, head in enumerate(self.heads):
51 self.add_module(f"head_{i}", head)
52
53 self.pool = MeanPooling()
54
55 def compute_results(self, preds, target, decimal=4):
56 if target == "task_type":
57 task_type = {}
58
59 top2_indices = torch.topk(preds, k=2, dim=1).indices
60 softmax_probs = torch.softmax(preds, dim=1)
61 top2_probs = softmax_probs.gather(1, top2_indices)
62 top2 = top2_indices.detach().cpu().tolist()
63 top2_prob = top2_probs.detach().cpu().tolist()
64
65 top2_strings = [
66 [self.task_type_map[str(idx)] for idx in sample] for sample in top2
67 ]
68 top2_prob_rounded = [
69 [round(value, 3) for value in sublist] for sublist in top2_prob
70 ]
71
72 counter = 0
73 for sublist in top2_prob_rounded:
74 if sublist[1] < 0.1:
75 top2_strings[counter][1] = "NA"
76 counter += 1
77
78 task_type_1 = [sublist[0] for sublist in top2_strings]
79 task_type_2 = [sublist[1] for sublist in top2_strings]
80 task_type_prob = [sublist[0] for sublist in top2_prob_rounded]
81
82 return (task_type_1, task_type_2, task_type_prob)
83
84 else:
85 preds = torch.softmax(preds, dim=1)
86
87 weights = np.array(self.weights_map[target])
88 weighted_sum = np.sum(np.array(preds.detach().cpu()) * weights, axis=1)
89 scores = weighted_sum / self.divisor_map[target]
90
91 scores = [round(value, decimal) for value in scores]
92 if target == "number_of_few_shots":
93 scores = [x if x >= 0.05 else 0 for x in scores]
94 return scores
95
96 def process_logits(self, logits):
97 result = {}
98
99 # Round 1: "task_type"
100 task_type_logits = logits[0]
101 task_type_results = self.compute_results(task_type_logits, target="task_type")
102 result["task_type_1"] = task_type_results[0]
103 result["task_type_2"] = task_type_results[1]
104 result["task_type_prob"] = task_type_results[2]
105
106 # Round 2: "creativity_scope"
107 creativity_scope_logits = logits[1]
108 target = "creativity_scope"
109 result[target] = self.compute_results(creativity_scope_logits, target=target)
110
111 # Round 3: "reasoning"
112 reasoning_logits = logits[2]
113 target = "reasoning"
114 result[target] = self.compute_results(reasoning_logits, target=target)
115
116 # Round 4: "contextual_knowledge"
117 contextual_knowledge_logits = logits[3]
118 target = "contextual_knowledge"
119 result[target] = self.compute_results(
120 contextual_knowledge_logits, target=target
121 )
122
123 # Round 5: "number_of_few_shots"
124 number_of_few_shots_logits = logits[4]
125 target = "number_of_few_shots"
126 result[target] = self.compute_results(number_of_few_shots_logits, target=target)
127
128 # Round 6: "domain_knowledge"
129 domain_knowledge_logits = logits[5]
130 target = "domain_knowledge"
131 result[target] = self.compute_results(domain_knowledge_logits, target=target)
132
133 # Round 7: "no_label_reason"
134 no_label_reason_logits = logits[6]
135 target = "no_label_reason"
136 result[target] = self.compute_results(no_label_reason_logits, target=target)
137
138 # Round 8: "constraint_ct"
139 constraint_ct_logits = logits[7]
140 target = "constraint_ct"
141 result[target] = self.compute_results(constraint_ct_logits, target=target)
142
143 # Round 9: "prompt_complexity_score"
144 result["prompt_complexity_score"] = [
145 round(
146 0.35 * creativity
147 + 0.25 * reasoning
148 + 0.15 * constraint
149 + 0.15 * domain_knowledge
150 + 0.05 * contextual_knowledge
151 + 0.05 * few_shots,
152 5,
153 )
154 for creativity, reasoning, constraint, domain_knowledge, contextual_knowledge, few_shots in zip(
155 result["creativity_scope"],
156 result["reasoning"],
157 result["constraint_ct"],
158 result["domain_knowledge"],
159 result["contextual_knowledge"],
160 result["number_of_few_shots"],
161 )
162 ]
163
164 return result
165
166 def forward(self, batch):
167 input_ids = batch["input_ids"]
168 attention_mask = batch["attention_mask"]
169 outputs = self.backbone(input_ids=input_ids, attention_mask=attention_mask)
170
171 last_hidden_state = outputs.last_hidden_state
172 mean_pooled_representation = self.pool(last_hidden_state, attention_mask)
173
174 logits = [
175 self.heads[k](mean_pooled_representation)
176 for k in range(len(self.target_sizes))
177 ]
178
179 return self.process_logits(logits)
180
181
182config = AutoConfig.from_pretrained("nvidia/prompt-task-and-complexity-classifier")
183tokenizer = AutoTokenizer.from_pretrained(
184 "nvidia/prompt-task-and-complexity-classifier"
185)
186model = CustomModel(
187 target_sizes=config.target_sizes,
188 task_type_map=config.task_type_map,
189 weights_map=config.weights_map,
190 divisor_map=config.divisor_map,
191).from_pretrained("nvidia/prompt-task-and-complexity-classifier")
192model.eval()
193
194prompt = ["Prompt: Write a Python script that uses a for loop."]
195
196encoded_texts = tokenizer(
197 prompt,
198 return_tensors="pt",
199 add_special_tokens=True,
200 max_length=512,
201 padding="max_length",
202 truncation=True,
203)
204
205result = model(encoded_texts)
206print(result)
207# {'task_type_1': ['Code Generation'], 'task_type_2': ['Text Generation'], 'task_type_prob': [0.767], 'creativity_scope': [0.0826], 'reasoning': [0.0632], 'contextual_knowledge': [0.056], 'number_of_few_shots': [0], 'domain_knowledge': [0.9803], 'no_label_reason': [0.0], 'constraint_ct': [0.5578], 'prompt_complexity_score': [0.27822]}
