""" Train DeepHealth with next-token / next-time-point supervision. The next-step dataset uses observed event histories, including CHECKUP state tokens, plus optional gap imputation. UTS training reads out only same-time group ends. """ from __future__ import annotations import argparse import json import logging import math import time from pathlib import Path from typing import Any, Dict import numpy as np import torch from torch.nn.utils import clip_grad_norm_ from torch.optim import AdamW from torch.utils.data import DataLoader, RandomSampler from tqdm.auto import tqdm from dataset import HealthDataset, collate_fn from losses import build_loss from models import DeepHealth, DeepHealthOutput from readouts import build_readout from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX from train_util import ( configure_torch_for_training, create_unique_run_dir, load_extra_info_types_file, resolve_device, save_checkpoint, save_config, set_optimizer_lr, set_seed, setup_logging, split_dataset, ) def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Train DeepHealth with next-token/point supervision") parser.add_argument("--data_prefix", type=str, default="ukb") parser.add_argument("--labels_file", type=str, default="labels.csv") parser.add_argument("--seed", type=int, default=42) parser.add_argument("--extra_info_types_file", type=str, default=None) parser.add_argument("--no_event_interval_years", type=float, default=5.0) parser.add_argument("--include_no_event_in_uts_target", action="store_true") 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("--n_embd", type=int, default=120) parser.add_argument("--n_head", type=int, default=10) parser.add_argument("--n_hist_layer", type=int, default=12) parser.add_argument("--n_tab_layer", type=int, default=4) parser.add_argument("--n_bins", type=int, default=16) parser.add_argument("--extra_pool_reduce", type=str, default="mean", choices=["mean", "sum"]) parser.add_argument("--time_mode", type=str, default="relative", choices=["relative", "absolute"]) parser.add_argument("--dropout", type=float, default=0.0) parser.add_argument("--target_mode", type=str, default="uts", choices=["delphi2m", "uts"]) parser.add_argument("--readout_name", type=str, default=None, choices=["token", "same_time_group_end", "last_valid"]) parser.add_argument("--readout_reduce", type=str, default="mean", choices=["mean", "sum"]) parser.add_argument("--t_min", type=float, default=0.0027378507871321013) parser.add_argument("--max_exp_input", type=float, default=60.0) parser.add_argument("--ce_weight", type=float, default=1.0) parser.add_argument("--time_weight", type=float, default=1.0) parser.add_argument("--ignore_no_event_in_delphi2m", action="store_true") parser.add_argument("--batch_size", type=int, default=128) parser.add_argument("--base_lr", type=float, default=3e-4) parser.add_argument("--weight_decay", type=float, default=0.1) parser.add_argument("--betas", type=float, nargs=2, default=(0.9, 0.99)) parser.add_argument("--grad_clip", type=float, default=1.0) parser.add_argument("--max_epochs", type=int, default=200) parser.add_argument("--warmup_epochs", type=int, default=10) parser.add_argument("--patience", type=int, default=15) parser.add_argument("--min_lr_ratio", type=float, default=0.1) parser.add_argument("--num_workers", type=int, default=4) parser.add_argument("--device", type=str, default="cuda") args = parser.parse_args() if 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.target_mode == "uts": args.readout_name = args.readout_name or "same_time_group_end" args.include_no_event_in_uts_target = True else: args.readout_name = args.readout_name or "token" args.extra_info_types = ( load_extra_info_types_file(args.extra_info_types_file) if args.extra_info_types_file is not None else None ) return args def get_lr(epoch: int, args: argparse.Namespace, adaptive_lr: float) -> float: if epoch < args.warmup_epochs: return adaptive_lr * (epoch + 1) / args.warmup_epochs progress = (epoch - args.warmup_epochs) / max(1, args.max_epochs - args.warmup_epochs) cosine = 0.5 * (1 + math.cos(math.pi * progress)) return adaptive_lr * (args.min_lr_ratio + cosine * (1 - args.min_lr_ratio)) def move_batch_to_device(batch: Dict[str, torch.Tensor], device: torch.device) -> Dict[str, torch.Tensor]: non_blocking = device.type == "cuda" return { key: value.to(device, non_blocking=non_blocking) if isinstance(value, torch.Tensor) else value for key, value in batch.items() } def build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth: return DeepHealth( vocab_size=dataset.vocab_size, n_embd=args.n_embd, n_head=args.n_head, n_hist_layer=args.n_hist_layer, n_tab_layer=args.n_tab_layer, n_types=dataset.n_types, n_cont_types=dataset.n_cont_types, n_categories=dataset.n_categories, cont_type_ids=dataset.cont_type_ids, n_bins=args.n_bins, extra_pool_reduce=args.extra_pool_reduce, target_mode="next_token", time_mode=args.time_mode, dist_mode="exponential", dropout=args.dropout, ) def build_next_step_readout(args: argparse.Namespace): if args.readout_name == "same_time_group_end": return build_readout("same_time_group_end", reduce=args.readout_reduce) return build_readout(args.readout_name) def build_next_step_loss(args: argparse.Namespace): if args.target_mode == "delphi2m": ignored_tokens = {PAD_IDX, CHECKUP_IDX} if args.ignore_no_event_in_delphi2m: ignored_tokens.add(NO_EVENT_IDX) return build_loss( "delphi2m", ignored_tokens=ignored_tokens, t_min=args.t_min, max_exp_input=args.max_exp_input, ce_weight=args.ce_weight, time_weight=args.time_weight, ) return build_loss( "uts", ignored_idx={PAD_IDX, CHECKUP_IDX}, t_min=args.t_min, max_exp_input=args.max_exp_input, ) def build_augmented_next_step_targets( batch: Dict[str, torch.Tensor], model_out: DeepHealthOutput, ) -> Dict[str, torch.Tensor]: hidden_len = model_out.hidden.size(1) event_len = int(model_out.event_len) extra_len = hidden_len - event_len if extra_len <= 0: return { "target_event_seq": batch["target_event_seq"], "target_time_seq": batch["target_time_seq"], "readout_mask": batch["readout_mask"], "target_dt_unique": batch["target_dt_unique"], "target_multi_hot": batch["target_multi_hot"], } device = model_out.hidden.device bsz, _seq_len, vocab_size = batch["target_multi_hot"].shape extra_mask = model_out.padding_mask[:, event_len:] extra_time = model_out.time_seq[:, event_len:] target_event_seq = torch.cat( [ batch["target_event_seq"], torch.full( (bsz, extra_len), PAD_IDX, dtype=batch["target_event_seq"].dtype, device=device, ), ], dim=1, ) target_time_seq = torch.cat( [ batch["target_time_seq"], torch.zeros( bsz, extra_len, dtype=batch["target_time_seq"].dtype, device=device, ), ], dim=1, ) readout_mask = torch.cat([batch["readout_mask"], extra_mask], dim=1) target_dt_unique = torch.cat( [ batch["target_dt_unique"], torch.zeros( bsz, extra_len, dtype=batch["target_dt_unique"].dtype, device=device, ), ], dim=1, ) target_multi_hot = torch.cat( [ batch["target_multi_hot"], torch.zeros( bsz, extra_len, vocab_size, dtype=batch["target_multi_hot"].dtype, device=device, ), ], dim=1, ) for b in range(bsz): valid_event = batch["padding_mask"][b].bool() if not valid_event.any(): continue n_event = int(valid_event.sum().item()) events = torch.cat( [ batch["event_seq"][b, :n_event], batch["target_event_seq"][b, n_event - 1:n_event], ] ) times = torch.cat( [ batch["time_seq"][b, :n_event], batch["target_time_seq"][b, n_event - 1:n_event], ] ) valid_full = events > PAD_IDX events = events[valid_full] times = times[valid_full] if events.numel() == 0: continue for j in range(extra_len): if not bool(extra_mask[b, j]): continue pos = event_len + j t = extra_time[b, j] future = times > t if not future.any(): readout_mask[b, pos] = False continue first_idx = int(torch.nonzero(future, as_tuple=False)[0].item()) next_time = times[first_idx] next_event = events[first_idx] target_event_seq[b, pos] = next_event target_time_seq[b, pos] = next_time same_next_time = times == next_time next_events = events[same_next_time] valid_next_events = next_events[ (next_events > PAD_IDX) & (next_events < vocab_size) ].long() if valid_next_events.numel() == 0: readout_mask[b, pos] = False continue target_multi_hot[b, pos, valid_next_events] = True target_dt_unique[b, pos] = next_time - t return { "target_event_seq": target_event_seq, "target_time_seq": target_time_seq, "readout_mask": readout_mask, "target_dt_unique": target_dt_unique, "target_multi_hot": target_multi_hot, } def compute_next_step_loss( args: argparse.Namespace, model: DeepHealth, readout, criterion, batch: Dict[str, torch.Tensor], device: torch.device, ) -> tuple[torch.Tensor, Dict[str, torch.Tensor]]: batch = move_batch_to_device(batch, device) model_out = model( event_seq=batch["event_seq"], time_seq=batch["time_seq"], sex=batch["sex"], padding_mask=batch["padding_mask"], other_type=batch["other_type"], other_value=batch["other_value"], other_value_kind=batch["other_value_kind"], other_time=batch["other_time"], target_mode="next_token", return_output=True, ) if not isinstance(model_out, DeepHealthOutput): raise TypeError("DeepHealth return_output=True must return DeepHealthOutput") targets = build_augmented_next_step_targets(batch=batch, model_out=model_out) readout_out = readout( hidden=model_out.hidden, time_seq=model_out.time_seq, padding_mask=model_out.padding_mask, readout_mask=targets["readout_mask"] if args.readout_name == "same_time_group_end" else None, ) logits = model.calc_risk(readout_out.hidden) if args.target_mode == "delphi2m": loss, parts = criterion( logits=logits, target_events=targets["target_event_seq"], target_times=targets["target_time_seq"], current_times=model_out.time_seq, padding_mask=readout_out.readout_mask, return_components=True, ) else: loss, parts = criterion( logits=logits, target_multi_hot=targets["target_multi_hot"], target_dt_unique=targets["target_dt_unique"], readout_mask=readout_out.readout_mask, return_components=True, ) if not torch.isfinite(loss): raise RuntimeError(f"Loss is not finite: {float(loss.detach().cpu())}") return loss, parts def run_epoch( logger: logging.Logger, args: argparse.Namespace, model: DeepHealth, readout, criterion, loader: DataLoader, optimizer: AdamW | None, device: torch.device, is_train: bool, ) -> float: model.train(is_train) readout.train(is_train) total = 0.0 n_batches = 0 skipped = 0 parts_sum: Dict[str, float] = {} desc = "train" if is_train else "val" progress = tqdm(loader, desc=desc, leave=False, dynamic_ncols=True) for batch_idx, batch in enumerate(progress): try: loss, parts = compute_next_step_loss(args, model, readout, criterion, batch, device) if is_train: if optimizer is None: raise ValueError("optimizer is required for training") optimizer.zero_grad(set_to_none=True) loss.backward() if args.grad_clip > 0: clip_grad_norm_(model.parameters(), args.grad_clip) optimizer.step() total += float(loss.detach().cpu()) n_batches += 1 for name, value in parts.items(): parts_sum[name] = parts_sum.get(name, 0.0) + float(value.detach().cpu()) avg = total / max(1, n_batches) postfix = {"loss": f"{float(loss.detach().cpu()):.4f}", "avg": f"{avg:.4f}", "skipped": skipped} for name, value in parts_sum.items(): postfix[name] = f"{value / max(1, n_batches):.4f}" progress.set_postfix(postfix) except RuntimeError as exc: if "Loss is not finite" not in str(exc): raise skipped += 1 logger.warning(f"Batch {batch_idx} skipped: {str(exc)[:120]}") if skipped: logger.info(f"Skipped {skipped} batches due to non-finite loss") return total / max(1, n_batches) if n_batches else float("inf") def build_metadata( args: argparse.Namespace, dataset: HealthDataset, run_name: str, train_subset, val_subset, test_subset, ) -> Dict[str, Any]: return { "run_name": run_name, "dataset_class": "NextStepHealthDataset", "collate_fn": "next_step_collate_fn", "model_class": "DeepHealth", "model_target_mode": "next_token", "target_mode": args.target_mode, "dist_mode": "exponential", "extra_info_types_file": ( Path(args.extra_info_types_file).name if args.extra_info_types_file is not None else None ), "extra_info_types": [int(x) for x in dataset.extra_info_types], "dataset_metadata": { "vocab_size": int(dataset.vocab_size), "n_types": int(dataset.n_types), "n_cont_types": int(dataset.n_cont_types), "n_categories": int(dataset.n_categories), "cont_type_ids": [int(x) for x in dataset.cont_type_ids], "extra_info_types": [int(x) for x in dataset.extra_info_types], }, "split_sizes": { "train": int(len(train_subset)), "val": int(len(val_subset)), "test": int(len(test_subset)), }, "resolved_readout_name": args.readout_name, "resolved_loss_name": args.target_mode, } def main() -> None: args = parse_args() set_seed(args.seed) device = resolve_device(args.device) configure_torch_for_training(device) run_dir, run_name = create_unique_run_dir( lambda timestamp: ( f"{args.time_mode}_exponential_next_token_{args.target_mode}_" f"gap_{args.no_event_interval_years:g}y_{timestamp}" ) ) logger = setup_logging(run_dir) 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"readout={args.readout_name}, target_mode={args.target_mode}") dataset = HealthDataset( data_prefix=args.data_prefix, labels_file=args.labels_file, no_event_interval_years=args.no_event_interval_years, include_no_event_in_uts_target=args.include_no_event_in_uts_target, extra_info_types=args.extra_info_types, ) train_subset, val_subset, test_subset = split_dataset( dataset=dataset, train_ratio=args.train_ratio, val_ratio=args.val_ratio, test_ratio=args.test_ratio, seed=args.seed, ) logger.info( f"Samples: train={len(train_subset)}, val={len(val_subset)}, test={len(test_subset)}" ) train_loader = DataLoader( train_subset, batch_size=args.batch_size, sampler=RandomSampler(train_subset, generator=torch.Generator().manual_seed(args.seed)), collate_fn=collate_fn, num_workers=args.num_workers, pin_memory=device.type == "cuda", persistent_workers=args.num_workers > 0, prefetch_factor=2 if args.num_workers > 0 else None, ) val_loader = DataLoader( val_subset, batch_size=args.batch_size, shuffle=False, collate_fn=collate_fn, num_workers=args.num_workers, pin_memory=device.type == "cuda", persistent_workers=args.num_workers > 0, prefetch_factor=2 if args.num_workers > 0 else None, ) test_loader = DataLoader( test_subset, batch_size=args.batch_size, shuffle=False, collate_fn=collate_fn, num_workers=args.num_workers, pin_memory=device.type == "cuda", persistent_workers=args.num_workers > 0, prefetch_factor=2 if args.num_workers > 0 else None, ) model = build_model(args, dataset).to(device) readout = build_next_step_readout(args).to(device) criterion = build_next_step_loss(args) optimizer = AdamW( model.parameters(), lr=args.base_lr, betas=tuple(args.betas), weight_decay=args.weight_decay, ) adaptive_lr = args.base_lr * math.sqrt(args.batch_size / 128) save_config( args, run_dir / "train_config.json", extra=build_metadata(args, dataset, run_name, train_subset, val_subset, test_subset), ) best_val = float("inf") patience = 0 history = [] best_model_path = run_dir / "best_model.pt" start = time.time() for epoch in range(args.max_epochs): lr = get_lr(epoch, args, adaptive_lr) set_optimizer_lr(optimizer, lr) train_loss = run_epoch(logger, args, model, readout, criterion, train_loader, optimizer, device, True) with torch.no_grad(): val_loss = run_epoch(logger, args, model, readout, criterion, val_loader, None, device, False) is_best = val_loss < best_val if is_best: best_val = val_loss patience = 0 save_checkpoint(model, best_model_path) else: patience += 1 logger.info( f"Epoch {epoch + 1}/{args.max_epochs} | lr={lr:.6f} | " f"train_loss={train_loss:.6f} | val_loss={val_loss:.6f} | " f"best_val_loss={best_val:.6f} | patience={patience}/{args.patience} | " f"elapsed={time.time() - start:.1f}s" ) history.append({ "epoch": epoch + 1, "lr": lr, "train_loss": train_loss, "val_loss": val_loss, "best_val_loss": best_val, "is_best": int(is_best), }) if patience >= args.patience: logger.info(f"Early stopping triggered at epoch {epoch + 1}") break with (run_dir / "history.json").open("w", encoding="utf-8") as f: json.dump(history, f, indent=2) logger.info("Evaluating best model on next-step test split...") model.load_state_dict(torch.load(best_model_path, map_location=device)) with torch.no_grad(): test_loss = run_epoch(logger, args, model, readout, criterion, test_loader, None, device, False) logger.info(f"Test loss: {test_loss:.6f}") logger.info(f"Best checkpoint: {best_model_path}") if __name__ == "__main__": main()