DeepHealthNew

这个目录包含 DeepHealth 的数据准备、数据集、模型、readout 和 loss 代码。当前版本的核心设计是：

疾病序列 stream + 统一的额外信息 token stream

疾病、死亡、checkup 事件仍然保存在事件序列里；性别单独保存在 basic_info；其他体检、暴露、生活方式等信息统一整理成 (type, value, value_kind, time) token。

数据准备

运行：

python prepare_data.py

输入文件：

• ukb_data.csv
• field_ids_enriched.csv
• icd10_codes_mod.tsv
• labels.csv

输出文件：

• ukb_event_data.npy

  • 形状为 (N, 3)
  • 每行是 (eid, days, label)
  • 包含疾病、死亡、checkup 事件

• ukb_basic_info.csv

  • index 为 eid
  • 当前只保留 sex

• ukb_other_info.npy

  • 形状为 (M, 5)
  • 每行是：

    (eid, type, value, value_kind, time)
    

  • type=0 预留给 padding
  • value_kind=1 表示连续变量
  • value_kind=2 表示分类变量
  • 缺失值不会生成 token
  • 当前 UKB 额外信息只有一个时间点，所以 time 暂时都是 date_of_assessment

• cate_types.csv

  • 分类变量元信息
  • 字段：

    type,name,n_categories
    

  • ukb_other_info.npy 里的分类 value 是变量内部的局部 id；global category id 在 dataset 中根据实验选择的变量动态计算。

Dataset

dataset.py 提供两个 dataset：

• NextStepHealthDataset

  • 用于 next-token / next-time-point 监督
  • 对应 Delphi2MLoss 和 UniqueTimeSetExponentialLoss

• AllFutureHealthDataset

  • 用于 query-conditioned all-future 监督
  • 对应 ExponentialLoss、WeibullLoss、MixedLoss

为了兼容旧训练入口：

HealthDataset = NextStepHealthDataset
collate_fn = next_step_collate_fn

all-future 训练入口会显式使用 AllFutureHealthDataset 和 all_future_collate_fn。

dataset 会输出：

event_seq
time_seq
padding_mask
sex
other_type
other_value
other_value_kind
other_time

其中 other_type=0 表示 padding，不额外传 other-token mask。

可以通过 extra_info_types 选择纳入哪些额外信息变量：

dataset = NextStepHealthDataset(extra_info_types=[1, 3, 7])

如果不传，则使用全部可用 other-info type。

dataset 会暴露模型初始化需要的元信息：

dataset.n_types
dataset.n_cont_types
dataset.n_categories
dataset.cont_type_ids
dataset.vocab_size

模型

模型主体定义在 models.py，通用网络模块定义在 backbones.py。

OtherInfoTokenizer

DeepHealth 内部的 OtherInfoTokenizer 编码统一的 other-info token：

other_type        # (B, K)
other_value       # (B, K)
other_value_kind  # (B, K)
other_time        # (B, K)

连续值使用 TokenAutoDiscretization：

type_id -> continuous type index -> soft bin embedding

分类值使用 dataset 动态计算后的 global category id：

selected type offsets + local category id

主序列内部保留展开的单个 extra-info token；疾病 token 因此可以直接关注到每个额外信息变量。只有在 next-token 训练需要把 extra-info 位置接入 readout/loss 时，同一患者、同一 other_time 的 extra hidden 才会被池化成一个监督用 token。默认池化方式是平均池化：

--extra_pool_reduce mean

也可以设为 sum。池化后，每个 extra-info 时间点最多产生一个用于预测和 loss 的 readout token。

DeepHealth

DeepHealth 的统一路径是：

disease tokens + expanded extra-info tokens
-> temporal backbone
-> risk head

extra-info 不再通过独立的 BaselineEncoder 或 CrossAttention 注入；它们作为主序列 token 直接参与同一个 temporal transformer。相对时间模式下，主序列内所有 token 共用 TimeRoPE 和 GaussianRBFTimeBasis。

两种目标模式的读出语义不同：

• next_token

  • 模型内部序列包含 disease tokens 和展开的单个 extra-info tokens
  • 默认 Tensor 返回值仍只返回 disease token hidden，形状为 (B, L, D)，兼容评估和旧调用
  • 训练时使用 return_output=True；此时同一 other_time 的 extra hidden 会在读出端池化，pooled extra readout tokens 也会产生 logits，并在有未来 disease target 时参与 prediction/loss 监督
  • next-token 模式下 risk_head.weight 与 token_embedding.weight 使用 weight tying

• all_future

  • 在合并后的主序列末尾拼接一个 query token
  • query token 的时间是 t_query
  • 只读出最后一个 query hidden，形状为 (B, D)
  • 展开的 extra-info tokens 只作为 query 的上下文输入，不单独读出、不参与 loss 监督
  • all-future 模式不使用 weight tying

如果需要拿到完整 next-token 输出，可使用结构化返回：

out = model(..., target_mode="next_token", return_output=True)
out.hidden        # disease tokens + pooled extra-info readout tokens
out.time_seq      # 与 hidden 对齐的时间
out.padding_mask  # 与 hidden 对齐的有效位置
out.event_len     # 原 disease token 长度

模型初始化示例：

model = DeepHealth(
    vocab_size=dataset.vocab_size,
    n_embd=120,
    n_head=10,
    n_hist_layer=12,
    n_tab_layer=4,  # 兼容旧配置；当前不再创建独立 tabular transformer
    n_types=dataset.n_types,
    n_cont_types=dataset.n_cont_types,
    n_categories=dataset.n_categories,
    cont_type_ids=dataset.cont_type_ids,
    extra_pool_reduce="mean",
)

Loss

losses.py 中保留：

next-token 监督：

• Delphi2MLoss
• UniqueTimeSetExponentialLoss

next-token 训练中，模型会请求 return_output=True，因此 loss 的预测位置包括：

• 原 disease token readout 位置
• 同一时间点 extra hidden 池化后的 pooled extra-info readout token

pooled extra-info readout token 的监督目标在训练时动态构造：对 pooled extra-info readout token 的时间 t，寻找该患者 t 之后的下一个 disease 事件时间；delphi2m 使用第一个未来事件作为 next-token target，uts 使用下一唯一时间点上的事件集合做 multi-hot target。若该 extra-info 时间点之后没有未来 disease target，则该位置不参与 loss。

all-future / query-conditioned 监督：

• ExponentialLoss
• WeibullLoss
• MixedLoss

all-future 训练只读出 t_query 对应的 query hidden。展开的 extra-info tokens 作为主序列上下文输入，但不会被单独读出，也不会被纳入 loss 监督。

UniqueTimeSetExponentialLoss 的 observed term 固定使用 sum reduction，不再暴露旧的 observed_reduction 参数。

训练

当前提供两类训练入口：

• train_next_step.py
  • 使用 NextStepHealthDataset
  • --target_mode delphi2m 默认搭配 Delphi2MLoss + token readout
  • --target_mode uts 默认搭配 UniqueTimeSetExponentialLoss + same_time_group_end readout
  • 当前 next-token 训练只支持 exponential time loss
  • 展开的 extra-info tokens 进入主序列；读出端 pooled extra-info tokens 会加入 prediction/loss 监督
• train_all_future.py
  • 使用 AllFutureHealthDataset
  • 不使用 readout，直接对 query hidden 计算风险
  • --dist_mode exponential/weibull/mixed 分别搭配 ExponentialLoss、WeibullLoss、MixedLoss
  • 展开的 extra-info tokens 只作为 query 上下文，不单独监督

当前 train_next_step.py / train_all_future.py 支持所有已有训练目标定义的组合：

训练模式	时间模式	分布/监督	默认 loss/readout
next_token	relative, absolute	target_mode=delphi2m, dist_mode=exponential	Delphi2MLoss + token
next_token	relative, absolute	target_mode=uts, dist_mode=exponential	UniqueTimeSetExponentialLoss + same_time_group_end
all_future	relative, absolute	dist_mode=exponential	ExponentialLoss，无 readout
all_future	relative, absolute	dist_mode=weibull	WeibullLoss，无 readout
all_future	relative, absolute	dist_mode=mixed	MixedLoss，无 readout

示例：

python train_next_step.py \
  --data_prefix ukb \
  --labels_file labels.csv \
  --target_mode uts \
  --n_embd 120 \
  --n_head 10 \
  --n_hist_layer 12 \
  --n_tab_layer 4 \
  --extra_pool_reduce mean

all-future 示例：

python train_all_future.py \
  --data_prefix ukb \
  --labels_file labels.csv \
  --dist_mode weibull \
  --time_mode relative \
  --extra_pool_reduce mean

选择额外信息变量：

python train_next_step.py --extra_info_types_file extra_info_types_smoking_alcohol_bmi.txt

train_next_step.py / train_all_future.py 只接受 --extra_info_types_file 指定变量列表，不接受在 CLI 里直接输入 type id。文件可以每行一个 type id，也可以带 # 注释；如果不传 --extra_info_types_file，默认使用全部 other-info type。

训练输出的 train_config.json 会记录：

• extra_info_types_file：训练时使用的列表文件名
• extra_info_types：解析后的实际 type id 列表，用于评估脚本复现变量选择
• extra_pool_reduce：同一 other_time 的 extra-info tokens 池化方式，默认为 mean
• model_target_mode、time_mode、dist_mode、dataset_class、collate_fn、resolved_loss_name：用于评估脚本重建模型和输入方式

评估 AUC

当前提供两个 AUC 评估入口，二者都已适配新的 DeepHealth 模型和统一的 other-info token 输入；AUC 的 DeLong 计算、病例/对照筛选和分层聚合逻辑保持原评估脚本口径。

评估脚本默认使用常规 Tensor 返回值：next-token checkpoint 只缓存 disease token/readout 位置的 hidden；all-future checkpoint 只缓存 t_query query hidden。训练中额外加入监督的 pooled extra-info readout tokens 不作为 AUC 评估位置单独输出。

evaluate_auc.py

evaluate_auc.py 评估的是 next-step / token-level 预测位置上的疾病 AUC。

查询方式由 train_config.json 中的 model_target_mode 决定：

• model_target_mode="next_token"：使用训练 readout 对应的历史 token hidden 作为预测点。
• model_target_mode="all_future"：不使用 readout token，直接把每个预测点年龄作为 t_query 传入模型，取 query hidden 作为预测点。

核心流程：

• 按训练配置重新构建 HealthDataset 和 DeepHealth。
• 对评估 split 中的患者做模型推理，缓存每个预测点的 hidden。
• 对疾病 token 分块投影到 risk_head，避免一次性保存全词表 logits。
• AUC score 使用疾病对应的 eta/logit 排序分数；dist_mode 只用于正确构建模型，不会把分数转换成 horizon-specific risk probability。
• 对每个疾病、性别、年龄段、prediction offset 分别计算 AUC。
• 输出未池化分层结果和按疾病汇总后的结果。

典型用法：

python evaluate_auc.py \
  --run_path runs/your_run_dir \
  --eval_split test \
  --offsets 0.1,1,5,10

主要输出：

• df_auc_unpooled.csv
  • 疾病 token 在 sex、age bracket、offset 分层下的 AUC。
• df_both.csv
  • 按疾病 token 和 offset 聚合后的 AUC。

适合回答的问题：

• “模型在历史序列中的某个预测 token 上，提前 offset 年预测未来疾病的区分能力如何？”
• “不同年龄段、性别、提前量下，next-step 训练模型的疾病预测 AUC 如何？”

evaluate_auc_v2.py

evaluate_auc_v2.py 评估的是 landmark fixed-horizon incident disease AUC。

它不是使用已有序列中的普通 readout 位置，而是在指定 landmark age 构造一个 landmark query，然后评估该 landmark 后固定 horizon 内是否发生 incident disease。

查询方式由 train_config.json 中的 model_target_mode 决定：

• model_target_mode="next_token"：在 landmark age 人工插入一个 <NO_EVENT> token，取该 token 的 hidden 做风险分数。
• model_target_mode="all_future"：不插入 <NO_EVENT>，直接把 landmark age 作为 t_query 传入模型，取 query hidden 做风险分数。

核心流程：

• 为每个患者和 landmark age 构造 landmark query 样本。
• 根据模型模式插入 <NO_EVENT> token 或直接传 t_query，取 landmark/query hidden。
• 对疾病 token 分块投影到 risk_head；score_mode="risk" 时会根据 dist_mode 把线性输出转换为固定 horizon 风险概率。
• 分布转换规则与 all-future 训练损失一致：exponential 使用 1 - exp(-rate * horizon)；weibull 使用 1 - exp(-rate * horizon ** rho)；mixed 中普通疾病使用 exponential，死亡 endpoint 使用 Weibull death rho。
• score_mode="eta" 是诊断用排序分数，不使用 rho，因此不区分不同分布的风险曲线。
• 按疾病、性别、landmark age、horizon 计算 incident disease AUC。
• 可选择排除 horizon 内先于目标疾病发生的死亡竞争风险。

典型用法：

python evaluate_auc_v2.py \
  --run_path runs/your_run_dir \
  --eval_split test \
  --landmark_start 40 \
  --landmark_stop 80 \
  --landmark_step 5 \
  --horizons 1,5,10

主要输出：

• df_auc_landmark_unpooled.csv
  • 疾病 token 在 sex、landmark age、horizon 分层下的 AUC。
• df_auc_landmark.csv
  • 按疾病 token 和 horizon 聚合后的 landmark AUC。

适合回答的问题：

• “一个人在 40/45/50/... 岁这个固定年龄点，如果此前未患某病，未来 1/5/10 年内发生该病的风险区分能力如何？”
• “模型能否作为 landmark risk prediction 模型使用？”

两者区别

项目	evaluate_auc.py	evaluate_auc_v2.py
评估口径	next-step/token-level 预测点	landmark fixed-horizon incident risk
查询位置	next-token 用满足 offset 条件的最新 readout token；all-future 直接用该预测点年龄作为 t_query	next-token 用人工插入的 <NO_EVENT> landmark token；all-future 直接用 t_query
时间参数	offsets：预测点至少早于目标事件多少年	landmark_* 和 horizons：固定年龄点与未来窗口
score 与分布	使用 eta/logit 排序分数；不按 dist_mode 转换风险概率	score_mode="risk" 按 dist_mode 区分 exponential / Weibull / mixed；score_mode="eta" 不区分分布
病例定义	target table 中出现目标疾病的患者/事件	landmark 后 horizon 内首次发生目标疾病
对照定义	从未出现该疾病的患者的 eligible target occurrence	landmark 时未患病，且 horizon 内未发病并有足够随访
分层	sex + age bracket + offset	sex + landmark age + horizon
输出文件	df_auc_unpooled.csv, df_both.csv	df_auc_landmark_unpooled.csv, df_auc_landmark.csv
适用问题	提前若干年预测未来目标事件的 token-level AUC	固定年龄点未来固定年限 incident disease risk AUC

简单选择：

• 想复现/延续旧的 next-token Delphi 风格 AUC：用 evaluate_auc.py。
• 想做临床上更像 “某年龄点未来 N 年发病风险” 的 landmark AUC：用 evaluate_auc_v2.py。

主要文件

• prepare_data.py

  • UKB 原始数据到模型输入文件的 ETL

• dataset.py

  • next-step 和 all-future dataset
  • 动态选择 other-info type
  • 动态计算 categorical global id

• models.py

  • DeepHealth
  • DeepHealthOutput
  • OtherInfoTokenizer

• backbones.py

  • TimeRoPE
  • GaussianRBFTimeBasis
  • TemporalAttention
  • GPTBlock
  • TokenAutoDiscretization
  • AgeSinusoidalEncoding

• losses.py

  • next-token 和 all-future losses

• readouts.py

  • token readout
  • same-time group readout
  • last-valid readout

• evaluate_auc.py

  • next-step/token-level 疾病 AUC 评估
  • 使用 prediction offset、sex、age bracket 分层
  • next-token 模型使用 readout hidden；all-future 模型使用 t_query hidden
  • score 是疾病 eta/logit，不按分布转换为固定 horizon 风险概率

• evaluate_auc_v2.py

  • landmark fixed-horizon incident disease AUC 评估
  • next-token 模型通过插入 <NO_EVENT> landmark token 查询固定年龄点风险
  • all-future 模型直接通过 t_query 查询固定年龄点风险
  • score_mode="risk" 按 exponential / Weibull / mixed 分布计算固定 horizon 风险