Use precomputed exposure embeddings

This commit is contained in:
2026-07-09 16:49:49 +08:00
parent 552e09aa01
commit e439c3c98c
7 changed files with 311 additions and 281 deletions

View File

@@ -323,14 +323,9 @@ def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], data
target_mode=model_target_mode,
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, "d_model", 64),
exposure_n_layers=int(cfg_get(args, cfg, "exposure_n_layers", 2)),
exposure_top_k=int(cfg_get(args, cfg, "exposure_top_k", 2)),
exposure_n_backbone_blocks=int(cfg_get(args, cfg, "exposure_n_backbone_blocks", 1)),
exposure_backbone_kernel_size=int(cfg_get(args, cfg, "exposure_backbone_kernel_size", 5)),
exposure_backbone_expansion=float(cfg_get(args, cfg, "exposure_backbone_expansion", 2.0)),
exposure_use_gate=bool(cfg_get(args, cfg, "exposure_use_gate", True)),
use_exposure_embeddings=bool(
cfg_get(args, cfg, "use_exposure_embeddings", False)
),
)
@@ -531,14 +526,9 @@ 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
exposure_embedding=(
batch_dev["exposure_embedding"][active]
if "exposure_embedding" in batch_dev
else None
),
target_mode="all_future",
@@ -551,8 +541,7 @@ 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"),
exposure_embedding=batch_dev.get("exposure_embedding"),
target_mode="next_token",
)
ro = readout(
@@ -1337,7 +1326,7 @@ def main() -> None:
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)),
exposure_embeddings_file=cfg.get("exposure_embeddings_file", None),
)
validate_dataset_metadata(dataset, cfg)