diff --git a/train_all_future.py b/train_all_future.py index 281d69f..9e1068f 100644 --- a/train_all_future.py +++ b/train_all_future.py @@ -32,6 +32,7 @@ from targets import CHECKUP_IDX, PAD_IDX from train_util import ( configure_torch_for_training, create_unique_run_dir, + format_extra_info_types, load_extra_info_types_file, resolve_device, save_checkpoint, @@ -340,7 +341,7 @@ def main() -> None: logger.info(f"Starting all-future training run: {run_name}") logger.info(f"Device: {device}") - logger.info(f"extra_info_types: {args.extra_info_types or 'all'}") + logger.info(f"extra_info_types: {format_extra_info_types(args.extra_info_types)}") logger.info("Loading all-future datasets...") train_dataset = AllFutureHealthDataset( diff --git a/train_next_step.py b/train_next_step.py index 8965bbd..35f86ff 100644 --- a/train_next_step.py +++ b/train_next_step.py @@ -30,6 +30,7 @@ from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX from train_util import ( configure_torch_for_training, create_unique_run_dir, + format_extra_info_types, load_extra_info_types_file, resolve_device, save_checkpoint, @@ -526,7 +527,7 @@ def main() -> None: logger.info(f"Starting next-step training run: {run_name}") logger.info(f"Device: {device}") - logger.info(f"extra_info_types: {args.extra_info_types or 'all'}") + logger.info(f"extra_info_types: {format_extra_info_types(args.extra_info_types)}") logger.info(f"readout={args.readout_name}, target_mode={args.target_mode}") dataset = HealthDataset( diff --git a/train_util.py b/train_util.py index 136e6ab..1fe578c 100644 --- a/train_util.py +++ b/train_util.py @@ -87,6 +87,15 @@ def load_extra_info_types_file(path: str) -> list[int]: raise ValueError(f"Invalid extra info type id in {path}") from exc +def format_extra_info_types(extra_info_types: Iterable[int] | None) -> str: + if extra_info_types is None: + return "all" + values = [int(x) for x in extra_info_types] + if not values: + return "none" + return str(values) + + def load_eid_file(path: str | Path) -> set[int]: file_path = Path(path) if not file_path.is_file():