Llama Prompt Guard 2 86M Onnx

🚀 Llama-Prompt-Guard-2-86M-onnx

This repository offers an ONNX converted and quantized version of meta-llama/Llama-Prompt-Guard-2-86M, which is useful for text classification tasks.

🚀 Quick Start

Prerequisites

• Ensure you have the necessary libraries installed. You can install them using pip install transformers optimum[onnxruntime] numpy.

Basic Usage

from transformers import AutoTokenizer
from optimum.onnxruntime import ORTModelForSequenceClassification
import numpy as np

# Load model and tokenizer using optimum
model = ORTModelForSequenceClassification.from_pretrained("gravitee-io/Llama-Prompt-Guard-2-86M-onnx", file_name="model.quant.onnx")
tokenizer = AutoTokenizer.from_pretrained("gravitee-io/Llama-Prompt-Guard-2-86M-onnx")

# Tokenize input
text = "Your comment here"
inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)

# Run inference
outputs = model(**inputs)
logits = outputs.logits

# Optional: convert to probabilities
probs = 1 / (1 + np.exp(-logits))
print(probs)

✨ Features

• Multilingual Support: Supports multiple languages including English, French, German, Hindi, Italian, Portuguese, Spanish, and Thai.
• Powerful Foundation Model: Built on Meta LLaMA models, providing strong classification capabilities.
• Efficient Inference: Utilizes ONNX and ONNX Runtime for efficient model export and inference.

📦 Installation

The installation mainly involves installing the required Python libraries. You can use the following command:

pip install transformers optimum[onnxruntime] numpy

💻 Usage Examples

Basic Usage

from transformers import AutoTokenizer
from optimum.onnxruntime import ORTModelForSequenceClassification
import numpy as np

# Load model and tokenizer using optimum
model = ORTModelForSequenceClassification.from_pretrained("gravitee-io/Llama-Prompt-Guard-2-86M-onnx", file_name="model.quant.onnx")
tokenizer = AutoTokenizer.from_pretrained("gravitee-io/Llama-Prompt-Guard-2-86M-onnx")

# Tokenize input
text = "Your comment here"
inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)

# Run inference
outputs = model(**inputs)
logits = outputs.logits

# Optional: convert to probabilities
probs = 1 / (1 + np.exp(-logits))
print(probs)

📚 Documentation

Built With

• Meta LLaMA: Foundation model powering the classifier
  • meta-llama/Llama-Prompt-Guard-2-22M
  • meta-llama/Llama-Prompt-Guard-2-86M
• Hugging Face Transformers: Used for model and tokenizer loading.
• ONNX: Model export and runtime format.
• ONNX Runtime: Efficient inference backend.

Evaluation Dataset

We use the jackhhao/jailbreak-classification dataset for evaluation.

Evaluation Results

GitHub Repository

You can find the full source code, CLI tools, and evaluation scripts in the official GitHub repository.

📄 License

This project is under the Llama4 license.