---
license: other
license_name: govtech-singapore
license_link: LICENSE
datasets:
- gabrielchua/off-topic
language:
- en
metrics:
- roc_auc
- f1
- precision
- recall
base_model:
- jinaai/jina-embeddings-v2-small-en
---

# Off-Topic Classification Model

This repository contains a fine-tuned **Jina Embeddings model** designed to perform binary classification. The model predicts whether a user prompt is **off-topic** based on the intended purpose defined in the system prompt.

## Model Highlights

- **Base Model**: [`jina-embeddings-v2-small-en`](https://huggingface.co/jinaai/jina-embeddings-v2-small-en)
- **Maximum Context Length**: 1024 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](https://huggingface.co/datasets/gabrielchua/off-topic).

| Approach                              | Model                          | ROC-AUC | F1   | Precision | Recall |
|---------------------------------------|--------------------------------|---------|------|-----------|--------|
| [Fine-tuned bi-encoder classifier](https://huggingface.co/govtech/jina-embeddings-v2-small-en-off-topic)      | jina-embeddings-v2-small-en    | 0.99    | 0.97 | 0.99      | 0.95   |
| 👉 [Fine-tuned cross-encoder classifier](https://huggingface.co/govtech/stsb-roberta-base-off-topic)   | 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](https://arxiv.org/abs/2411.12946).

## Usage
1. Clone this repository and install the required dependencies:

    ```bash
    pip install -r requirements.txt
    ```

2. 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](https://arxiv.org/abs/2411.12946).