botirk/tiny-prompt-task-complexity-classifier

Quantized ONNX model for botirk/tiny-prompt-task-complexity-classifier

This repository contains the quantized ONNX version of the \ nvidia/prompt-task-and-complexity-classifier model.

Model Description

This is a multi-headed model which classifies English text prompts across task \ types and complexity dimensions. This version has been quantized to INT8 \ using dynamic quantization with the 🤗 Optimum \ library, resulting in a smaller footprint and faster CPU inference.

For more details on the model architecture, tasks, and complexity dimensions, \ please refer to the [original model card]\ (https://huggingface.co/nvidia/prompt-task-and-complexity-classifier).

How to Use

You can use this model directly with optimum.onnxruntime for accelerated \ inference.

First, install the required libraries:

pip install optimum[onnxruntime] transformers

Then, you can use the model in a pipeline:

from optimum.onnxruntime import ORTModelForSequenceClassification
from transformers import AutoTokenizer, pipeline

repo_id = "botirk/tiny-prompt-task-complexity-classifier"
model = ORTModelForSequenceClassification.from_pretrained(repo_id)
tokenizer = AutoTokenizer.from_pretrained(repo_id)

# Note: The pipeline task is a simplification.
# For full multi-headed output, you need to process the logits manually.
classifier = pipeline("text-classification", model=model, tokenizer=tokenizer)

prompt = "Write a mystery set in a small town where an everyday object goes missing."
results = classifier(prompt)
print(results)