gameboy_worlds.setup_data

  1"""
  2Syncs data with Hugging Face Hub repositories for easy sharing of files.
  3Run with:
  4```bash
  5python -m gameboy_worlds.setup_data --help
  6```
  7
  8or
  9```bash
 10python -m gameboy_worlds.setup_data push --game pokemon_red
 11```
 12
 13or perhaps:
 14```bash
 15python -m gameboy_worlds.setup_data pull --game pokemon_red
 16```
 17"""
 18
 19from gameboy_worlds.utils import log_error, log_info, log_warn, load_parameters
 20from gameboy_worlds import AVAILABLE_GAMES
 21import click
 22from huggingface_hub import HfApi
 23import os
 24
 25loaded_parameters = load_parameters()
 26repo_namespace = "DJ-Research"
 27
 28
 29def check_variant(game, parameters):
 30    if game.lower() not in AVAILABLE_GAMES + ["all"]:
 31        log_error(
 32            f"Game {game} is not in the list of available variants: {AVAILABLE_GAMES}. Add it to the emulation registry before you proceed.",
 33            parameters,
 34        )
 35    games = [game] if game != "all" else AVAILABLE_GAMES
 36    for game in games:
 37        if f"{game}_rom_data_path" not in parameters:
 38            log_error(
 39                f"{game}_rom_data not found in parameters. You must add it to a config file.",
 40                parameters,
 41            )
 42
 43
 44@click.command()
 45@click.option(
 46    "--game",
 47    type=str,
 48    required=True,
 49    help="The game or game variant to create the repo for.",
 50)
 51@click.pass_obj
 52def create_hub_repo(parameters, game):
 53    """
 54    Create a new repo on the Hugging Face Hub. This should only be called once for each game_variant.
 55    Once set up, the repo will contain a mirror of what we expect (not including the ROM file itself) in the rom_data directory of that variant.
 56    """
 57    check_variant(game, parameters)
 58    games = [game] if game != "all" else AVAILABLE_GAMES
 59    for game in games:
 60        repo_name = f"GameBoy-{game}"
 61        repo_id = f"{repo_namespace}/{repo_name}"
 62        api = parameters["api"]
 63        api.create_repo(
 64            repo_id=repo_id, repo_type="dataset", exist_ok=False, private=False
 65        )
 66        log_info(
 67            f"Successfully created repo {repo_id} on the Hugging Face Hub.", parameters
 68        )
 69
 70
 71@click.command()
 72@click.option(
 73    "--game",
 74    type=str,
 75    required=True,
 76    help="The game or game variant to sync the repo for.",
 77)
 78@click.pass_obj
 79def sync(parameters, game):
 80    """
 81    Update the local repo with data from the HuggingFace Hub.
 82    """
 83    check_variant(game, parameters)
 84    games = [game] if game != "all" else AVAILABLE_GAMES
 85    for game in games:
 86        repo_name = f"GameBoy-{game}"
 87        repo_id = f"{repo_namespace}/{repo_name}"
 88
 89        rom_data_path = parameters[f"{game}_rom_data_path"] + "/"
 90        if not os.path.exists(rom_data_path):
 91            os.makedirs(rom_data_path)
 92        api = parameters["api"]
 93        api.snapshot_download(
 94            repo_id=repo_id, repo_type="dataset", local_dir=rom_data_path
 95        )
 96        log_info(
 97            f"Tried to sync repo at {repo_namespace}/{repo_name} with directory {rom_data_path}. Check output above for success",
 98            parameters,
 99        )
100
101
102@click.command()
103@click.option(
104    "--game",
105    type=str,
106    required=True,
107    help="The game or game variant to push the repo for.",
108)
109@click.pass_obj
110def push_data_to_hub(parameters, game):
111    """
112    Upload local data to the Hugging Face Hub.
113    """
114    check_variant(game, parameters)
115    if game == "all":
116        raise ValueError(
117            f"Pushing to all is unsafe. Please specify a game variant to push."
118        )
119    games = [game] if game != "all" else AVAILABLE_GAMES
120    for game in games:
121        repo_name = f"GameBoy-{game}"
122        api = parameters["api"]
123        repo_id = f"{repo_namespace}/{repo_name}"
124        rom_data_path = parameters[f"{game}_rom_data_path"] + "/"
125        if not os.path.exists(rom_data_path):
126            log_error(
127                f"Rom data path {rom_data_path} does not exist. Cannot push to hub.",
128                parameters,
129            )
130        api.upload_large_folder(
131            repo_id=repo_id,
132            repo_type="dataset",
133            folder_path=rom_data_path,
134            ignore_patterns=["*.gb", "*.gbc"],
135        )
136        log_info(
137            f"Tried to push data from {rom_data_path} to repo {repo_namespace}/{repo_name} on the Hugging Face Hub.",
138            parameters,
139        )
140
141
142@click.group()
143@click.pass_context
144def main(ctx, **input_parameters):
145    api = HfApi()
146    loaded_parameters["api"] = api
147    ctx.obj = loaded_parameters
148
149
150main.add_command(create_hub_repo, name="init")
151main.add_command(sync, name="pull")
152main.add_command(push_data_to_hub, name="push")
153
154if __name__ == "__main__":
155    main()