Align disease attribution batching with survival evaluator

evaluate_single_disease_mortality_attribution.py

@@ -196,48 +196,6 @@ def write_summary_csv(
     return len(rows)
 
 
-def concat_padded_tensor_batches(chunks: list[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
-    if not chunks:
-        raise ValueError("Cannot concatenate an empty chunk list")
-
-    fill_values = {
-        "event_seq": PAD_IDX,
-        "time_seq": 0.0,
-        "readout_mask": False,
-        "padding_mask": False,
-        "other_type": 0,
-        "other_value": 0.0,
-        "other_value_kind": 0,
-        "other_time": 0.0,
-    }
-    out: Dict[str, torch.Tensor] = {}
-    for key in chunks[0]:
-        tensors = [chunk[key] for chunk in chunks]
-        shapes = [tuple(t.shape) for t in tensors]
-        if len(set(shapes)) == 1:
-            out[key] = torch.cat(tensors, dim=0)
-            continue
-
-        if any(t.ndim == 0 for t in tensors):
-            raise ValueError(f"Cannot concatenate scalar tensor key={key!r} with mismatched shapes")
-        max_shape = list(shapes[0])
-        for shape in shapes[1:]:
-            if len(shape) != len(max_shape):
-                raise ValueError(f"Cannot concatenate key={key!r} with shapes {shapes}")
-            max_shape = [max(a, b) for a, b in zip(max_shape, shape)]
-
-        padded: list[torch.Tensor] = []
-        fill = fill_values.get(key, 0)
-        for tensor in tensors:
-            target_shape = [int(tensor.shape[0]), *max_shape[1:]]
-            padded_tensor = tensor.new_full(target_shape, fill)
-            slices = tuple(slice(0, int(size)) for size in tensor.shape)
-            padded_tensor[slices] = tensor
-            padded.append(padded_tensor)
-        out[key] = torch.cat(padded, dim=0)
-    return out
-
-
 def build_disease_ablated_slice(
     batch: Dict[str, torch.Tensor],
     row_indices: torch.Tensor,
@@ -469,6 +427,12 @@ def parse_args() -> argparse.Namespace:
         default=1e-7,
         help="Small lower bound for ratio denominators.",
     )
+    parser.add_argument(
+        "--shard_rows",
+        type=int,
+        default=200_000,
+        help="Approximate number of detailed rows to buffer before writing one .npz shard.",
+    )
     return parser.parse_args()
 
 
@@ -561,6 +525,8 @@ def main() -> None:
         raise ValueError("attribution_batch_size must be positive")
     if float(args.ratio_eps) <= 0:
         raise ValueError("--ratio_eps must be positive")
+    if int(args.shard_rows) <= 0:
+        raise ValueError("--shard_rows must be positive")
 
     num_workers = int(cfg_get(args, cfg, "num_workers", 4))
     loader = DataLoader(
@@ -611,51 +577,27 @@ def main() -> None:
     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_orig_risk_chunks: list[torch.Tensor] = []
-    pending_orig_hazard_chunks: list[torch.Tensor] = []
-    pending_n = 0
+    shard_frames: list[pd.DataFrame] = []
+    shard_buffer_rows = 0
 
-    def flush_pending() -> None:
-        nonlocal written_rows, shard_index, pending_batch_chunks, pending_meta_chunks
-        nonlocal pending_orig_risk_chunks, pending_orig_hazard_chunks, pending_n
-        if pending_n == 0:
+    def flush_shard_buffer() -> None:
+        nonlocal written_rows, shard_index, shard_frames, shard_buffer_rows
+        if shard_buffer_rows == 0:
             return
+        table = pd.concat(shard_frames, ignore_index=True).reindex(columns=OUTPUT_COLUMNS)
+        shard_name = f"part-{shard_index:06d}.npz"
+        shard_path = output_dir / shard_name
+        shard_rows = write_compressed_npz_table(shard_path, table)
+        shards.append({"file": shard_name, "rows": int(shard_rows)})
+        shard_index += 1
+        written_rows += int(shard_rows)
+        shard_frames = []
+        shard_buffer_rows = 0
 
-        ablated_batch = concat_padded_tensor_batches(pending_batch_chunks)
-        meta_rows = [row for chunk in pending_meta_chunks for row in chunk]
-        orig_risk = torch.cat(pending_orig_risk_chunks, dim=0).to(device)
-        orig_hazard = torch.cat(pending_orig_hazard_chunks, dim=0).to(device)
-        with torch.no_grad():
-            ablated_risk = death_risk_for_batch(
-                model=model,
-                batch=ablated_batch,
-                device=device,
-                model_target_mode=model_target_mode,
-                readout_name=readout_name,
-                readout_reduce=readout_reduce,
-                dist_mode=dist_mode,
-                tau=tau,
-            )
-        ablated_hazard = mortality_hazard_from_risk(ablated_risk)
-        attr_prob = orig_risk - ablated_risk
-        attr_hazard = orig_hazard - ablated_hazard
-        ratio_prob = safe_ratio(orig_risk, ablated_risk, eps=float(args.ratio_eps))
-        ratio_hazard = safe_ratio(orig_hazard, ablated_hazard, eps=float(args.ratio_eps))
-        value_block = torch.stack(
-            [
-                orig_risk,
-                orig_hazard,
-                ablated_risk,
-                ablated_hazard,
-                attr_prob,
-                attr_hazard,
-                ratio_prob,
-                ratio_hazard,
-            ],
-            dim=1,
-        ).detach().cpu().numpy()
+    def write_result_rows(meta_rows: list[dict[str, Any]], value_block: np.ndarray) -> None:
+        nonlocal shard_buffer_rows
+        if not meta_rows:
+            return
 
         for i, row in enumerate(meta_rows):
             row["death_risk"] = float(value_block[i, 0])
@@ -669,17 +611,10 @@ def main() -> None:
 
         table = pd.DataFrame(meta_rows).reindex(columns=OUTPUT_COLUMNS)
         update_summary_accumulator(summary_accumulator, table)
-        shard_name = f"part-{shard_index:06d}.npz"
-        shard_path = output_dir / shard_name
-        shard_rows = write_compressed_npz_table(shard_path, table)
-        shards.append({"file": shard_name, "rows": int(shard_rows)})
-        shard_index += 1
-        written_rows += len(meta_rows)
-        pending_batch_chunks = []
-        pending_meta_chunks = []
-        pending_orig_risk_chunks = []
-        pending_orig_hazard_chunks = []
-        pending_n = 0
+        shard_frames.append(table)
+        shard_buffer_rows += len(table)
+        if shard_buffer_rows >= int(args.shard_rows):
+            flush_shard_buffer()
 
     def get_row_base(row_idx: int) -> dict[str, Any]:
         cached = row_base_cache.get(row_idx)
@@ -756,8 +691,7 @@ def main() -> None:
 
         pair_offset = 0
         while pair_offset < int(pair_indices.shape[0]):
-            capacity = int(attribution_batch_size) - pending_n
-            pair_stop = min(int(pair_indices.shape[0]), pair_offset + capacity)
+            pair_stop = min(int(pair_indices.shape[0]), pair_offset + int(attribution_batch_size))
             pair_chunk = pair_indices[pair_offset:pair_stop]
             local_rows = pair_chunk[:, 0].long()
             disease_token_ids = scanned_disease_tensor[pair_chunk[:, 1]].long()
@@ -766,6 +700,37 @@ def main() -> None:
                 row_indices=local_rows,
                 token_ids=disease_token_ids,
             )
+            with torch.no_grad():
+                ablated_risk = death_risk_for_batch(
+                    model=model,
+                    batch=ablated_chunk,
+                    device=device,
+                    model_target_mode=model_target_mode,
+                    readout_name=readout_name,
+                    readout_reduce=readout_reduce,
+                    dist_mode=dist_mode,
+                    tau=tau,
+                )
+            orig_risk = death_risk_tensor[local_rows]
+            orig_hazard = death_hazard_tensor[local_rows]
+            ablated_hazard = mortality_hazard_from_risk(ablated_risk)
+            attr_prob = orig_risk - ablated_risk
+            attr_hazard = orig_hazard - ablated_hazard
+            ratio_prob = safe_ratio(orig_risk, ablated_risk, eps=float(args.ratio_eps))
+            ratio_hazard = safe_ratio(orig_hazard, ablated_hazard, eps=float(args.ratio_eps))
+            value_block = torch.stack(
+                [
+                    orig_risk,
+                    orig_hazard,
+                    ablated_risk,
+                    ablated_hazard,
+                    attr_prob,
+                    attr_hazard,
+                    ratio_prob,
+                    ratio_hazard,
+                ],
+                dim=1,
+            ).detach().cpu().numpy()
 
             meta_chunk: list[dict[str, Any]] = []
             row_ids = batch_dev["row_idx"][local_rows].detach().cpu().numpy().astype(np.int64)
@@ -810,18 +775,10 @@ def main() -> None:
                     }
                 )
 
-            if meta_chunk:
-                pending_batch_chunks.append(ablated_chunk)
-                pending_meta_chunks.append(meta_chunk)
-                pending_orig_risk_chunks.append(death_risk_tensor[local_rows].detach())
-                pending_orig_hazard_chunks.append(death_hazard_tensor[local_rows].detach())
-                pending_n += len(meta_chunk)
+            write_result_rows(meta_chunk, value_block)
             pair_offset = pair_stop
 
-            if pending_n >= int(attribution_batch_size):
-                flush_pending()
-
-    flush_pending()
+    flush_shard_buffer()
     if not shards:
         empty_path = output_dir / "part-000000.npz"
         write_compressed_npz_table(empty_path, pd.DataFrame(columns=OUTPUT_COLUMNS))