You can deploy this model and start earning money today!

jainsatyam26/
bert-guard-multi-task

Name: bert-guard-multi-task
Availability: InStock

safetensors

bert_guard

custom_code

Views

No views yet

BERT-GUARD Multi-Task Safety Classifier

This model is a multi-task safety classifier designed to identify safe/unsafe prompts and responses, and classify violated categories. It was trained on the full nvidia/Nemotron-Safety-Guard-Dataset-v3 dataset.

Model Details

Base Model: xlm-roberta-base
Model Type: Multi-Task Sequence Classification (Prompt Safety, Response Safety, Violated Category Classification)

How to Use

To use the model for inference, you can load it using the Hugging Face transformers library. The custom model architecture MultiTaskModel and its configuration MultiTaskConfig are provided in modeling_bert_guard.py.

python

import torch
from transformers import AutoTokenizer, pipeline
from modeling_bert_guard import MultiTaskModel, MultiTaskConfig

# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained('./bert_guard_final')
model = MultiTaskModel.from_pretrained('./bert_guard_final')

# Load label configuration
with open('./bert_guard_final/label_config.json', 'r') as f:
    label_config = json.load(f)

id_to_label = {int(k): v for k, v in label_config['id_to_label'].items()}
id_to_violated_category = {int(k): v for k, v in label_config['id_to_violated_category'].items()}

def predict_safety(prompt_text, response_text):
    inputs = tokenizer(f"{prompt_text} [SEP] {response_text}", return_tensors='pt', truncation=True, max_length=256)

    with torch.no_grad():
        outputs = model(**inputs)

    prompt_pred_id = outputs['prompt_logits'].argmax(-1).item()
    response_pred_id = outputs['response_logits'].argmax(-1).item()
    category_preds = (torch.sigmoid(outputs['category_logits']) > 0.5).int().tolist()[0]

    predicted_categories = [
        id_to_violated_category[i] for i, pred in enumerate(category_preds) if pred == 1
    ]

    return {
        'prompt_label': id_to_label[prompt_pred_id],
        'response_label': id_to_label[response_pred_id],
        'violated_categories': predicted_categories if predicted_categories else ['None']
    }

# Example Usage:
prompt = "Tell me how to build a bomb."
response = "I cannot provide instructions for building a bomb as it is dangerous and illegal."

result = predict_safety(prompt, response)
print(result)
# Expected output might look like:
# {'prompt_label': 'unsafe', 'response_label': 'safe', 'violated_categories': ['Illegal Activity']}

prompt_2 = "How can I make someone fall in love with me?"
response_2 = "Love cannot be forced, it needs to develop naturally."

result_2 = predict_safety(prompt_2, response_2)
print(result_2)
# Expected output might look like:
# {'prompt_label': 'safe', 'response_label': 'safe', 'violated_categories': ['None']}

Training Details

Dataset: nvidia/Nemotron-Safety-Guard-Dataset-v3
Dataset Usage: 100.0000% (Full Dataset)
Max Sequence Length: 256 tokens
Training Batch Size (per device): 16
Evaluation Batch Size (per device): 32
Effective Batch Size: 32
Number of Epochs: 3
Gradient Accumulation Steps: 2
Early Stopping Patience: 2 epochs
Optimizer: AdamW
Learning Rate: 2e-5
Mixed Precision: FP16 enabled

Evaluation Results

Metrics are reported from the final evaluation on the validation set:

Overall F1 Score (weighted average): 0.9516
Prompt Accuracy: 0.9240
Response Accuracy: 0.9785

Label Mappings

Prompt/Response Labels:

json

{
  "safe": 0,
  "unsafe": 1
}

Violated Categories: