AIEnergyScore/launch-computation-example

1import argparse2import os3 4from datasets import load_dataset, Dataset5from huggingface_hub import HfApi6 7TOKEN = os.environ.get("DEBUG")8api = HfApi(token=TOKEN)9 10parser = argparse.ArgumentParser()11parser.add_argument(12    "--run_dir",13    default=None,14    type=str,15    required=True,16    help="Path to the run directory.",17)18parser.add_argument(19    "--model_name",20    default=None,21    type=str,22    required=True,23    help="Model to benchmark.",24)25 26args = parser.parse_args()27 28# Updating request29dataset = load_dataset("AIEnergyScore/requests_debug", split="test",30                       token=TOKEN).to_pandas()31 32# Set benchmark to failed33# TODO: This doesn't have to be try-except, we could actually check if the file is there.34try:35    # Read error message36    with open(f"{args.run_dir}/error.log", 'r') as file:37        for f in file.readlines():38            if 'Traceback (most recent call last):' in f:39                error_message = f40                dataset.loc[dataset["model"].isin([args.model_name]), [41                    'status']] = "FAILED"42                print("Status set to FAILED")43            else:44                dataset.loc[dataset["model"].isin([args.model_name]), [45                    'status']] = "COMPLETED"46    # Add a new column for the error message if necessary47except FileNotFoundError as e:48    print(f"Could not find {args.run_dir}/error.log")49 50updated_dataset = Dataset.from_pandas(dataset)51updated_dataset.push_to_hub("AIEnergyScore/requests_debug", split="test",52                            token=TOKEN)53