You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
49 lines
1.6 KiB
Python
49 lines
1.6 KiB
Python
from __future__ import annotations
|
|
import argparse
|
|
import multiprocessing as mp
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
sys.path.append(str(ROOT))
|
|
|
|
from src.common.config import Config
|
|
from src.common.experiment_paths import config_for_stage
|
|
from src.data.dataset_generation import ParallelDatasetGenerator
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--config", default=None)
|
|
parser.add_argument(
|
|
"--stage",
|
|
choices=["fixed_case", "case_neighborhood", "family_random", "family_random_hard", "family_random_v2_q"],
|
|
default=None,
|
|
)
|
|
parser.add_argument("--n-samples", type=int, default=None)
|
|
parser.add_argument("--n-workers", type=int, default=None)
|
|
parser.add_argument("--seed", type=int, default=None)
|
|
parser.add_argument("--method", type=str, default=None)
|
|
parser.add_argument("--dataset-tag", type=str, default=None, help="Optional tag injected into output dataset filename")
|
|
args = parser.parse_args()
|
|
|
|
config_path = args.config
|
|
if config_path is None:
|
|
config_path = str(config_for_stage(args.stage) or Path("configs/data_gen.yaml"))
|
|
|
|
cfg = Config(config_path)
|
|
cfg.ensure_dirs()
|
|
path = ParallelDatasetGenerator(cfg=cfg, n_workers=args.n_workers).generate(
|
|
n_samples=args.n_samples, method=args.method, random_seed=args.seed, dataset_tag=args.dataset_tag
|
|
)
|
|
print(path)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
mp.freeze_support()
|
|
try:
|
|
mp.set_start_method("spawn", force=True)
|
|
except Exception:
|
|
pass
|
|
main()
|