Refactor code structure for improved readability and maintainability

This commit is contained in:
2026-06-18 13:07:35 +08:00
parent 1ea72e9133
commit aa8ec5c3ac
7 changed files with 1005002 additions and 18 deletions

View File

@@ -40,6 +40,7 @@ from train_util import (
set_seed,
setup_logging,
split_all_future_datasets,
split_all_future_datasets_by_eid_files,
)
@@ -55,6 +56,9 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--train_ratio", type=float, default=0.7)
parser.add_argument("--val_ratio", type=float, default=0.15)
parser.add_argument("--test_ratio", type=float, default=0.15)
parser.add_argument("--train_eid_file", type=str, default="ukb_train_eid.csv")
parser.add_argument("--val_eid_file", type=str, default="ukb_val_eid.csv")
parser.add_argument("--test_eid_file", type=str, default="ukb_test_eid.csv")
parser.add_argument("--min_history_events", type=int, default=1)
parser.add_argument("--min_future_events", type=int, default=1)
parser.add_argument("--validation_query_seed", type=int, default=None)
@@ -89,7 +93,11 @@ def parse_args() -> argparse.Namespace:
raise ValueError("min_history_events must be >= 1")
if args.min_future_events < 1:
raise ValueError("min_future_events must be >= 1")
if not np.isclose(args.train_ratio + args.val_ratio + args.test_ratio, 1.0):
use_eid_split = all(
getattr(args, name)
for name in ("train_eid_file", "val_eid_file", "test_eid_file")
)
if not use_eid_split and not np.isclose(args.train_ratio + args.val_ratio + args.test_ratio, 1.0):
raise ValueError("train_ratio + val_ratio + test_ratio must equal 1.0")
if args.validation_query_seed is None:
args.validation_query_seed = int(args.seed)
@@ -335,15 +343,33 @@ def main() -> None:
validation_query_seed=args.validation_query_seed,
extra_info_types=args.extra_info_types,
)
train_subset, val_subset, test_subset = split_all_future_datasets(
train_dataset=train_dataset,
val_dataset=val_dataset,
test_dataset=test_dataset,
train_ratio=args.train_ratio,
val_ratio=args.val_ratio,
test_ratio=args.test_ratio,
seed=args.seed,
)
if args.train_eid_file and args.val_eid_file and args.test_eid_file:
train_subset, val_subset, test_subset = split_all_future_datasets_by_eid_files(
train_dataset=train_dataset,
val_dataset=val_dataset,
test_dataset=test_dataset,
train_eid_file=args.train_eid_file,
val_eid_file=args.val_eid_file,
test_eid_file=args.test_eid_file,
)
logger.info(
"Using eid split files: "
f"train={args.train_eid_file}, val={args.val_eid_file}, test={args.test_eid_file}"
)
else:
train_subset, val_subset, test_subset = split_all_future_datasets(
train_dataset=train_dataset,
val_dataset=val_dataset,
test_dataset=test_dataset,
train_ratio=args.train_ratio,
val_ratio=args.val_ratio,
test_ratio=args.test_ratio,
seed=args.seed,
)
logger.info(
f"Using random ratio split: train={args.train_ratio}, "
f"val={args.val_ratio}, test={args.test_ratio}, seed={args.seed}"
)
logger.info(
f"Patients/queries: train={len(train_subset)}, val={len(val_subset)}, test={len(test_subset)}"
)