| 145 | |
| 146 | |
| 147 | def parse_args(): |
| 148 | parser = argparse.ArgumentParser() |
| 149 | parser.add_argument("--model", type=str, required=True, |
| 150 | help="Path to the model") |
| 151 | parser.add_argument("--dataset", type=str, required=True, |
| 152 | choices=["boolq", "piqa", "social_i_qa", "hellaswag", |
| 153 | "winogrande", "ARC-Challenge", "ARC-Easy", "openbookqa"], |
| 154 | help="Dataset to evaluate on") |
| 155 | parser.add_argument("--data_file", type=str, default=None, |
| 156 | help="Path to the dataset file") |
| 157 | parser.add_argument("--start", type=int, default=0, |
| 158 | help="Start index for evaluation") |
| 159 | parser.add_argument("--end", type=int, default=MAX_INT, |
| 160 | help="End index for evaluation") |
| 161 | parser.add_argument("--batch_size", type=int, default=32, |
| 162 | help="Batch size for evaluation") |
| 163 | parser.add_argument("--tensor_parallel_size", type=int, default=1, |
| 164 | help="Tensor parallel size for model") |
| 165 | parser.add_argument("--run_dir", type=str, |
| 166 | help="Directory containing the wandb run ID") |
| 167 | |
| 168 | args = parser.parse_args() |
| 169 | |
| 170 | # Set default data file path if not provided |
| 171 | if args.data_file is None: |
| 172 | args.data_file = f'data/commonsense/{args.dataset}/test.json' |
| 173 | |
| 174 | # Initialize wandb |
| 175 | if args.run_dir: |
| 176 | try: |
| 177 | with open(os.path.join(args.run_dir, "wandb_run_id.txt"), "r") as f: |
| 178 | wandb_run_id = f.read().strip() |
| 179 | wandb.init( |
| 180 | id=wandb_run_id, |
| 181 | project="project-name", |
| 182 | resume="must" |
| 183 | ) |
| 184 | except FileNotFoundError: |
| 185 | print("WandB run ID file not found, starting new run") |
| 186 | wandb.init(project="project-name") |
| 187 | |
| 188 | return args |
| 189 | |
| 190 | |
| 191 | if __name__ == "__main__": |