code: def create_json_map(
      test_files_to_run: List[str], json_output_file: str
  ):
      """
      Creates a map from a list of tests to run to easily split them by category, when running parallelism of slow tests.
  
      Args:
          test_files_to_run (`List[str]`): The list of tests to run.
          json_output_file (`str`): The path where to store the built json map.
      """
      if json_output_file is None:
          return
  
      test_map = {}
      for test_file in test_files_to_run:
          # `test_file` is a path to a test folder/file, starting with `tests/`. For example,
          #   - `tests/models/bert/test_modeling_bert.py` or `tests/models/bert`
          #   - `tests/trainer/test_trainer.py` or `tests/trainer`
          #   - `tests/test_modeling_common.py`
          names = test_file.split(os.path.sep)
          if names[1] == "models":
              # take the part like `models/bert` for modeling tests
              key = os.path.sep.join(names[1:3])
          elif len(names) > 2 or not test_file.endswith(".py"):
              # test folders under `tests` or python files under them
              # take the part like tokenization, `pipeline`, etc. for other test categories
              key = os.path.sep.join(names[1:2])
          else:
              # common test files directly under `tests/`
              key = "common"
  
          if key not in test_map:
              test_map[key] = []
          test_map[key].append(test_file)
  
      # sort the keys & values
      keys = sorted(test_map.keys())
      test_map = {k: " ".join(sorted(test_map[k])) for k in keys}
      with open(json_output_file, "w", encoding="UTF-8") as fp:
          json.dump(test_map, fp, ensure_ascii=False)
  
imports: import argparse
  import collections
  import importlib.util
  import json
  import os
  import re
  import tempfile
  from contextlib import contextmanager
  from pathlib import Path
  from typing import Dict, List, Optional, Tuple, Union
  from git import Repo
label_desc: Use image node_id 24 for calling a global function with example usage: create_json_map(test_files_to_run, json_output_file) without return types