| |
| """Calibrate outgoing-span confidence risk on three Extended-251 prompts.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import math |
| import os |
| import sys |
| import tempfile |
| from pathlib import Path |
| from typing import Any |
|
|
|
|
| def preparse_gpu() -> str: |
| parser = argparse.ArgumentParser(add_help=False) |
| parser.add_argument("--gpu", default="5") |
| args, _ = parser.parse_known_args() |
| os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu) |
| os.environ.setdefault( |
| "MPLCONFIGDIR", tempfile.mkdtemp(prefix="self_forcing_span_calibration_") |
| ) |
| return str(args.gpu) |
|
|
|
|
| PHYSICAL_GPU = preparse_gpu() |
|
|
| import torch |
| from omegaconf import OmegaConf |
| from safetensors.torch import load_file |
|
|
| REPO_ROOT = Path(__file__).resolve().parents[1] |
| if str(REPO_ROOT) not in sys.path: |
| sys.path.insert(0, str(REPO_ROOT)) |
|
|
| from predictor_training.confidence import PredictorConfidenceHead |
| from scripts import evaluate_single_block_fppf as base |
| from scripts import generate_vbench8_extended_strategies as extended |
| from utils.wan_wrapper import WanTextEncoder |
|
|
|
|
| MAPPING_DEFAULT = REPO_ROOT / "assets/vbench8_extended_subset_mapping.json" |
| EXPERIMENT_DEFAULT = ( |
| REPO_ROOT / "confidence_experiments/layer17_stage1_step2000_20260831" |
| ) |
| OUTPUT_DEFAULT = ( |
| REPO_ROOT |
| / "confidence_experiments/layer17_stage1_step2000_spanrisk_vbench3_20260901" |
| ) |
| BETAS = (0.0, 0.5, 1.0, 1.5, 2.0) |
| TARGETS = { |
| "step12": (6, 8, 10), |
| "step123": (6, 9, 12, 15), |
| } |
| CANDIDATE_STEPS = { |
| "step12": (1, 2), |
| "step123": (1, 2, 3), |
| } |
|
|
|
|
| def atomic_json(path: Path, value: Any) -> None: |
| path.parent.mkdir(parents=True, exist_ok=True) |
| temporary = path.with_suffix(path.suffix + ".tmp") |
| temporary.write_text( |
| json.dumps(value, ensure_ascii=False, indent=2, allow_nan=False) + "\n", |
| encoding="utf-8", |
| ) |
| os.replace(temporary, path) |
|
|
|
|
| def quantile(values: list[float], fraction: float) -> float: |
| if not values: |
| raise ValueError("Cannot take a quantile of an empty list") |
| ordered = sorted(values) |
| position = fraction * (len(ordered) - 1) |
| lower = int(math.floor(position)) |
| upper = int(math.ceil(position)) |
| if lower == upper: |
| return ordered[lower] |
| weight = position - lower |
| return ordered[lower] * (1.0 - weight) + ordered[upper] * weight |
|
|
|
|
| def beta_name(beta: float) -> str: |
| return str(beta).replace(".", "p") |
|
|
|
|
| def read_mapping_row(path: Path, global_index: int) -> dict[str, Any]: |
| rows = extended.read_mapping(path) |
| matches = [row for row in rows if int(row["global_index"]) == global_index] |
| if len(matches) != 1: |
| raise ValueError(f"Expected one mapping row for global index {global_index}") |
| return matches[0] |
|
|
|
|
| def load_models( |
| experiment_root: Path, device: torch.device |
| ) -> tuple[Any, Any, Any, Any, Any]: |
| config = OmegaConf.merge( |
| OmegaConf.load(REPO_ROOT / "configs/default_config.yaml"), |
| OmegaConf.load(REPO_ROOT / "configs/self_forcing_sid.yaml"), |
| ) |
| pipeline = base.build_pipeline( |
| config, |
| REPO_ROOT / "checkpoints/self_forcing_dmd.pt", |
| torch.nn.Identity(), |
| device, |
| ) |
| predictor = base.load_predictor( |
| pipeline.generator.model, |
| { |
| "source_layer": 17, |
| "weights": REPO_ROOT |
| / "training_runs/layer17_stage1_1000p_4gpu_b16_2000steps" |
| / "checkpoint_step_2000/predictor.safetensors", |
| "gate_mode": "baseline", |
| }, |
| device, |
| ) |
| text_encoder = WanTextEncoder().to(device=device, dtype=torch.bfloat16).eval() |
| text_encoder.requires_grad_(False) |
| step12_head = PredictorConfidenceHead(num_steps=2).to(device=device).eval() |
| step12_head.load_state_dict( |
| load_file( |
| str(experiment_root / "confidence_step12/confidence_best.safetensors"), |
| device="cpu", |
| ), |
| strict=True, |
| ) |
| step12_head.requires_grad_(False) |
| step123_head = PredictorConfidenceHead(num_steps=3).to(device=device).eval() |
| step123_head.load_state_dict( |
| load_file( |
| str(experiment_root / "confidence_step123/confidence_best.safetensors"), |
| device="cpu", |
| ), |
| strict=True, |
| ) |
| step123_head.requires_grad_(False) |
| return pipeline, predictor, text_encoder, step12_head, step123_head |
|
|
|
|
| def probe(args: argparse.Namespace) -> None: |
| if args.global_index is None: |
| raise ValueError("--global-index is required in probe mode") |
| row = read_mapping_row(args.mapping, args.global_index) |
| destination = args.output_root / "probes" / f"global_{args.global_index:04d}.json" |
| if destination.is_file() and not args.overwrite: |
| print(f"[cached] {destination}", flush=True) |
| return |
| device = torch.device("cuda") |
| torch.set_grad_enabled(False) |
| pipeline, predictor, text_encoder, step12_head, step123_head = load_models( |
| args.experiment_root, device |
| ) |
| prompt = str(row["extended_prompt"]) |
| conditional = text_encoder(text_prompts=[prompt]) |
| probes: dict[str, Any] = {} |
| for family, head in (("step12", step12_head), ("step123", step123_head)): |
| config = { |
| "name": f"{family}_span_probe", |
| "policy": "dynamic", |
| "candidate_steps": list(CANDIDATE_STEPS[family]), |
| "beta": 0.0, |
| "threshold": -math.inf, |
| "target_accepts": 0, |
| "head": family, |
| "risk_mode": "outgoing_span", |
| } |
| _, diagnostic = extended.generate_rollout( |
| pipeline=pipeline, |
| conditional_dict=conditional, |
| seed=args.seed, |
| device=device, |
| predictor=predictor, |
| head=head, |
| config=config, |
| ) |
| decisions = diagnostic.pop("decisions") |
| if any(decision["accepted"] for decision in decisions): |
| raise RuntimeError("Probe unexpectedly accepted a Predictor result") |
| probes[family] = {"generation": diagnostic, "decisions": decisions} |
| print( |
| f"[probe] global={args.global_index} family={family} " |
| f"decisions={len(decisions)}", |
| flush=True, |
| ) |
| atomic_json( |
| destination, |
| { |
| "status": "complete", |
| "physical_gpu": PHYSICAL_GPU, |
| "global_index": int(row["global_index"]), |
| "prompt_suite": row["prompt_suite"], |
| "suite_index": int(row["suite_index"]), |
| "prompt": prompt, |
| "seed": args.seed, |
| "probes": probes, |
| }, |
| ) |
|
|
|
|
| def prepare(args: argparse.Namespace) -> None: |
| paths = sorted((args.output_root / "probes").glob("global_*.json")) |
| if len(paths) != 3: |
| raise ValueError(f"Expected exactly three probe files, found {len(paths)}") |
| probes = [json.loads(path.read_text(encoding="utf-8")) for path in paths] |
| configurations: list[dict[str, Any]] = [] |
| span_values: dict[int, set[float]] = {1: set(), 2: set(), 3: set()} |
| for family in ("step12", "step123"): |
| decisions = [ |
| decision |
| for probe_record in probes |
| for decision in probe_record["probes"][family]["decisions"] |
| ] |
| expected = 3 * 6 * len(CANDIDATE_STEPS[family]) |
| if len(decisions) != expected: |
| raise ValueError( |
| f"Expected {expected} {family} decisions, got {len(decisions)}" |
| ) |
| for decision in decisions: |
| span_values[int(decision["step"])].add(float(decision["outgoing_span"])) |
| max_calls = 6 * len(CANDIDATE_STEPS[family]) |
| for beta in BETAS: |
| risks = [ |
| float(decision["predicted_local_error"]) |
| * float(decision["span_weight"]) |
| * (1.0 + beta * float(decision["chunk_alpha"])) |
| for decision in decisions |
| ] |
| for target in TARGETS[family]: |
| threshold = quantile(risks, target / max_calls) |
| configurations.append( |
| { |
| "name": f"{family}_span_b{beta_name(beta)}_k{target:02d}", |
| "policy": "dynamic", |
| "candidate_steps": list(CANDIDATE_STEPS[family]), |
| "beta": beta, |
| "threshold": threshold, |
| "target_accepts": target, |
| "head": family, |
| "risk_mode": "outgoing_span", |
| } |
| ) |
| actual_spans = { |
| str(step): sorted(values) for step, values in span_values.items() if values |
| } |
| expected_spans = {1: 104.1666, 2: 208.3333, 3: 625.0} |
| for step, expected in expected_spans.items(): |
| values = span_values[step] |
| if values and any(abs(value - expected) > 0.02 for value in values): |
| raise ValueError(f"Unexpected outgoing span at step {step}: {values}") |
| atomic_json( |
| args.output_root / "scan_configs.json", |
| { |
| "status": "ready", |
| "risk_definition": ( |
| "predicted_local_error * (outgoing_span / 1000) * " |
| "(1 + beta * chunk_alpha)" |
| ), |
| "probe_global_indices": [int(value["global_index"]) for value in probes], |
| "outgoing_spans": actual_spans, |
| "betas": list(BETAS), |
| "targets": {key: list(value) for key, value in TARGETS.items()}, |
| "num_configs": len(configurations), |
| "configs": configurations, |
| }, |
| ) |
| print(f"[prepare] configurations={len(configurations)}", flush=True) |
|
|
|
|
| def prepare_threshold_only(args: argparse.Namespace) -> None: |
| probe_root = args.probe_root or args.output_root |
| paths = sorted((probe_root / "probes").glob("global_*.json")) |
| if len(paths) != 3: |
| raise ValueError(f"Expected exactly three probe files, found {len(paths)}") |
| probes = [json.loads(path.read_text(encoding="utf-8")) for path in paths] |
| configurations: list[dict[str, Any]] = [] |
| span_values: dict[int, set[float]] = {1: set(), 2: set(), 3: set()} |
| beta = 2.0 |
| for family in ("step12", "step123"): |
| decisions = [ |
| decision |
| for probe_record in probes |
| for decision in probe_record["probes"][family]["decisions"] |
| ] |
| expected = 3 * 6 * len(CANDIDATE_STEPS[family]) |
| if len(decisions) != expected: |
| raise ValueError( |
| f"Expected {expected} {family} decisions, got {len(decisions)}" |
| ) |
| for decision in decisions: |
| span_values[int(decision["step"])].add(float(decision["outgoing_span"])) |
| max_calls = 6 * len(CANDIDATE_STEPS[family]) |
| risks = [ |
| float(decision["predicted_local_error"]) |
| * float(decision["span_weight"]) |
| * (1.0 + beta * float(decision["chunk_alpha"])) |
| for decision in decisions |
| ] |
| for rank in range(1, max_calls): |
| configurations.append( |
| { |
| "name": f"{family}_span_beta2_thr{rank:02d}", |
| "policy": "dynamic", |
| "candidate_steps": list(CANDIDATE_STEPS[family]), |
| "beta": beta, |
| "threshold": quantile(risks, rank / max_calls), |
| "target_accepts": rank, |
| "threshold_rank": rank, |
| "head": family, |
| "risk_mode": "outgoing_span", |
| } |
| ) |
| atomic_json( |
| args.output_root / "scan_configs.json", |
| { |
| "status": "ready", |
| "risk_definition": ( |
| "predicted_local_error * (outgoing_span / 1000) * " |
| "(1 + 2.0 * chunk_alpha)" |
| ), |
| "beta": beta, |
| "probe_global_indices": [int(value["global_index"]) for value in probes], |
| "outgoing_spans": { |
| str(step): sorted(values) |
| for step, values in span_values.items() |
| if values |
| }, |
| "threshold_search": "quantile rank 1..max_candidate_calls-1", |
| "targets": {key: list(value) for key, value in TARGETS.items()}, |
| "num_configs": len(configurations), |
| "configs": configurations, |
| }, |
| ) |
| print( |
| f"[prepare-threshold-only] beta={beta} configurations={len(configurations)}", |
| flush=True, |
| ) |
|
|
|
|
| def latent_nrmse(predicted: torch.Tensor, reference: torch.Tensor) -> float: |
| numerator = (predicted.float() - reference.float()).square().sum() |
| denominator = reference.float().square().sum().clamp_min(1e-8) |
| return float(torch.sqrt(numerator / denominator)) |
|
|
|
|
| def scan(args: argparse.Namespace) -> None: |
| if args.global_index is None: |
| raise ValueError("--global-index is required in scan mode") |
| row = read_mapping_row(args.mapping, args.global_index) |
| config_path = args.output_root / "scan_configs.json" |
| config_manifest = json.loads(config_path.read_text(encoding="utf-8")) |
| configurations = list(config_manifest["configs"]) |
| destination = args.output_root / "scans" / f"global_{args.global_index:04d}.json" |
| existing: dict[str, Any] = {} |
| if destination.is_file() and not args.overwrite: |
| existing = json.loads(destination.read_text(encoding="utf-8")) |
| if existing.get("status") == "complete": |
| print(f"[cached] {destination}", flush=True) |
| return |
| completed = {record["name"]: record for record in existing.get("results", [])} |
| device = torch.device("cuda") |
| torch.set_grad_enabled(False) |
| pipeline, predictor, text_encoder, step12_head, step123_head = load_models( |
| args.experiment_root, device |
| ) |
| prompt = str(row["extended_prompt"]) |
| conditional = text_encoder(text_prompts=[prompt]) |
| reference, reference_diagnostic = extended.generate_rollout( |
| pipeline=pipeline, |
| conditional_dict=conditional, |
| seed=args.seed, |
| device=device, |
| predictor=predictor, |
| head=None, |
| config={ |
| "name": "ffff", |
| "policy": "ffff", |
| "candidate_steps": [], |
| "beta": None, |
| "threshold": None, |
| "target_accepts": 0, |
| "head": None, |
| "risk_mode": "outgoing_span", |
| }, |
| ) |
| for position, config in enumerate(configurations, start=1): |
| name = str(config["name"]) |
| if name in completed: |
| continue |
| head = step12_head if config["head"] == "step12" else step123_head |
| latent, diagnostic = extended.generate_rollout( |
| pipeline=pipeline, |
| conditional_dict=conditional, |
| seed=args.seed, |
| device=device, |
| predictor=predictor, |
| head=head, |
| config=config, |
| ) |
| decisions = diagnostic.pop("decisions") |
| record = { |
| **config, |
| "accepted_predictor_calls": diagnostic["accepted_predictor_calls"], |
| "full_calls": diagnostic["full_calls"], |
| "policy_latency_ms": diagnostic["policy_latency_ms"], |
| "context_dit_time_ms": diagnostic["context_dit_time_ms"], |
| "latent_nrmse_vs_ffff": latent_nrmse(latent, reference), |
| "decisions": decisions, |
| } |
| completed[name] = record |
| atomic_json( |
| destination, |
| { |
| "status": "running", |
| "physical_gpu": PHYSICAL_GPU, |
| "global_index": int(row["global_index"]), |
| "prompt_suite": row["prompt_suite"], |
| "suite_index": int(row["suite_index"]), |
| "prompt": prompt, |
| "seed": args.seed, |
| "ffff": { |
| key: value |
| for key, value in reference_diagnostic.items() |
| if key != "decisions" |
| }, |
| "results": [completed[key] for key in sorted(completed)], |
| }, |
| ) |
| print( |
| f"[scan] global={args.global_index} {position}/{len(configurations)} " |
| f"name={name} K={record['accepted_predictor_calls']} " |
| f"latent_nrmse={record['latent_nrmse_vs_ffff']:.6f}", |
| flush=True, |
| ) |
| del latent |
| atomic_json( |
| destination, |
| { |
| "status": "complete", |
| "physical_gpu": PHYSICAL_GPU, |
| "global_index": int(row["global_index"]), |
| "prompt_suite": row["prompt_suite"], |
| "suite_index": int(row["suite_index"]), |
| "prompt": prompt, |
| "seed": args.seed, |
| "ffff": { |
| key: value |
| for key, value in reference_diagnostic.items() |
| if key != "decisions" |
| }, |
| "results": [completed[key] for key in sorted(completed)], |
| }, |
| ) |
|
|
|
|
| def summarize(args: argparse.Namespace) -> None: |
| paths = sorted((args.output_root / "scans").glob("global_*.json")) |
| if len(paths) != 3: |
| raise ValueError(f"Expected exactly three scan files, found {len(paths)}") |
| scans = [json.loads(path.read_text(encoding="utf-8")) for path in paths] |
| if any(scan.get("status") != "complete" for scan in scans): |
| raise RuntimeError("At least one scan is incomplete") |
| by_name = [ |
| {record["name"]: record for record in scan["results"]} for scan in scans |
| ] |
| names = sorted(set(by_name[0])) |
| if any(set(records) != set(names) for records in by_name[1:]): |
| raise RuntimeError("Scan configuration sets differ") |
| summary: list[dict[str, Any]] = [] |
| for name in names: |
| rows = [records[name] for records in by_name] |
| first = rows[0] |
| summary.append( |
| { |
| "name": name, |
| "head": first["head"], |
| "candidate_steps": first["candidate_steps"], |
| "beta": first["beta"], |
| "threshold": first["threshold"], |
| "target_accepts": first["target_accepts"], |
| "mean_accepted_predictor_calls": sum( |
| float(row["accepted_predictor_calls"]) for row in rows |
| ) |
| / len(rows), |
| "mean_full_calls": sum(float(row["full_calls"]) for row in rows) |
| / len(rows), |
| "mean_policy_latency_ms": sum( |
| float(row["policy_latency_ms"]) for row in rows |
| ) |
| / len(rows), |
| "mean_latent_nrmse_vs_ffff": sum( |
| float(row["latent_nrmse_vs_ffff"]) for row in rows |
| ) |
| / len(rows), |
| "per_prompt_accepts": [ |
| int(row["accepted_predictor_calls"]) for row in rows |
| ], |
| "per_prompt_latent_nrmse": [ |
| float(row["latent_nrmse_vs_ffff"]) for row in rows |
| ], |
| } |
| ) |
| selected: list[dict[str, Any]] = [] |
| for family, targets in TARGETS.items(): |
| for target in targets: |
| candidates = [ |
| row |
| for row in summary |
| if row["head"] == family and int(row["target_accepts"]) == target |
| ] |
| within_budget = [ |
| row |
| for row in candidates |
| if abs(float(row["mean_accepted_predictor_calls"]) - target) <= 0.5 |
| ] |
| pool = within_budget or candidates |
| choice = min( |
| pool, |
| key=lambda row: ( |
| float(row["mean_latent_nrmse_vs_ffff"]), |
| abs(float(row["mean_accepted_predictor_calls"]) - target), |
| float(row["mean_policy_latency_ms"]), |
| float(row["beta"]), |
| ) |
| if within_budget |
| else ( |
| abs(float(row["mean_accepted_predictor_calls"]) - target), |
| float(row["mean_latent_nrmse_vs_ffff"]), |
| float(row["mean_policy_latency_ms"]), |
| float(row["beta"]), |
| ), |
| ) |
| selected.append(choice) |
| atomic_json( |
| args.output_root / "calibration_summary.json", |
| { |
| "status": "complete", |
| "global_indices": [int(scan["global_index"]) for scan in scans], |
| "selection_rule": ( |
| "within mean K +/-0.5 choose lowest mean latent nRMSE; if no " |
| "candidate is within budget, choose nearest K first" |
| ), |
| "summary": summary, |
| "selected": selected, |
| }, |
| ) |
| print("[selected]", flush=True) |
| for row in selected: |
| print( |
| f" {row['head']} K{row['target_accepts']:02d}: beta={row['beta']} " |
| f"threshold={row['threshold']:.8f} " |
| f"actual_K={row['mean_accepted_predictor_calls']:.3f} " |
| f"latent_nrmse={row['mean_latent_nrmse_vs_ffff']:.6f}", |
| flush=True, |
| ) |
|
|
|
|
| def summarize_threshold_only(args: argparse.Namespace) -> None: |
| paths = sorted((args.output_root / "scans").glob("global_*.json")) |
| if len(paths) != 3: |
| raise ValueError(f"Expected exactly three scan files, found {len(paths)}") |
| scans = [json.loads(path.read_text(encoding="utf-8")) for path in paths] |
| if any(scan.get("status") != "complete" for scan in scans): |
| raise RuntimeError("At least one scan is incomplete") |
| by_name = [ |
| {record["name"]: record for record in scan["results"]} for scan in scans |
| ] |
| names = sorted(set(by_name[0])) |
| if any(set(records) != set(names) for records in by_name[1:]): |
| raise RuntimeError("Scan configuration sets differ") |
| summary: list[dict[str, Any]] = [] |
| for name in names: |
| rows = [records[name] for records in by_name] |
| first = rows[0] |
| summary.append( |
| { |
| "name": name, |
| "head": first["head"], |
| "candidate_steps": first["candidate_steps"], |
| "beta": first["beta"], |
| "threshold": first["threshold"], |
| "threshold_rank": first.get("threshold_rank"), |
| "mean_accepted_predictor_calls": sum( |
| float(row["accepted_predictor_calls"]) for row in rows |
| ) |
| / len(rows), |
| "mean_full_calls": sum(float(row["full_calls"]) for row in rows) |
| / len(rows), |
| "mean_policy_latency_ms": sum( |
| float(row["policy_latency_ms"]) for row in rows |
| ) |
| / len(rows), |
| "mean_latent_nrmse_vs_ffff": sum( |
| float(row["latent_nrmse_vs_ffff"]) for row in rows |
| ) |
| / len(rows), |
| "per_prompt_accepts": [ |
| int(row["accepted_predictor_calls"]) for row in rows |
| ], |
| } |
| ) |
| selected: list[dict[str, Any]] = [] |
| for family, targets in TARGETS.items(): |
| for target in targets: |
| candidates = [row for row in summary if row["head"] == family] |
| within_budget = [ |
| row |
| for row in candidates |
| if abs(float(row["mean_accepted_predictor_calls"]) - target) <= 0.5 |
| ] |
| pool = within_budget or candidates |
| choice = min( |
| pool, |
| key=lambda row: ( |
| float(row["mean_latent_nrmse_vs_ffff"]), |
| abs(float(row["mean_accepted_predictor_calls"]) - target), |
| float(row["threshold"]), |
| ) |
| if within_budget |
| else ( |
| abs(float(row["mean_accepted_predictor_calls"]) - target), |
| float(row["mean_latent_nrmse_vs_ffff"]), |
| float(row["threshold"]), |
| ), |
| ) |
| selected.append( |
| { |
| **choice, |
| "target_accepts": target, |
| } |
| ) |
| atomic_json( |
| args.output_root / "threshold_summary.json", |
| { |
| "status": "complete", |
| "beta": 2.0, |
| "global_indices": [int(scan["global_index"]) for scan in scans], |
| "selection_rule": ( |
| "fixed beta=2; for each target K choose threshold with mean K " |
| "within +/-0.5 and lowest latent nRMSE" |
| ), |
| "summary": summary, |
| "selected": selected, |
| }, |
| ) |
| print("[selected-threshold-only]", flush=True) |
| for row in selected: |
| print( |
| f" {row['head']} K{row['target_accepts']:02d}: " |
| f"threshold={row['threshold']:.8f} " |
| f"actual_K={row['mean_accepted_predictor_calls']:.3f} " |
| f"latent_nrmse={row['mean_latent_nrmse_vs_ffff']:.6f}", |
| flush=True, |
| ) |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument( |
| "--mode", |
| choices=( |
| "probe", "prepare", "prepare_threshold", "scan", "summarize", |
| "summarize_threshold", |
| ), |
| required=True, |
| ) |
| parser.add_argument("--gpu", default=PHYSICAL_GPU) |
| parser.add_argument("--global-index", type=int, default=None) |
| parser.add_argument("--mapping", type=Path, default=MAPPING_DEFAULT) |
| parser.add_argument("--experiment-root", type=Path, default=EXPERIMENT_DEFAULT) |
| parser.add_argument("--output-root", type=Path, default=OUTPUT_DEFAULT) |
| parser.add_argument("--probe-root", type=Path, default=None) |
| parser.add_argument("--seed", type=int, default=0) |
| parser.add_argument("--overwrite", action="store_true") |
| args = parser.parse_args() |
| args.mapping = args.mapping.resolve() |
| args.experiment_root = args.experiment_root.resolve() |
| args.output_root = args.output_root.resolve() |
| args.probe_root = args.probe_root.resolve() if args.probe_root is not None else None |
| return args |
|
|
|
|
| def main() -> None: |
| args = parse_args() |
| args.output_root.mkdir(parents=True, exist_ok=True) |
| if args.mode == "probe": |
| probe(args) |
| elif args.mode == "prepare": |
| prepare(args) |
| elif args.mode == "prepare_threshold": |
| prepare_threshold_only(args) |
| elif args.mode == "scan": |
| scan(args) |
| else: |
| if args.mode == "summarize_threshold": |
| summarize_threshold_only(args) |
| else: |
| summarize(args) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
