code: def rm_and_merge(
      dp_size: int,
      save_path: str,
      model_names: List[str],
      dataset_names: Dict[str, List],
      dataset_classes: Dict[str, List],
  ) -> None:
      """
      Remove inference result per rank and merge them into one file.
  
      Args:
          dp_size: Number of groups for data parallel.
          save_path: The folder for storing inference results.
          model_names: Names of models for inference.
          dataset_names: Names of dataset for inference.
          dataset_classes: Dataset class for different inference results. We need to save dataset class to smooth the evaluation process.
  
      """
  
      for model_name in model_names:
          for dataset_name, categories in dataset_names.items():
              all_answers_with_dataset_class = {}
              all_answers_with_dataset_class[
                  "dataset_class"
              ] = dataset_classes[dataset_name]
  
              all_answers = {}
              for category in categories:
                  all_answers[category] = {"data": []}
                  answers = {"data": []}
  
                  for r in range(dp_size):
                      directory = os.path.join(
                          save_path,
                          model_name,
                          f"{dataset_name}_{category}_inference_results_dp_rank{r}.json",
                      )
                      if not os.path.exists(directory):
                          raise Exception(
                              f"Directory {directory} not found. There may be an error during inference time."
                          )
                      else:
                          rank_answers = utils.jload(directory)
                          answers["data"].extend(rank_answers["data"])
                          answers["inference_kwargs"] = rank_answers[
                              "inference_kwargs"
                          ]
  
                  for r in range(dp_size):
                      try:
                          directory = os.path.join(
                              save_path,
                              model_name,
                              f"{dataset_name}_{category}_inference_results_dp_rank{r}.json",
                          )
                          os.remove(directory)
                      except Exception as e:
                          print(e)
  
                  all_answers[category] = answers
  
              all_answers_with_dataset_class[
                  "inference_results"
              ] = all_answers
  
              logger.info(
                  f"Save inference results of model {model_name} on dataset {dataset_name}."
              )
              utils.jdump(
                  all_answers_with_dataset_class,
                  os.path.join(
                      save_path,
                      model_name,
                      f"{dataset_name}_inference_results.json",
                  ),
              )
  
          logger.info(
              f"Save inference results of model {model_name} for all dataset."
          )
      logger.info(
          f"Save inference results of all models for all dataset."
      )
  
file_path: ColossalAI/applications/ColossalEval/examples/gpt_evaluation/inference.py
imports: import argparse
  import copy
  import os
  from typing import Dict, List
  import torch
  import torch.distributed
  from colossal_eval import dataset, models, utils
  import colossalai
  from colossalai.cluster import ProcessGroupMesh
  from colossalai.logging import get_dist_logger
  from colossalai.shardformer import ShardConfig
label_desc: Use image node_id 1 for calling a global function with example usage: rm_and_merge(dp_size, save_path, model_names, dataset_names, dataset_classes) without return types