Implementation:Hpcaitech ColossalAI CMMLUDataset
| Knowledge Sources | |
|---|---|
| Domains | Evaluation, Benchmarking |
| Last Updated | 2026-02-09 00:00 GMT |
Overview
CMMLUDataset is a dataset wrapper class that loads and converts the CMMLU (Chinese Massive Multitask Language Understanding) benchmark into the ColossalEval inference format, spanning 67 subjects across diverse academic and professional domains.
Description
The class extends BaseDataset and provides a static load method that reads CSV files from "dev" and "test" subdirectories. Each file corresponds to a subject mapped through the cmmlu_subject_mapping dictionary, which translates English subject keys to their Chinese equivalents. Questions are formatted as Chinese single-choice prompts with four options (A-D), and the module supports few-shot evaluation by prepending dev-split examples. Default inference kwargs use loss calculation with all_classes set to ["A", "B", "C", "D"].
Usage
Use this class when you need to evaluate a language model on the CMMLU benchmark within the ColossalEval framework. It expects the CMMLU dataset organized with "dev" and "test" subdirectories containing per-subject CSV files.
Code Reference
Source Location
- Repository: Hpcaitech_ColossalAI
- File: applications/ColossalEval/colossal_eval/dataset/cmmlu.py
- Lines: 1-145
Signature
class CMMLUDataset(BaseDataset):
@staticmethod
def load(path: str, logger: DistributedLogger, few_shot: bool, *args, **kwargs) -> List[Dict]:
Import
from colossal_eval.dataset.cmmlu import CMMLUDataset
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| path | str | Yes | Path to the directory containing "dev" and "test" subdirectories with per-subject CSV 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 subject categories with "data" (list of data samples with fields dataset, split, category, instruction, input, output, target) and "inference_kwargs" (calculate_loss=True, all_classes=["A","B","C","D"], language="Chinese", max_new_tokens=32) |
Usage Examples
from colossal_eval.dataset.cmmlu import CMMLUDataset
from colossalai.logging import DistributedLogger
logger = DistributedLogger("cmmlu")
dataset = CMMLUDataset(path="/path/to/cmmlu/data", logger=logger, few_shot=True)
dataset.save("/path/to/output.json")