DeepHealthExpo/train_next_step.py

"""
Train DeepHealth with next-token / next-time-point supervision.

The next-step dataset uses observed event histories, including CHECKUP state
tokens, plus optional gap <NO_EVENT> 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 typing import Any, Dict, Iterator, List

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, Sampler
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,
    resolve_device,
    save_checkpoint,
    save_config,
    set_optimizer_lr,
    set_seed,
    setup_logging,
    split_dataset,
    split_dataset_by_eid_files,
)


MODEL_INPUT_KEYS = (
    "event_seq",
    "time_seq",
    "sex",
    "padding_mask",
)

EXPOSURE_INPUT_KEYS = (
    "exposure_daily",
    "exposure_monthly",
)


class ExposureLocalityBatchSampler(Sampler[List[int]]):
    """Randomized batches with within-buffer sorting by exposure parquet locality."""

    def __init__(
        self,
        data_source,
        batch_size: int,
        buffer_size: int,
        seed: int,
        drop_last: bool = False,
    ) -> None:
        self.data_source = data_source
        self.batch_size = int(batch_size)
        self.buffer_size = max(int(buffer_size), self.batch_size)
        self.seed = int(seed)
        self.drop_last = bool(drop_last)
        self.epoch = 0

    def __iter__(self) -> Iterator[List[int]]:
        n = len(self.data_source)
        generator = torch.Generator().manual_seed(self.seed + self.epoch)
        self.epoch += 1
        shuffled = torch.randperm(n, generator=generator).tolist()
        for start in range(0, n, self.buffer_size):
            buffer = shuffled[start:start + self.buffer_size]
            buffer.sort(key=self._locality_key)
            for batch_start in range(0, len(buffer), self.batch_size):
                batch = buffer[batch_start:batch_start + self.batch_size]
                if len(batch) < self.batch_size and self.drop_last:
                    continue
                yield batch

    def __len__(self) -> int:
        n = len(self.data_source)
        if self.drop_last:
            return n // self.batch_size
        return (n + self.batch_size - 1) // self.batch_size

    def _locality_key(self, local_idx: int) -> tuple[int, int, int]:
        dataset = self.data_source
        raw_idx = int(local_idx)
        if hasattr(dataset, "dataset") and hasattr(dataset, "indices"):
            raw_idx = int(dataset.indices[local_idx])
            dataset = dataset.dataset
        sample = getattr(dataset, "samples", [])[raw_idx]
        exposure_index = sample.get("exposure_index")
        exposure_cache = getattr(dataset, "exposure_cache", None)
        if exposure_index is None or exposure_cache is None:
            return (2**31 - 1, 2**31 - 1, raw_idx)
        block_id, block_offset = exposure_cache.locality_key(exposure_index)
        return (block_id, block_offset, raw_idx)


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("--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("--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("--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("--dropout", type=float, default=0.0)
    parser.add_argument("--exposure_cache_dir", type=str, default=None)
    parser.add_argument("--mask_onset_exposure", action="store_true")
    parser.add_argument("--exposure_d_model", type=int, default=None)
    parser.add_argument("--exposure_n_layers", type=int, default=2)
    parser.add_argument("--exposure_top_k", type=int, default=3)
    parser.add_argument("--exposure_n_convnext_blocks", type=int, default=2)
    parser.add_argument("--exposure_conv_kernel_size", type=int, default=7)
    parser.add_argument("--exposure_mlp_ratio", type=float, default=4.0)
    parser.add_argument("--no_exposure_gate", action="store_true")
    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(
        "--prefetch_factor",
        type=int,
        default=4,
        help="DataLoader batches prefetched per worker when num_workers > 0.",
    )
    parser.add_argument(
        "--exposure_locality_buffer_size",
        type=int,
        default=4096,
        help=(
            "Training-only shuffle buffer sorted by exposure parquet locality. "
            "Set 0 to use the standard RandomSampler."
        ),
    )
    parser.add_argument("--device", type=str, default="cuda")
    parser.add_argument("--progress_interval", type=int, default=20)

    args = parser.parse_args()
    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.num_workers > 0 and args.prefetch_factor <= 0:
        raise ValueError("prefetch_factor must be positive when num_workers > 0")
    if args.exposure_locality_buffer_size < 0:
        raise ValueError("exposure_locality_buffer_size must be non-negative")
    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"
    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,
        target_mode="next_token",
        dist_mode="exponential",
        dropout=args.dropout,
        use_exposure_encoder=args.exposure_cache_dir is not None,
        exposure_d_model=args.exposure_d_model,
        exposure_n_layers=args.exposure_n_layers,
        exposure_top_k=args.exposure_top_k,
        exposure_n_convnext_blocks=args.exposure_n_convnext_blocks,
        exposure_conv_kernel_size=args.exposure_conv_kernel_size,
        exposure_mlp_ratio=args.exposure_mlp_ratio,
        exposure_use_gate=not args.no_exposure_gate,
    )


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_cpu: Dict[str, torch.Tensor],
    model_out: DeepHealthOutput,
    include_uts_targets: bool,
) -> Dict[str, torch.Tensor]:
    device = model_out.hidden.device
    non_blocking = device.type == "cuda"
    targets = {
        "target_event_seq": batch_cpu["target_event_seq"].to(device, non_blocking=non_blocking),
        "target_time_seq": batch_cpu["target_time_seq"].to(device, non_blocking=non_blocking),
        "readout_mask": batch_cpu["readout_mask"].to(device, non_blocking=non_blocking),
    }
    if include_uts_targets:
        targets["target_dt_unique"] = batch_cpu["target_dt_unique"].to(
            device, non_blocking=non_blocking
        )
        targets["target_multi_hot"] = batch_cpu["target_multi_hot"].to(
            device, non_blocking=non_blocking
        )
    return targets


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_cpu = batch
    input_keys = list(MODEL_INPUT_KEYS)
    input_keys.extend(key for key in EXPOSURE_INPUT_KEYS if key in batch_cpu)
    batch = move_batch_to_device(
        {key: batch_cpu[key] for key in input_keys},
        device,
    )
    model_kwargs = {
        "event_seq": batch["event_seq"],
        "time_seq": batch["time_seq"],
        "sex": batch["sex"],
        "padding_mask": batch["padding_mask"],
        "target_mode": "next_token",
        "return_output": True,
    }
    if "exposure_daily" in batch:
        model_kwargs["exposure_daily"] = batch["exposure_daily"]
        model_kwargs["exposure_monthly"] = batch["exposure_monthly"]
    model_out = model(**model_kwargs)
    if not isinstance(model_out, DeepHealthOutput):
        raise TypeError("DeepHealth return_output=True must return DeepHealthOutput")
    targets = build_augmented_next_step_targets(
        batch_cpu=batch_cpu,
        model_out=model_out,
        include_uts_targets=args.target_mode == "uts",
    )
    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 = torch.zeros((), device=device)
    n_batches = 0
    skipped = 0
    parts_sum: Dict[str, torch.Tensor] = {}
    desc = "train" if is_train else "val"
    progress_interval = max(1, int(args.progress_interval))

    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 = total + loss.detach()
            n_batches += 1
            for name, value in parts.items():
                parts_sum[name] = parts_sum.get(name, torch.zeros((), device=device)) + value.detach()
            if (batch_idx + 1) % progress_interval == 0:
                avg = total / max(1, n_batches)
                postfix = {
                    "loss": f"{float(loss.detach().cpu()):.4f}",
                    "avg": f"{float(avg.detach().cpu()):.4f}",
                    "skipped": skipped,
                }
                for name, value in parts_sum.items():
                    postfix[name] = f"{float((value / max(1, n_batches)).detach().cpu()):.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 float((total / max(1, n_batches)).detach().cpu()) 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",
        "dataset_metadata": {
            "vocab_size": int(dataset.vocab_size),
        },
        "use_exposure_encoder": args.exposure_cache_dir is not None,
        "exposure_cache_dir": args.exposure_cache_dir,
        "mask_onset_exposure": bool(args.mask_onset_exposure),
        "exposure_d_model": args.exposure_d_model,
        "exposure_n_layers": int(args.exposure_n_layers),
        "exposure_top_k": int(args.exposure_top_k),
        "exposure_n_convnext_blocks": int(args.exposure_n_convnext_blocks),
        "exposure_conv_kernel_size": int(args.exposure_conv_kernel_size),
        "exposure_mlp_ratio": float(args.exposure_mlp_ratio),
        "exposure_use_gate": not bool(args.no_exposure_gate),
        "num_workers": int(args.num_workers),
        "prefetch_factor": int(args.prefetch_factor),
        "exposure_locality_buffer_size": int(args.exposure_locality_buffer_size),
        "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"absolute_exponential_next_token_{args.target_mode}_"
            f"gap_{args.no_event_interval_years:g}y_"
            f"{'exposure' if args.exposure_cache_dir else 'noexposure'}_"
            f"{timestamp}"
        )
    )
    logger = setup_logging(run_dir)

    logger.info(f"Starting next-step training run: {run_name}")
    logger.info(f"Device: {device}")
    logger.info(f"readout={args.readout_name}, target_mode={args.target_mode}")
    logger.info(f"exposure_cache_dir={args.exposure_cache_dir}")
    logger.info(
        "DataLoader IO: "
        f"num_workers={args.num_workers}, "
        f"prefetch_factor={args.prefetch_factor if args.num_workers > 0 else None}, "
        f"exposure_locality_buffer_size={args.exposure_locality_buffer_size}"
    )

    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,
        exposure_cache_dir=args.exposure_cache_dir,
        mask_onset_exposure=args.mask_onset_exposure,
    )
    if args.train_eid_file and args.val_eid_file and args.test_eid_file:
        train_subset, val_subset, test_subset = split_dataset_by_eid_files(
            dataset=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_dataset(
            dataset=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"Samples: train={len(train_subset)}, val={len(val_subset)}, test={len(test_subset)}"
    )

    dataloader_kwargs = {
        "collate_fn": collate_fn,
        "num_workers": args.num_workers,
        "pin_memory": device.type == "cuda",
        "persistent_workers": args.num_workers > 0,
        "prefetch_factor": args.prefetch_factor if args.num_workers > 0 else None,
    }
    use_locality_sampler = (
        args.exposure_cache_dir is not None
        and args.exposure_locality_buffer_size > 0
    )
    if use_locality_sampler:
        logger.info("Using exposure-locality batch sampler for training")
        train_loader = DataLoader(
            train_subset,
            batch_sampler=ExposureLocalityBatchSampler(
                train_subset,
                batch_size=args.batch_size,
                buffer_size=args.exposure_locality_buffer_size,
                seed=args.seed,
            ),
            **dataloader_kwargs,
        )
    else:
        train_loader = DataLoader(
            train_subset,
            batch_size=args.batch_size,
            sampler=RandomSampler(
                train_subset,
                generator=torch.Generator().manual_seed(args.seed),
            ),
            **dataloader_kwargs,
        )
    val_loader = DataLoader(
        val_subset,
        batch_size=args.batch_size,
        shuffle=False,
        **dataloader_kwargs,
    )
    test_loader = DataLoader(
        test_subset,
        batch_size=args.batch_size,
        shuffle=False,
        **dataloader_kwargs,
    )

    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()