Implementation:Hpcaitech ColossalAI SafetyBenchENDataset
| Knowledge Sources | |
|---|---|
| Domains | Evaluation, Benchmarking |
| Last Updated | 2026-02-09 00:00 GMT |
Overview
SafetyBenchENDataset is a dataset wrapper class that loads and converts the English portion of the SafetyBench safety evaluation benchmark into the ColossalEval inference format, assessing model safety via multiple-choice questions.
Description
The class extends BaseDataset and provides a static load method that reads English dev and test JSON files from the SafetyBench dataset. Questions are formatted as multiple-choice with four options (A-D), and options with fewer than four choices are padded with "NULL" entries when PAD_CHOICES is enabled (default True). The module includes helper functions process_test for formatting test samples by category, process_dev for formatting dev samples with correct answers, get_few_shot_data for assembling few-shot examples, and add_few_shot_to_test for injecting dev examples into the test inference kwargs.
Usage
Use this class when you need to evaluate a language model on English safety-related multiple-choice questions within the ColossalEval framework. It supports few-shot evaluation using dev-split examples.
Code Reference
Source Location
- Repository: Hpcaitech_ColossalAI
- File: applications/ColossalEval/colossal_eval/dataset/safetybench_en.py
- Lines: 1-152
Signature
class SafetyBenchENDataset(BaseDataset):
@staticmethod
def load(path: str, logger: DistributedLogger, few_shot: bool, *args, **kwargs) -> List[Dict]:
Import
from colossal_eval.dataset.safetybench_en import SafetyBenchENDataset
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| path | str | Yes | Path to the directory containing "dev_en.json" and "test_en.json" files |
| logger | DistributedLogger | Yes | Logger instance for distributed logging |
| few_shot | bool | Yes | Whether to prepend dev-split examples as few-shot demonstrations for the test split |
Outputs
| Name | Type | Description |
|---|---|---|
| dataset | Dict[str, Dict] | A nested dictionary with "dev" and "test" splits, each containing per-category entries with "data" (list of data samples with fields dataset, split, category, instruction, input, output, target, id) and "inference_kwargs" (calculate_loss=False, all_classes=["A","B","C","D"], language="English", max_new_tokens=32) |
Usage Examples
from colossal_eval.dataset.safetybench_en import SafetyBenchENDataset
from colossalai.logging import DistributedLogger
logger = DistributedLogger("safetybench_en")
dataset = SafetyBenchENDataset(path="/path/to/safetybench/data", logger=logger, few_shot=True)
dataset.save("/path/to/output.json")