| 17 | |
| 18 | |
| 19 | def execute_config( |
| 20 | model: str, |
| 21 | task: str, |
| 22 | batch_size: int, |
| 23 | limit: int, |
| 24 | output_dir: str, |
| 25 | num_fewshot: int, |
| 26 | context_length: int = 1000, |
| 27 | answer_length: int = 50, |
| 28 | cutting_context: bool = False, |
| 29 | decode_mode: str = "default", |
| 30 | ): |
| 31 | # Save the original standard output |
| 32 | import subprocess |
| 33 | |
| 34 | output_dir = os.path.join(output_dir, model, task) |
| 35 | |
| 36 | args = [ |
| 37 | "lm_eval", |
| 38 | "--model", "based_lm", |
| 39 | "--model_args", f"checkpoint_name={model}", |
| 40 | "--tasks", task, |
| 41 | "--device", "cuda:0", |
| 42 | "--batch_size", str(batch_size), |
| 43 | "--log_samples", |
| 44 | "--output_path", output_dir, |
| 45 | "--decode_mode", decode_mode, |
| 46 | "--num_fewshot", str(num_fewshot), |
| 47 | # , |
| 48 | |
| 49 | ] |
| 50 | |
| 51 | if cutting_context: |
| 52 | args.extend(["--cutting_context"]) |
| 53 | args.extend(["--context_length", str(context_length)]) |
| 54 | args.extend(["--answer_length", str(answer_length)]) |
| 55 | args.extend(["--context_key", "text"]) |
| 56 | |
| 57 | if 'squad' not in task: |
| 58 | args.extend(["--answer_key", "key", "value"]) |
| 59 | else: |
| 60 | args.extend(["--answer_key", "value"]) |
| 61 | |
| 62 | if limit is not None: |
| 63 | args.extend(["--limit", str(limit)]) |
| 64 | |
| 65 | subprocess.run(args) |
| 66 | |
| 67 | print(f"Decoded with mode: {decode_mode}") |
| 68 | |
| 69 | |
| 70 | @click.command() |