Add training scripts for all-future and next-step supervision with DeepHealth

- Implement `train_all_future.py` for training with query-conditioned all-future supervision.
- Implement `train_next_step.py` for training with next-token/next-time-point supervision.
- Introduce `train_util.py` for shared utility functions including logging, dataset splitting, and model checkpointing.
- Enhance argument parsing for both training scripts to accommodate new parameters.
- Update loss functions and model configurations to support the new training paradigms.
This commit is contained in:
2026-06-13 11:42:04 +08:00
parent 034d8065a7
commit 46a3dfe628
12 changed files with 1927 additions and 1273 deletions

View File

@@ -28,6 +28,7 @@ from torch.utils.data import DataLoader, Dataset
from tqdm.auto import tqdm
from dataset import HealthDataset
from eval_data import load_sequence_eval_dataset
from models import DeepHealth
from readouts import build_readout
from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX
@@ -1391,11 +1392,14 @@ def main() -> None:
labels_meta = pd.read_csv(str(labels_meta_path))
print("Loading dataset...")
dataset = HealthDataset(
dataset = load_sequence_eval_dataset(
model_target_mode=model_target_mode,
data_prefix=data_prefix,
labels_file=labels_file,
no_event_interval_years=float(no_event_interval_years),
include_no_event_in_uts_target=bool(include_no_event_in_uts_target),
min_history_events=int(cfg.get("all_future_min_history_events", 1)),
min_future_events=int(cfg.get("all_future_min_future_events", 1)),
extra_info_types=parse_int_list(cfg.get("extra_info_types", None)),
)
validate_dataset_metadata(dataset, cfg)