| 17 | |
| 18 | |
| 19 | def execute_config( |
| 20 | model: str, |
| 21 | task: str, |
| 22 | batch_size: int, |
| 23 | limit: int, |
| 24 | output_dir: str, |
| 25 | context_length: int = 1000, |
| 26 | cutting_context: bool = False, |
| 27 | answer_length: int=50, |
| 28 | ): |
| 29 | # Save the original standard output |
| 30 | import subprocess |
| 31 | |
| 32 | output_dir = os.path.join(output_dir, model, task) |
| 33 | |
| 34 | if 'mamba' in model.lower() and 'rw' not in model.lower(): model_name = "mamba_ssm" |
| 35 | else: model_name = 'hf-auto' |
| 36 | |
| 37 | args = [ |
| 38 | "lm_eval", |
| 39 | "--model", f"{model_name}", |
| 40 | "--model_args", f"checkpoint_name={model}", |
| 41 | "--tasks", task, |
| 42 | "--device", "cuda:0", |
| 43 | "--batch_size", str(batch_size), |
| 44 | "--log_samples", |
| 45 | "--output_path", output_dir |
| 46 | ] |
| 47 | |
| 48 | if cutting_context: |
| 49 | args.extend(["--cutting_context"]) |
| 50 | args.extend(["--context_length", str(context_length)]) |
| 51 | args.extend(["--answer_length", str(answer_length)]) |
| 52 | args.extend(["--context_key", "text"]) |
| 53 | |
| 54 | if limit is not None: |
| 55 | args.extend(["--limit", str(limit)]) |
| 56 | |
| 57 | subprocess.run(args) |
| 58 | |
| 59 | |
| 60 | |