Add exposure cache and keep absolute time only

This commit is contained in:
2026-07-07 17:21:52 +08:00
parent a0379daf29
commit 45a857d1a6
9 changed files with 690 additions and 198 deletions

View File

@@ -321,9 +321,16 @@ def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], data
n_head=int(cfg_get(args, cfg, "n_head", 10)),
n_hist_layer=int(cfg_get(args, cfg, "n_hist_layer", 12)),
target_mode=model_target_mode,
time_mode=str(cfg_get(args, cfg, "time_mode", "relative")),
dist_mode=str(cfg_get(args, cfg, "dist_mode", "exponential")),
dropout=float(cfg_get(args, cfg, "dropout", 0.0)),
use_exposure_encoder=bool(cfg_get(args, cfg, "use_exposure_encoder", False)),
exposure_d_model=cfg_get(args, cfg, "exposure_d_model", None),
exposure_n_layers=int(cfg_get(args, cfg, "exposure_n_layers", 2)),
exposure_top_k=int(cfg_get(args, cfg, "exposure_top_k", 3)),
exposure_n_convnext_blocks=int(cfg_get(args, cfg, "exposure_n_convnext_blocks", 2)),
exposure_conv_kernel_size=int(cfg_get(args, cfg, "exposure_conv_kernel_size", 7)),
exposure_mlp_ratio=float(cfg_get(args, cfg, "exposure_mlp_ratio", 4.0)),
exposure_use_gate=bool(cfg_get(args, cfg, "exposure_use_gate", True)),
)
@@ -524,6 +531,16 @@ def infer_readout_hidden(
sex=batch_dev["sex"][active],
padding_mask=padding_mask[active],
t_query=time_seq[active, pos],
exposure_daily=(
batch_dev["exposure_daily"][active]
if "exposure_daily" in batch_dev
else None
),
exposure_monthly=(
batch_dev["exposure_monthly"][active]
if "exposure_monthly" in batch_dev
else None
),
target_mode="all_future",
)
hidden[active, pos, :] = hidden_pos.float()
@@ -534,6 +551,8 @@ def infer_readout_hidden(
time_seq=time_seq,
sex=batch_dev["sex"],
padding_mask=padding_mask,
exposure_daily=batch_dev.get("exposure_daily"),
exposure_monthly=batch_dev.get("exposure_monthly"),
target_mode="next_token",
)
ro = readout(
@@ -1317,6 +1336,8 @@ def main() -> None:
include_no_event_in_uts_target=include_no_event,
min_history_events=int(cfg.get("all_future_min_history_events", 1)),
min_future_events=int(cfg.get("all_future_min_future_events", 1)),
exposure_cache_dir=cfg.get("exposure_cache_dir", None),
mask_onset_exposure=bool(cfg.get("mask_onset_exposure", False)),
)
validate_dataset_metadata(dataset, cfg)