From f6bde7e16790749103778341bec3c83445994714 Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Wed, 8 Jul 2026 17:29:25 +0800 Subject: [PATCH] Add t-query logits and hidden export script --- export_tquery_logits_hidden.py | 325 +++++++++++++++++++++++++++++++++ 1 file changed, 325 insertions(+) create mode 100644 export_tquery_logits_hidden.py diff --git a/export_tquery_logits_hidden.py b/export_tquery_logits_hidden.py new file mode 100644 index 0000000..692e805 --- /dev/null +++ b/export_tquery_logits_hidden.py @@ -0,0 +1,325 @@ +"""Export landmark risk logits and hidden states for t_query ages. + +This script follows evaluate_event_free_survival.py's data loading, +landmark construction, checkpoint loading, and readout logic, but only exports: + +* all token/disease risk logits from ``model.calc_risk(hidden)``; +* the corresponding landmark hidden state. + +The two large arrays are saved separately as .npy files. Row metadata is saved +as a CSV with matching row order. +""" +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any, Optional + +import numpy as np +import pandas as pd +import torch +from torch.utils.data import DataLoader +from tqdm.auto import tqdm + +from evaluate_auc_v2 import ( + LandmarkDataset, + build_model_from_dataset, + cfg_get, + load_checkpoint_state_dict, + load_json_config, + load_model_state, + resolve_dist_mode_for_checkpoint, + resolve_eval_device, + validate_dataset_metadata, +) +from evaluate_event_free_survival import ( + IndexedLandmarkDataset, + build_first_occurrence_maps_for_landmarks, + collate_indexed_landmark_fn, + infer_landmark_hidden, + load_eval_sequence_dataset, + make_landmark_ages, +) + + +def numpy_float_dtype(name: str) -> np.dtype: + key = str(name).lower() + if key in {"float16", "fp16", "half"}: + return np.dtype(np.float16) + if key in {"float32", "fp32", "single"}: + return np.dtype(np.float32) + raise ValueError(f"dtype must be float16 or float32, got {name!r}") + + +def output_paths_for_run( + run_path: Path, + eval_split: str, + landmark_start: float, + landmark_stop: float, + landmark_step: float, + output_dir: Optional[str], +) -> tuple[Path, Path, Path, Path]: + suffix = f"{eval_split}_t{landmark_start:g}-{landmark_stop:g}_step{landmark_step:g}" + base_dir = Path(output_dir) if output_dir else run_path + return ( + base_dir / f"tquery_logits_{suffix}.npy", + base_dir / f"tquery_hidden_{suffix}.npy", + base_dir / f"tquery_metadata_{suffix}.csv", + base_dir / f"tquery_export_config_{suffix}.json", + ) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Export landmark risk logits and hidden states for t_query ages." + ) + parser.add_argument("--run_path", type=str, required=True) + parser.add_argument( + "--output_dir", + type=str, + default=None, + help="Directory for output files. Defaults to run_path.", + ) + parser.add_argument("--logits_path", type=str, default=None) + parser.add_argument("--hidden_path", type=str, default=None) + parser.add_argument("--metadata_path", type=str, default=None) + parser.add_argument("--export_config_path", type=str, default=None) + parser.add_argument("--eval_split", type=str, default=None) + parser.add_argument("--dataset_subset_size", type=int, default=None) + parser.add_argument("--train_eid_file", type=str, default=None) + parser.add_argument("--val_eid_file", type=str, default=None) + parser.add_argument("--test_eid_file", type=str, default=None) + parser.add_argument("--landmark_start", type=float, default=40.0) + parser.add_argument("--landmark_stop", type=float, default=80.0) + parser.add_argument( + "--landmark_step", + type=float, + default=1.0, + help="t_query grid step in years. Default exports every integer age 40..80.", + ) + parser.add_argument("--min_history_events", type=int, default=None) + parser.add_argument("--batch_size", type=int, default=None) + parser.add_argument("--num_workers", type=int, default=None) + parser.add_argument("--device", type=str, default=None) + parser.add_argument("--extra_info_types", type=str, default=None) + parser.add_argument( + "--logits_dtype", + type=str, + default="float32", + choices=["float16", "float32"], + ) + parser.add_argument( + "--hidden_dtype", + type=str, + default="float32", + choices=["float16", "float32"], + ) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + run_path = Path(args.run_path) + config_path = run_path / "train_config.json" + checkpoint_path = run_path / "best_model.pt" + if not config_path.exists(): + raise FileNotFoundError(f"train_config.json not found: {config_path}") + if not checkpoint_path.exists(): + raise FileNotFoundError(f"best_model.pt not found: {checkpoint_path}") + + cfg = load_json_config(config_path) + model_target_mode = str(cfg.get("model_target_mode", "next_token")).lower() + if model_target_mode not in {"next_token", "all_future"}: + raise ValueError(f"Unsupported model_target_mode: {model_target_mode!r}") + + target_mode = str(cfg.get("target_mode", "uts")) + attn_mask_mode = str( + cfg.get( + "attn_mask_mode", + "non_strict_time" if target_mode == "uts" else "target_aware", + ) + ) + readout_name = str( + cfg.get("readout_name", "same_time_group_end" if target_mode == "uts" else "token") + ) + readout_reduce = str(cfg.get("readout_reduce", "mean")) + + dataset, subset_indices, eval_split, split_source = load_eval_sequence_dataset( + args, + cfg, + ) + validate_dataset_metadata(dataset, cfg) + + landmark_ages = make_landmark_ages( + float(args.landmark_start), + float(args.landmark_stop), + float(args.landmark_step), + ) + first_occurrence_by_token = build_first_occurrence_maps_for_landmarks( + dataset, + subset_indices, + ) + death_idx = int(dataset.vocab_size) - 1 + landmark_dataset = LandmarkDataset( + dataset=dataset, + subset_indices=subset_indices, + landmark_ages=landmark_ages, + attn_mask_mode=attn_mask_mode, + model_target_mode=model_target_mode, + min_history_events=int(cfg_get(args, cfg, "min_history_events", 1)), + first_occurrence_by_token=first_occurrence_by_token, + death_token_ids=[death_idx], + ) + + state_dict = load_checkpoint_state_dict(checkpoint_path, map_location="cpu") + dist_mode = resolve_dist_mode_for_checkpoint( + str(cfg.get("dist_mode", "exponential")), + state_dict, + ) + cfg_model = dict(cfg) + cfg_model["dist_mode"] = dist_mode + device = resolve_eval_device(args.device) + model = build_model_from_dataset(args, cfg_model, dataset).to(device) + load_model_state(model, state_dict) + model.eval() + + default_logits_path, default_hidden_path, default_metadata_path, default_config_path = ( + output_paths_for_run( + run_path=run_path, + eval_split=eval_split, + landmark_start=float(args.landmark_start), + landmark_stop=float(args.landmark_stop), + landmark_step=float(args.landmark_step), + output_dir=args.output_dir, + ) + ) + logits_path = Path(args.logits_path) if args.logits_path else default_logits_path + hidden_path = Path(args.hidden_path) if args.hidden_path else default_hidden_path + metadata_path = Path(args.metadata_path) if args.metadata_path else default_metadata_path + export_config_path = ( + Path(args.export_config_path) if args.export_config_path else default_config_path + ) + for path in (logits_path, hidden_path, metadata_path, export_config_path): + path.parent.mkdir(parents=True, exist_ok=True) + + n_rows = len(landmark_dataset) + vocab_size = int(dataset.vocab_size) + hidden_dim = int(getattr(model, "n_embd", cfg_get(args, cfg_model, "n_embd", 120))) + logits_dtype = numpy_float_dtype(args.logits_dtype) + hidden_dtype = numpy_float_dtype(args.hidden_dtype) + + logits_memmap = np.lib.format.open_memmap( + logits_path, + mode="w+", + dtype=logits_dtype, + shape=(n_rows, vocab_size), + ) + hidden_memmap = np.lib.format.open_memmap( + hidden_path, + mode="w+", + dtype=hidden_dtype, + shape=(n_rows, hidden_dim), + ) + + batch_size = int(cfg_get(args, cfg, "batch_size", 128)) + num_workers = int(cfg_get(args, cfg, "num_workers", 4)) + loader = DataLoader( + IndexedLandmarkDataset(landmark_dataset), + batch_size=batch_size, + shuffle=False, + collate_fn=collate_indexed_landmark_fn, + num_workers=num_workers, + pin_memory=device.type == "cuda", + persistent_workers=num_workers > 0, + prefetch_factor=2 if num_workers > 0 else None, + ) + + print(f"Eval split: {eval_split}") + print(f"Split source: {split_source}") + print(f"Selected patients: {len(subset_indices)}") + print(f"t_query ages: {landmark_ages.tolist()}") + print(f"Dist mode: {dist_mode}") + print(f"Device: {device}") + print(f"Landmark rows: {n_rows}") + print(f"Logits: {logits_path} shape={(n_rows, vocab_size)} dtype={logits_dtype}") + print(f"Hidden: {hidden_path} shape={(n_rows, hidden_dim)} dtype={hidden_dtype}") + print(f"Metadata: {metadata_path}") + + meta_rows: list[dict[str, Any]] = [] + written = 0 + with torch.no_grad(): + for batch in tqdm(loader, desc="Export logits/hidden", dynamic_ncols=True): + hidden = infer_landmark_hidden( + model=model, + batch=batch, + device=device, + model_target_mode=model_target_mode, + readout_name=readout_name, + readout_reduce=readout_reduce, + ) + logits = model.calc_risk(hidden) + row_indices = batch["row_idx"].detach().cpu().numpy().astype(np.int64) + if not np.array_equal(row_indices, np.arange(written, written + len(row_indices))): + raise RuntimeError("DataLoader row order changed; export requires shuffle=False.") + + batch_n = int(logits.shape[0]) + logits_memmap[written : written + batch_n] = ( + logits.detach().cpu().numpy().astype(logits_dtype, copy=False) + ) + hidden_memmap[written : written + batch_n] = ( + hidden.detach().cpu().numpy().astype(hidden_dtype, copy=False) + ) + + for row_idx in row_indices.tolist(): + meta = landmark_dataset.rows[int(row_idx)] + sample = dataset.samples[int(meta["dataset_index"])] + meta_rows.append( + { + "row_index": int(row_idx), + "patient_id": int(meta["patient_id"]), + "dataset_index": int(meta["dataset_index"]), + "eid": int(sample.get("eid", -1)), + "sex": int(meta["sex"]), + "t_query": float(meta["t_query"]), + "landmark_age": float(meta["landmark_age"]), + "followup_end_time": float(meta["followup_end_time"]), + "death_time": float(meta["death_time"]), + } + ) + written += batch_n + + logits_memmap.flush() + hidden_memmap.flush() + pd.DataFrame(meta_rows).to_csv(metadata_path, index=False) + + export_config = { + "run_path": str(run_path), + "eval_split": eval_split, + "split_source": split_source, + "model_target_mode": model_target_mode, + "target_mode": target_mode, + "attn_mask_mode": attn_mask_mode, + "readout_name": readout_name, + "readout_reduce": readout_reduce, + "dist_mode": dist_mode, + "landmark_ages": [float(x) for x in landmark_ages.tolist()], + "n_rows": int(n_rows), + "vocab_size": int(vocab_size), + "hidden_dim": int(hidden_dim), + "death_token": int(death_idx), + "logits_path": str(logits_path), + "hidden_path": str(hidden_path), + "metadata_path": str(metadata_path), + "logits_dtype": str(logits_dtype), + "hidden_dtype": str(hidden_dtype), + } + with export_config_path.open("w", encoding="utf-8") as f: + json.dump(export_config, f, indent=2) + + print(f"Wrote {written} rows.") + print(f"Wrote export config: {export_config_path}") + + +if __name__ == "__main__": + main()