diff --git a/evaluate_single_disease_mortality_attribution.py b/evaluate_single_disease_mortality_attribution.py index e254259..29c1672 100644 --- a/evaluate_single_disease_mortality_attribution.py +++ b/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))