Modelpublic

govtech/stsb-roberta-base-off-topic

sourceHugging Faceotherupdated 2y agoView on Hugging Face

4likes1.7kdownloads

Model Card

Off-Topic Classification Model

This model leverages a fine-tuned Cross Encoder STSB Roberta Base to perform binary classification, determining whether a user prompt is off-topic in relation to the system's intended purpose as defined by the system prompt.

Model Highlights

—Base Model: `stsb-roberta-base`
—Maximum Context Length: 514 tokens
—Task: Binary classification (on-topic/off-topic)

Performance

We evaluated our fine-tuned models on synthetic data modelling system and user prompt pairs reflecting real world enterprise use cases of LLMs. The dataset is available here.

Approach	Model	ROC-AUC	F1	Precision	Recall
👉 Fine-tuned bi-encoder classifier	jina-embeddings-v2-small-en	0.99	0.97	0.99	0.95
Fine-tuned cross-encoder classifier	stsb-roberta-base	0.99	0.99	0.99	0.99
Pre-trained cross-encoder	stsb-roberta-base	0.73	0.68	0.53	0.93
Prompt Engineering	GPT 4o (2024-08-06)	-	0.95	0.94	0.97
Prompt Engineering	GPT 4o Mini (2024-07-18)	-	0.91	0.85	0.91
Zero-shot Classification	GPT 4o Mini (2024-07-18)	0.99	0.97	0.95	0.99

Further evaluation results on additional synthetic and external datasets (e.g.,JailbreakBench, HarmBench, TrustLLM) are available in our technical report.

Usage

1.Clone this repository and install the required dependencies:

bash

    pip install -r requirements.txt

1.You can run the model using two options:

Option 1: Using inference_onnx.py with the ONNX Model.

        python inference_onnx.py '[
            ["System prompt example 1", "User prompt example 1"],
            ["System prompt example 2", "System prompt example 2]
        ]'

Option 2: Using inference_safetensors.py with PyTorch and SafeTensors.

        python inference_safetensors.py '[
            ["System prompt example 1", "User prompt example 1"],
            ["System prompt example 2", "System prompt example 2]
        ]'

Read more about this model in our technical report.