326 lines
12 KiB
Python
326 lines
12 KiB
Python
"""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()
|