From 8928688aca2b4ae9a0db68d678374c1a5eaa99bd Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Sat, 27 Jun 2026 15:43:27 +0800 Subject: [PATCH] Optimize per-disease attribution scanning --- ...te_single_disease_mortality_attribution.py | 78 ++++++++++++------- 1 file changed, 51 insertions(+), 27 deletions(-) diff --git a/evaluate_single_disease_mortality_attribution.py b/evaluate_single_disease_mortality_attribution.py index 69abb29..0104198 100644 --- a/evaluate_single_disease_mortality_attribution.py +++ b/evaluate_single_disease_mortality_attribution.py @@ -438,7 +438,6 @@ def main() -> None: if not scanned_disease_items: raise ValueError("No diseases selected for attribution") scanned_disease_tokens = [token for token, _meta in scanned_disease_items] - scanned_disease_token_set = set(scanned_disease_tokens) landmark_ages = make_landmark_ages( float(args.landmark_start), @@ -476,6 +475,11 @@ def main() -> None: cfg_model = dict(cfg) cfg_model["dist_mode"] = dist_mode device = resolve_eval_device(args.device) + scanned_disease_tensor = torch.as_tensor( + scanned_disease_tokens, + dtype=torch.long, + device=device, + ) model = build_model_from_dataset(args, cfg_model, dataset).to(device) load_model_state(model, state_dict) model.eval() @@ -537,6 +541,7 @@ def main() -> None: shard_index = 0 shards: list[dict[str, Any]] = [] summary_accumulator: dict[tuple[Any, ...], dict[str, float]] = {} + row_base_cache: dict[int, dict[str, Any]] = {} pending_batch_chunks: list[Dict[str, torch.Tensor]] = [] pending_meta_chunks: list[list[dict[str, Any]]] = [] pending_n = 0 @@ -602,6 +607,36 @@ def main() -> None: pending_meta_chunks = [] pending_n = 0 + def get_row_base(row_idx: int) -> dict[str, Any]: + cached = row_base_cache.get(row_idx) + if cached is not None: + return cached + + meta = landmark_dataset.rows[int(row_idx)] + dataset_index = int(meta["dataset_index"]) + sample = dataset.samples[dataset_index] + hist_tokens = np.asarray(meta["event_seq"], dtype=np.int64) + total_count, group_counts = historical_counts_by_group( + hist_tokens, + death_idx=death_idx, + token_to_group=token_to_group, + group_names=group_names, + ) + cached = { + "patient_id": int(meta["patient_id"]), + "dataset_index": dataset_index, + "eid": int(sample.get("eid", -1)), + "sex": int(meta["sex"]), + "landmark_age": float(meta["landmark_age"]), + "tau": tau, + "followup_end_time": float(meta["followup_end_time"]), + "history_disease_count": int(total_count), + "_hist_tokens": hist_tokens, + "_group_counts": group_counts, + } + row_base_cache[row_idx] = cached + return cached + for batch in tqdm(loader, desc="Per-disease mortality attribution", dynamic_ncols=True): hidden = infer_landmark_hidden( model=model, @@ -629,14 +664,10 @@ def main() -> None: vocab_size=int(dataset.vocab_size), device=device, ) - event_values = batch["event_seq"].detach().cpu().numpy().reshape(-1) - batch_disease_tokens = sorted( - { - int(token) - for token in event_values.tolist() - if int(token) in scanned_disease_token_set - } - ) + active_token_mask = occurred[:, scanned_disease_tensor].any(dim=0) + batch_disease_tokens = scanned_disease_tensor[ + active_token_mask + ].detach().cpu().tolist() if not batch_disease_tokens: continue @@ -663,16 +694,9 @@ def main() -> None: meta_chunk: list[dict[str, Any]] = [] row_ids = batch["row_idx"][row_indices.cpu()].cpu().numpy().astype(np.int64) for local_pos, row_idx in enumerate(row_ids.tolist()): - meta = landmark_dataset.rows[int(row_idx)] - dataset_index = int(meta["dataset_index"]) - sample = dataset.samples[dataset_index] - hist_tokens = np.asarray(meta["event_seq"], dtype=np.int64) - total_count, group_counts = historical_counts_by_group( - hist_tokens, - death_idx=death_idx, - token_to_group=token_to_group, - group_names=group_names, - ) + row_base = get_row_base(int(row_idx)) + hist_tokens = row_base["_hist_tokens"] + group_counts = row_base["_group_counts"] disease_history_count = int((hist_tokens == int(disease_token)).sum()) if disease_history_count <= 0: continue @@ -680,14 +704,14 @@ def main() -> None: orig_row = int(row_indices[local_pos].detach().cpu()) meta_chunk.append( { - "patient_id": int(meta["patient_id"]), - "dataset_index": dataset_index, - "eid": int(sample.get("eid", -1)), - "sex": int(meta["sex"]), - "landmark_age": float(meta["landmark_age"]), - "tau": tau, - "followup_end_time": float(meta["followup_end_time"]), - "history_disease_count": int(total_count), + "patient_id": row_base["patient_id"], + "dataset_index": row_base["dataset_index"], + "eid": row_base["eid"], + "sex": row_base["sex"], + "landmark_age": row_base["landmark_age"], + "tau": row_base["tau"], + "followup_end_time": row_base["followup_end_time"], + "history_disease_count": row_base["history_disease_count"], "selected_disease_history_count": disease_history_count, "selected_disease_token_id": int(disease_token), "selected_disease_code": str(disease_meta.get("code", "")),