botirk/tiny-prompt-task-complexity-classifier

1---2license: apache-2.03language: en4library_name: optimum5tags:6- onnx7- quantized8- text-classification9- nvidia10- nemotron11pipeline_tag: text-classification12---13 14# Quantized ONNX model for botirk/tiny-prompt-task-complexity-classifier15 16This repository contains the quantized ONNX version of the \17[nvidia/prompt-task-and-complexity-classifier](https://huggingface.co/nvidia/prompt-task-and-complexity-classifier) model.18 19## Model Description20 21This is a multi-headed model which classifies English text prompts across task \22types and complexity dimensions. This version has been quantized to `INT8` \23using dynamic quantization with the [🤗 Optimum](https://github.com/huggingface/optimum) \24library, resulting in a smaller footprint and faster CPU inference.25 26For more details on the model architecture, tasks, and complexity dimensions, \27please refer to the [original model card]\28(https://huggingface.co/nvidia/prompt-task-and-complexity-classifier).29 30## How to Use31 32You can use this model directly with `optimum.onnxruntime` for accelerated \33inference.34 35First, install the required libraries:36```bash37pip install optimum[onnxruntime] transformers38```39 40Then, you can use the model in a pipeline:41```python42from optimum.onnxruntime import ORTModelForSequenceClassification43from transformers import AutoTokenizer, pipeline44 45repo_id = "botirk/tiny-prompt-task-complexity-classifier"46model = ORTModelForSequenceClassification.from_pretrained(repo_id)47tokenizer = AutoTokenizer.from_pretrained(repo_id)48 49# Note: The pipeline task is a simplification.50# For full multi-headed output, you need to process the logits manually.51classifier = pipeline("text-classification", model=model, tokenizer=tokenizer)52 53prompt = "Write a mystery set in a small town where an everyday object goes missing."54results = classifier(prompt)55print(results)56```57