(
model: List[str],
task: List[str],
batch_size: int,
limit: Optional[int],
parallelize: bool,
gpus: str,
num_fewshot: int = 0,
output_dir: str = "output",
context_length: int = 1000,
answer_length: int = 50,
cutting_context: bool = False,
decode_mode: str = 'default'
)
| 79 | @click.option("--cutting_context", is_flag=True) |
| 80 | @click.option("--decode_mode", default="default", type=str) |
| 81 | def main( |
| 82 | model: List[str], |
| 83 | task: List[str], |
| 84 | batch_size: int, |
| 85 | limit: Optional[int], |
| 86 | parallelize: bool, |
| 87 | gpus: str, |
| 88 | num_fewshot: int = 0, |
| 89 | output_dir: str = "output", |
| 90 | context_length: int = 1000, |
| 91 | answer_length: int = 50, |
| 92 | cutting_context: bool = False, |
| 93 | decode_mode: str = 'default' |
| 94 | ): |
| 95 | |
| 96 | if limit < 0: limit = None |
| 97 | |
| 98 | if gpus is not None: |
| 99 | os.environ["CUDA_VISIBLE_DEVICES"] = gpus |
| 100 | |
| 101 | # Load the given Python file as a module |
| 102 | configs = [ |
| 103 | {"model": m, "task": t} for m in model for t in task |
| 104 | ] |
| 105 | |
| 106 | use_ray = parallelize and len(configs) > 0 |
| 107 | if use_ray: |
| 108 | import ray |
| 109 | # ray was killing workers due to OOM, but it didn't seem to be necessary |
| 110 | os.environ["RAY_memory_monitor_refresh_ms"] = "0" |
| 111 | ray.init(ignore_reinit_error=True, log_to_driver=True) |
| 112 | |
| 113 | print(f"Running sweep with {len(configs)} configs") |
| 114 | |
| 115 | output_dir = f"{output_dir}/{datetime.now().strftime('%y-%m-%d_%H-%M')}" |
| 116 | |
| 117 | # Run each script in parallel using Ray |
| 118 | if not use_ray: |
| 119 | for config in configs: |
| 120 | execute_config( |
| 121 | **config, |
| 122 | batch_size=batch_size, |
| 123 | limit=limit, |
| 124 | output_dir=output_dir, |
| 125 | num_fewshot=num_fewshot, |
| 126 | context_length=context_length, |
| 127 | answer_length=answer_length, |
| 128 | cutting_context=cutting_context, |
| 129 | decode_mode=decode_mode, |
| 130 | ) |
| 131 | else: |
| 132 | completed = 0 |
| 133 | total = len(configs) |
| 134 | print(f"Completed: {completed} ({completed / total:0.1%}) | Total: {total}") |
| 135 | |
| 136 | remote = ray.remote(num_gpus=(1 // MAX_WORKERS_PER_GPU))(execute_config) |
| 137 | futures = [remote.remote( |
| 138 | **config, batch_size=batch_size, limit=limit, output_dir=output_dir, num_fewshot=num_fewshot, |
no test coverage detected