code: def add_few_shot_to_test(dataset):
      categories = list(dataset["test"].keys())
      for category in categories:
          original_category = category.split("_")[0]
          # Add a 'few_shot_data' field to each category of the test set
          dataset["test"][category]["inference_kwargs"][
              "few_shot_data"
          ] = get_few_shot_data(
              dataset["dev"][original_category]["data"]
          )
      return dataset
  
file_path: colossalai/applications/ColossalEval/colossal_eval/dataset/safetybench_en.py
imports: import json
  import os
  from copy import deepcopy
  from typing import Dict, List
  from colossalai.logging import DistributedLogger
  from .base import BaseDataset
label_desc: Use image node_id 5 for calling a global function with example usage: add_few_shot_to_test(dataset) and returns: dataset