Refactor code structure for improved readability and maintainability
This commit is contained in:
@@ -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)}"
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user