nvidia/prompt-task-and-complexity-classifier

1---2tags:3- model_hub_mixin4- pytorch_model_hub_mixin5license: other6---7 8# NemoCurator Prompt Task and Complexity Classifier9 10# Model Overview11 12<img src="https://huggingface.co/nvidia/prompt-task-and-complexity-classifier/resolve/main/image.png" alt="image" style="width:600px;">13 14This is a multi-headed model which classifies English text prompts across task types and complexity dimensions. Tasks are classified across 11 common categories. Complexity is evaluated across 6 dimensions and ensembled to create an overall complexity score. Further information on the taxonomies can be found below.15 16This model is ready for commercial use.17 18**Task types:**19* Open QA: A question where the response is based on general knowledge20* Closed QA: A question where the response is based on text/data provided with the prompt21* Summarization22* Text Generation23* Code Generation24* Chatbot25* Classification26* Rewrite27* Brainstorming28* Extraction29* Other30 31**Complexity dimensions:**32* Overall Complexity Score: The weighted sum of the complexity dimensions. Calculated as 0.35\*CreativityScore + 0.25\*ReasoningScore + 0.15\*ConstraintScore + 0.15\*DomainKnowledgeScore + 0.05\*ContextualKnowledgeScore + 0.05\*NumberOfFewShots33* Creativity: The level of creativity needed to respond to a prompt. Score range of 0-1, with a higher score indicating more creativity.34* Reasoning: The extent of logical or cognitive effort required to respond to a prompt. Score range of 0-1, with a higher score indicating more reasoning35* Contextual Knowledge: The background information necessary to respond to a prompt. Score range of 0-1, with a higher score indicating more contextual knowledge required outside of prompt.36* Domain Knowledge: The amount of specialized knowledge or expertise within a specific subject area needed to respond to a prompt. Score range of 0-1, with a higher score indicating more domain knowledge is required.37* Constraints: The number of constraints or conditions provided with the prompt. Score range of 0-1, with a higher score indicating more constraints in the prompt.38* Number of Few Shots: The number of examples provided with the prompt. Score range of 0-n, with a higher score indicating more examples provided in the prompt.39 40# License41This model is released under the [NVIDIA Open Model License Agreement](https://developer.download.nvidia.com/licenses/nvidia-open-model-license-agreement-june-2024.pdf).42 43# Model Architecture44The model architecture uses a DeBERTa backbone and incorporates multiple classification heads, each dedicated to a task categorization or complexity dimension. This approach enables the training of a unified network, allowing it to predict simultaneously during inference. Deberta-v3-base can theoretically handle up to 12k tokens, but default context length is set at 512 tokens.45 46# How to Use in NVIDIA NeMo Curator47NeMo Curator improves generative AI model accuracy by processing text, image, and video data at scale for training and customization. It also provides pre-built pipelines for generating synthetic data to customize and evaluate generative AI systems.48 49The inference code for this model is available through the NeMo Curator GitHub repository. Check out this [example notebook](https://github.com/NVIDIA-NeMo/Curator/blob/main/tutorials/text/distributed-data-classification/prompt-task-complexity-classification.ipynb) to get started.50 51# Input & Output52## Input53* Input Type: Text54* Input Format: String55* Input Parameters: 1D56* Other Properties Related to Input: Token Limit of 512 tokens57 58## Output59* Output Type: Text/Numeric Classifications60* Output Format: String & Numeric61* Output Parameters: 1D62* Other Properties Related to Output: None63 64## Examples65 66```67Prompt: Write a mystery set in a small town where an everyday object goes missing, causing a ripple of curiosity and suspicion. Follow the investigation and reveal the surprising truth behind the disappearance.68```69 70| Task            | Complexity | Creativity | Reasoning | Contextual Knowledge | Domain Knowledge | Constraints | # of Few Shots |71|------------------|------------|------------|-----------|-----------------------|------------------|-------------|----------------|72| Text Generation | 0.472      | 0.867      | 0.056     | 0.048                | 0.226            | 0.785       | 0              |73 74```75Prompt: Antibiotics are a type of medication used to treat bacterial infections. They work by either killing the bacteria or preventing them from reproducing, allowing the body’s immune system to fight off the infection. Antibiotics are usually taken orally in the form of pills, capsules, or liquid solutions, or sometimes administered intravenously. They are not effective against viral infections, and using them inappropriately can lead to antibiotic resistance. Explain the above in one sentence.76```77 78| Task           | Complexity | Creativity | Reasoning | Contextual Knowledge | Domain Knowledge | Constraints | # of Few Shots |79|-----------------|------------|------------|-----------|-----------------------|------------------|-------------|----------------|80| Summarization  | 0.133      | 0.003      | 0.014     | 0.003                | 0.644            | 0.211       | 0              |81 82# Software Integration83* Runtime Engine: Python 3.10 and NeMo Curator84* Supported Hardware Microarchitecture Compatibility: NVIDIA GPU, Volta™ or higher (compute capability 7.0+), CUDA 12 (or above)85* Preferred/Supported Operating System(s): Ubuntu 22.04/20.0486 87# Model Version88NemoCurator Prompt Task and Complexity Classifier v1.189 90# Training, Testing, and Evaluation Datasets91## Training Data92* 4024 English prompts with task distribution outlined below93* Prompts were annotated by humans according to task and complexity taxonomies94 95Task distribution:96| Task            | Count |97|------------------|-------|98| Open QA         | 1214  |99| Closed QA       | 786   |100| Text Generation | 480   |101| Chatbot         | 448   |102| Classification  | 267   |103| Summarization   | 230   |104| Code Generation | 185   |105| Rewrite         | 169   |106| Other           | 104   |107| Brainstorming   | 81    |108| Extraction      | 60    |109| Total       | 4024  |110 111## Evaluation112For evaluation, Top-1 accuracy metric was used, which involves matching the category with the highest probability to the expected answer. Additionally, n-fold cross-validation was used to produce n different values for this metric to verify the consistency of the results. The table below displays the average of the top-1 accuracy values for the N folds calculated for each complexity dimension separately.113 114| | Task Accuracy    | Creative Accuracy | Reasoning Accuracy | Contextual Accuracy | FewShots Accuracy | Domain Accuracy | Constraint Accuracy |115|-|------------------|-------------------|--------------------|---------------------|-------------------|-----------------|---------------------|116| Average of 10 Folds | 0.981            | 0.996            | 0.997             | 0.981              | 0.979            | 0.937          | 0.991 |117 118# Inference119* Engine: PyTorch120* Test Hardware: A10G121 122# How to Use in Transformers123To use the prompt task and complexity classifier, use the following code:124 125```python126import numpy as np127import torch128import torch.nn as nn129from huggingface_hub import PyTorchModelHubMixin130from transformers import AutoConfig, AutoModel, AutoTokenizer131 132 133class MeanPooling(nn.Module):134    def __init__(self):135        super(MeanPooling, self).__init__()136 137    def forward(self, last_hidden_state, attention_mask):138        input_mask_expanded = (139            attention_mask.unsqueeze(-1).expand(last_hidden_state.size()).float()140        )141        sum_embeddings = torch.sum(last_hidden_state * input_mask_expanded, 1)142 143        sum_mask = input_mask_expanded.sum(1)144        sum_mask = torch.clamp(sum_mask, min=1e-9)145 146        mean_embeddings = sum_embeddings / sum_mask147        return mean_embeddings148 149 150class MulticlassHead(nn.Module):151    def __init__(self, input_size, num_classes):152        super(MulticlassHead, self).__init__()153        self.fc = nn.Linear(input_size, num_classes)154 155    def forward(self, x):156        x = self.fc(x)157        return x158 159 160class CustomModel(nn.Module, PyTorchModelHubMixin):161    def __init__(self, target_sizes, task_type_map, weights_map, divisor_map):162        super(CustomModel, self).__init__()163 164        self.backbone = AutoModel.from_pretrained("microsoft/DeBERTa-v3-base")165        self.target_sizes = target_sizes.values()166        self.task_type_map = task_type_map167        self.weights_map = weights_map168        self.divisor_map = divisor_map169 170        self.heads = [171            MulticlassHead(self.backbone.config.hidden_size, sz)172            for sz in self.target_sizes173        ]174 175        for i, head in enumerate(self.heads):176            self.add_module(f"head_{i}", head)177 178        self.pool = MeanPooling()179 180    def compute_results(self, preds, target, decimal=4):181        if target == "task_type":182            task_type = {}183 184            top2_indices = torch.topk(preds, k=2, dim=1).indices185            softmax_probs = torch.softmax(preds, dim=1)186            top2_probs = softmax_probs.gather(1, top2_indices)187            top2 = top2_indices.detach().cpu().tolist()188            top2_prob = top2_probs.detach().cpu().tolist()189 190            top2_strings = [191                [self.task_type_map[str(idx)] for idx in sample] for sample in top2192            ]193            top2_prob_rounded = [194                [round(value, 3) for value in sublist] for sublist in top2_prob195            ]196 197            counter = 0198            for sublist in top2_prob_rounded:199                if sublist[1] < 0.1:200                    top2_strings[counter][1] = "NA"201                counter += 1202 203            task_type_1 = [sublist[0] for sublist in top2_strings]204            task_type_2 = [sublist[1] for sublist in top2_strings]205            task_type_prob = [sublist[0] for sublist in top2_prob_rounded]206 207            return (task_type_1, task_type_2, task_type_prob)208 209        else:210            preds = torch.softmax(preds, dim=1)211 212            weights = np.array(self.weights_map[target])213            weighted_sum = np.sum(np.array(preds.detach().cpu()) * weights, axis=1)214            scores = weighted_sum / self.divisor_map[target]215 216            scores = [round(value, decimal) for value in scores]217            if target == "number_of_few_shots":218                scores = [x if x >= 0.05 else 0 for x in scores]219            return scores220 221    def process_logits(self, logits):222        result = {}223 224        # Round 1: "task_type"225        task_type_logits = logits[0]226        task_type_results = self.compute_results(task_type_logits, target="task_type")227        result["task_type_1"] = task_type_results[0]228        result["task_type_2"] = task_type_results[1]229        result["task_type_prob"] = task_type_results[2]230 231        # Round 2: "creativity_scope"232        creativity_scope_logits = logits[1]233        target = "creativity_scope"234        result[target] = self.compute_results(creativity_scope_logits, target=target)235 236        # Round 3: "reasoning"237        reasoning_logits = logits[2]238        target = "reasoning"239        result[target] = self.compute_results(reasoning_logits, target=target)240 241        # Round 4: "contextual_knowledge"242        contextual_knowledge_logits = logits[3]243        target = "contextual_knowledge"244        result[target] = self.compute_results(245            contextual_knowledge_logits, target=target246        )247 248        # Round 5: "number_of_few_shots"249        number_of_few_shots_logits = logits[4]250        target = "number_of_few_shots"251        result[target] = self.compute_results(number_of_few_shots_logits, target=target)252 253        # Round 6: "domain_knowledge"254        domain_knowledge_logits = logits[5]255        target = "domain_knowledge"256        result[target] = self.compute_results(domain_knowledge_logits, target=target)257 258        # Round 7: "no_label_reason"259        no_label_reason_logits = logits[6]260        target = "no_label_reason"261        result[target] = self.compute_results(no_label_reason_logits, target=target)262 263        # Round 8: "constraint_ct"264        constraint_ct_logits = logits[7]265        target = "constraint_ct"266        result[target] = self.compute_results(constraint_ct_logits, target=target)267 268        # Round 9: "prompt_complexity_score"269        result["prompt_complexity_score"] = [270            round(271                0.35 * creativity272                + 0.25 * reasoning273                + 0.15 * constraint274                + 0.15 * domain_knowledge275                + 0.05 * contextual_knowledge276                + 0.05 * few_shots,277                5,278            )279            for creativity, reasoning, constraint, domain_knowledge, contextual_knowledge, few_shots in zip(280                result["creativity_scope"],281                result["reasoning"],282                result["constraint_ct"],283                result["domain_knowledge"],284                result["contextual_knowledge"],285                result["number_of_few_shots"],286            )287        ]288 289        return result290 291    def forward(self, batch):292        input_ids = batch["input_ids"]293        attention_mask = batch["attention_mask"]294        outputs = self.backbone(input_ids=input_ids, attention_mask=attention_mask)295 296        last_hidden_state = outputs.last_hidden_state297        mean_pooled_representation = self.pool(last_hidden_state, attention_mask)298 299        logits = [300            self.heads[k](mean_pooled_representation)301            for k in range(len(self.target_sizes))302        ]303 304        return self.process_logits(logits)305 306 307config = AutoConfig.from_pretrained("nvidia/prompt-task-and-complexity-classifier")308tokenizer = AutoTokenizer.from_pretrained(309    "nvidia/prompt-task-and-complexity-classifier"310)311model = CustomModel(312    target_sizes=config.target_sizes,313    task_type_map=config.task_type_map,314    weights_map=config.weights_map,315    divisor_map=config.divisor_map,316).from_pretrained("nvidia/prompt-task-and-complexity-classifier")317model.eval()318 319prompt = ["Prompt: Write a Python script that uses a for loop."]320 321encoded_texts = tokenizer(322    prompt,323    return_tensors="pt",324    add_special_tokens=True,325    max_length=512,326    padding="max_length",327    truncation=True,328)329 330result = model(encoded_texts)331print(result)332# {'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]}333```334 335# References336* [DeBERTaV3: Improving DeBERTa using ELECTRA-Style Pre-Training with Gradient-Disentangled Embedding Sharing](https://arxiv.org/abs/2111.09543)337* [DeBERTa: Decoding-enhanced BERT with Disentangled Attention](https://github.com/microsoft/DeBERTa)338* [Training language models to follow instructions with human feedback](https://arxiv.org/pdf/2203.02155)339 340# Ethical Considerations341NVIDIA believes Trustworthy AI is a shared responsibility and we have established policies and practices to enable development for a wide array of AI applications.  When downloaded or used in accordance with our terms of service, developers should work with their internal model team to ensure this model meets requirements for the relevant industry and use case and addresses unforeseen product misuse.342 343Please report security vulnerabilities or NVIDIA AI Concerns [here](https://www.nvidia.com/en-us/support/submit-security-vulnerability).344