| |
| |
| |
| """Run released pre-decoders on the five T0-T4 simulated noise tasks.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import sys |
| from pathlib import Path |
|
|
| CODE_ROOT = Path(__file__).resolve().parents[1] |
| if str(CODE_ROOT) not in sys.path: |
| sys.path.insert(0, str(CODE_ROOT)) |
|
|
| from scripts.qadapt_example_utils import ( |
| InferenceJob, |
| TASK_CONFIGS, |
| add_common_inference_args, |
| build_paired_command, |
| parse_gpus, |
| run_jobs, |
| ) |
|
|
|
|
| TASK_BY_ID = { |
| f"T{index}": (task_key, config_name) |
| for index, (task_key, config_name) in enumerate(TASK_CONFIGS) |
| } |
|
|
|
|
| def parse_distances(value: str) -> list[int]: |
| result = [int(item.strip()) for item in value.split(",") if item.strip()] |
| if not result or result != sorted(set(result)): |
| raise argparse.ArgumentTypeError( |
| "distances must be a non-empty, increasing comma-separated list" |
| ) |
| return result |
|
|
|
|
| def parse_tasks(value: str) -> list[str]: |
| result = [item.strip().upper() for item in value.split(",") if item.strip()] |
| if not result or len(result) != len(set(result)): |
| raise argparse.ArgumentTypeError( |
| "tasks must be a non-empty comma-separated subset of T0,T1,T2,T3,T4" |
| ) |
| unknown = [item for item in result if item not in TASK_BY_ID] |
| if unknown: |
| raise argparse.ArgumentTypeError(f"unknown task(s): {','.join(unknown)}") |
| return result |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument( |
| "--distances", |
| type=parse_distances, |
| default=[9], |
| help=( |
| "Comma-separated distances. Use 7,9 with --tasks T0 for the " |
| "paper's mapped-noise geometry; the default is release coverage at d=9." |
| ), |
| ) |
| parser.add_argument( |
| "--tasks", |
| type=parse_tasks, |
| default=list(TASK_BY_ID), |
| help="Comma-separated task subset; defaults to T0,T1,T2,T3,T4.", |
| ) |
| parser.add_argument("--n-rounds", type=int, default=9) |
| add_common_inference_args( |
| parser, |
| default_output_dir=Path("outputs/examples/released_models/t0_t4"), |
| ) |
| return parser.parse_args() |
|
|
|
|
| def main() -> None: |
| args = parse_args() |
| jobs = [] |
| for distance in args.distances: |
| for task_id in args.tasks: |
| task_key, config_name = TASK_BY_ID[task_id] |
| label = f"d{distance}_{task_key}" |
| output_path = args.output_dir / f"d{distance}" / f"{task_key}.json" |
| jobs.append( |
| InferenceJob( |
| label=label, |
| command=build_paired_command( |
| args, |
| config_name=config_name, |
| output_path=output_path, |
| distance=distance, |
| n_rounds=args.n_rounds, |
| ), |
| output_path=output_path, |
| ) |
| ) |
| run_jobs( |
| jobs, |
| gpus=parse_gpus(args.gpus), |
| parallelism=args.parallelism, |
| resume=args.resume, |
| dry_run=args.dry_run, |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
