Add target construction and training script for DeepHealth model

- Implemented target construction in `targets.py` for next-token and unique-time set supervision. - Added validation functions and utility methods for target building. - Created a comprehensive training script in `train.py` that includes data loading, model building, optimizer setup, and training loop with early stopping and logging. - Integrated loss functions and readout mechanisms based on target modes. - Established dataset splitting and DataLoader configurations for training, validation, and testing.
2026-06-12 10:28:16 +08:00
commit 5e979e061b
13 changed files with 9165 additions and 0 deletions
--- a/README.md
+++ b/README.md
@@ -0,0 +1,280 @@
+# DeepHealthNew
+
+这个目录包含 DeepHealth 的数据准备、数据集、模型、readout 和 loss 代码。当前版本的核心设计是：
+
+```text
+疾病序列 stream + 统一的额外信息 token stream
+```
+
+疾病、死亡、checkup 事件仍然保存在事件序列里；性别单独保存在 `basic_info`；其他体检、暴露、生活方式等信息统一整理成 `(type, value, value_kind, time)` token。
+
+## 数据准备
+
+运行：
+
+```bash
+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)`
+  - 每行是：
+    ```text
+    (eid, type, value, value_kind, time)
+    ```
+  - `type=0` 预留给 padding
+  - `value_kind=1` 表示连续变量
+  - `value_kind=2` 表示分类变量
+  - 缺失值不会生成 token
+  - 当前 UKB 额外信息只有一个时间点，所以 `time` 暂时都是 `date_of_assessment`
+
+- `cate_types.csv`
+  - 分类变量元信息
+  - 字段：
+    ```text
+    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`
+
+为了兼容旧训练入口：
+
+```python
+HealthDataset = NextStepHealthDataset
+collate_fn = next_step_collate_fn
+```
+
+dataset 会输出：
+
+```python
+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` 选择纳入哪些额外信息变量：
+
+```python
+dataset = NextStepHealthDataset(extra_info_types=[1, 3, 7])
+```
+
+如果不传，则使用全部可用 other-info type。
+
+dataset 会暴露模型初始化需要的元信息：
+
+```python
+dataset.n_types
+dataset.n_cont_types
+dataset.n_categories
+dataset.cont_type_ids
+dataset.vocab_size
+```
+
+## 模型
+
+模型主体定义在 `models.py`，通用网络模块定义在 `backbones.py`。
+
+### BaselineEncoder
+
+`BaselineEncoder` 编码统一的 other-info token：
+
+```python
+other_type        # (B, K)
+other_value       # (B, K)
+other_value_kind  # (B, K)
+```
+
+它暂时不直接使用 `other_time`。时间信息保留给后续 `CrossAttention`，用于建模疾病/query 与 other-info token 的相对时间关系。
+
+连续值使用 `TokenAutoDiscretization`：
+
+```text
+type_id -> continuous type index -> soft bin embedding
+```
+
+分类值使用 dataset 动态计算后的 global category id：
+
+```text
+selected type offsets + local category id
+```
+
+### CrossAttention
+
+`CrossAttention` 让 disease-side hidden state 注意到 other-info token：
+
+```python
+h_disease  # (B, L, D)
+t_disease  # (B, L)
+h_token    # (B, K, D)
+t_token    # (B, K)
+```
+
+时间信息通过两种方式进入注意力：
+
+- `TimeRoPE`
+  - 使用 query time 和 key time 旋转 q/k
+  - 让 q-k 相似度带有时间位置信息
+
+- `GaussianRBFTimeBasis`
+  - 对 `t_disease - t_token` 做 RBF 编码
+  - 投影成每个 attention head 的时间 bias
+
+注意力是时间因果的：
+
+```text
+other_info_time <= disease_or_query_time
+```
+
+如果某个 disease/query 位置没有任何可见 other-info token，则该位置保持原 hidden 不变。
+
+### DeepHealth
+
+`DeepHealth` 的统一路径是：
+
+```text
+disease-side sequence
+-> disease temporal backbone
+-> CrossAttention 到 other-info tokens
+-> risk head
+```
+
+两种目标模式共用同一套语义：
+
+- `next_token`
+  - `h_disease` 长度为 `L`
+  - 输出 `(B, L, D)`
+
+- `all_future`
+  - 在 disease-side 序列末尾拼接一个 query token
+  - query token 的时间是 `t_query`
+  - `h_disease` 长度为 `L + 1`
+  - 输出最后一个 query hidden，形状为 `(B, D)`
+
+模型初始化示例：
+
+```python
+model = DeepHealth(
+    vocab_size=dataset.vocab_size,
+    n_embd=120,
+    n_head=10,
+    n_hist_layer=12,
+    n_tab_layer=4,
+    n_types=dataset.n_types,
+    n_cont_types=dataset.n_cont_types,
+    n_categories=dataset.n_categories,
+    cont_type_ids=dataset.cont_type_ids,
+)
+```
+
+## Loss
+
+`losses.py` 中保留：
+
+next-token 监督：
+
+- `Delphi2MLoss`
+- `UniqueTimeSetExponentialLoss`
+
+all-future / query-conditioned 监督：
+
+- `ExponentialLoss`
+- `WeibullLoss`
+- `MixedLoss`
+
+`UniqueTimeSetExponentialLoss` 的 observed term 固定使用 sum reduction，不再暴露旧的 `observed_reduction` 参数。
+
+## 训练
+
+当前 `train.py` 是 next-step 训练入口，使用：
+
+```python
+HealthDataset = NextStepHealthDataset
+```
+
+示例：
+
+```bash
+python train.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
+```
+
+选择额外信息变量：
+
+```bash
+python train.py --extra_info_types 1 3 7
+```
+
+如果不传 `--extra_info_types`，默认使用全部 other-info type。
+
+## 主要文件
+
+- `prepare_data.py`
+  - UKB 原始数据到模型输入文件的 ETL
+
+- `dataset.py`
+  - next-step 和 all-future dataset
+  - 动态选择 other-info type
+  - 动态计算 categorical global id
+
+- `models.py`
+  - `DeepHealth`
+
+- `backbones.py`
+  - `TimeRoPE`
+  - `GaussianRBFTimeBasis`
+  - `TemporalAttention`
+  - `GPTBlock`
+  - `TokenAutoDiscretization`
+  - `BaselineEncoder`
+  - `CrossAttention`
+  - `AgeSinusoidalEncoding`
+
+- `losses.py`
+  - next-token 和 all-future losses
+
+- `readouts.py`
+  - token readout
+  - same-time group readout
+  - last-valid readout