JiaruiLi/DeepHealth
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Batch extra-info disease attribution stats

Browse Source
This commit is contained in:
Jiarui Li
2026-07-01 12:47:30 +08:00
parent d3e68ac09b
commit e3c9ecd19f
1 changed files with 149 additions and 115 deletions
Download Patch File Download Diff File Expand all files Collapse all files

264
evaluate_extra_info_attribution.py
View File

@@ -212,38 +212,6 @@ def death_distribution_parameters(
return "weibull", torch.stack([nan, scale, shape], dim=1) return "weibull", torch.stack([nan, scale, shape], dim=1)
def disease_distribution_parameters(
model,
hidden: torch.Tensor,
*,
token_ids: Sequence[int],
dist_mode: str,
logits: torch.Tensor | None = None,
rho: torch.Tensor | None = None,
eps: float = 1e-8,
) -> tuple[str, torch.Tensor]:
ids = [int(x) for x in token_ids]
if not ids:
empty = hidden.new_empty((hidden.shape[0], 0, 3))
return "none", empty
all_logits = model.calc_risk(hidden) if logits is None else logits
disease_lambda = F.softplus(all_logits[:, ids]) + float(eps)
if dist_mode in {"exponential", "mixed"}:
nan = torch.full_like(disease_lambda, float("nan"))
return "exponential", torch.stack([disease_lambda, nan, nan], dim=2)
if dist_mode == "weibull":
all_rho = model.calc_weibull_rho(hidden) if rho is None else rho
shape = all_rho[:, ids].to(dtype=disease_lambda.dtype).clamp_min(float(eps))
scale = torch.pow(disease_lambda.clamp_min(float(eps)), -1.0 / shape)
nan = torch.full_like(disease_lambda, float("nan"))
return "weibull", torch.stack([nan, scale, shape], dim=2)
raise ValueError(f"Unsupported dist_mode={dist_mode!r}")
def parameter_pair_block(original: torch.Tensor, ablated: torch.Tensor) -> torch.Tensor: def parameter_pair_block(original: torch.Tensor, ablated: torch.Tensor) -> torch.Tensor:
return torch.stack( return torch.stack(
[ [
@@ -258,17 +226,70 @@ def parameter_pair_block(original: torch.Tensor, ablated: torch.Tensor) -> torch
) )
def disease_parameter_pair_block(original: torch.Tensor, ablated: torch.Tensor) -> torch.Tensor: def all_disease_parameter_pair_block(
return torch.stack( *,
[ original_logits: torch.Tensor,
original[:, :, 0], ablated_logits: torch.Tensor,
ablated[:, :, 0], dist_mode: str,
original[:, :, 1], original_rho: torch.Tensor | None = None,
ablated[:, :, 1], ablated_rho: torch.Tensor | None = None,
original[:, :, 2], eps: float = 1e-8,
ablated[:, :, 2], ) -> torch.Tensor:
], original_lambda = F.softplus(original_logits) + float(eps)
dim=2, ablated_lambda = F.softplus(ablated_logits) + float(eps)
if dist_mode in {"exponential", "mixed"}:
nan = torch.full_like(original_lambda, float("nan"))
return torch.stack(
[
original_lambda,
ablated_lambda,
nan,
nan,
nan,
nan,
],
dim=2,
)
if dist_mode == "weibull":
if original_rho is None or ablated_rho is None:
raise ValueError("rho tensors are required for weibull disease parameters")
original_shape = original_rho.to(dtype=original_lambda.dtype).clamp_min(float(eps))
ablated_shape = ablated_rho.to(dtype=ablated_lambda.dtype).clamp_min(float(eps))
original_scale = torch.pow(original_lambda.clamp_min(float(eps)), -1.0 / original_shape)
ablated_scale = torch.pow(ablated_lambda.clamp_min(float(eps)), -1.0 / ablated_shape)
nan = torch.full_like(original_lambda, float("nan"))
return torch.stack(
[
nan,
nan,
original_scale,
ablated_scale,
original_shape,
ablated_shape,
],
dim=2,
)
raise ValueError(f"Unsupported dist_mode={dist_mode!r}")
def grouped_parameter_stats(
values: torch.Tensor,
group_token_mask: torch.Tensor,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
finite = torch.isfinite(values)
values64 = values.to(dtype=torch.float64)
safe_values = torch.where(finite, values64, torch.zeros_like(values64))
mask = group_token_mask.to(device=values.device, dtype=torch.float64)
sums = torch.einsum("nvc,gv->ngc", safe_values, mask)
sumsq = torch.einsum("nvc,gv->ngc", safe_values * safe_values, mask)
counts = torch.einsum("nvc,gv->ngc", finite.to(dtype=torch.float64), mask)
return (
sums.detach().cpu().numpy().astype(np.float64, copy=False),
sumsq.detach().cpu().numpy().astype(np.float64, copy=False),
counts.detach().cpu().numpy().astype(np.float64, copy=False),
) )
@@ -401,39 +422,44 @@ def update_death_summary(
acc[f"sumsq__{column}"] += float(np.square(vals).sum()) acc[f"sumsq__{column}"] += float(np.square(vals).sum())
def update_disease_parameter_summary( def update_disease_parameter_summary_from_group_stats(
summary: dict[tuple[Any, ...], dict[str, float]], summary: dict[tuple[Any, ...], dict[str, float]],
*, *,
key_rows: pd.DataFrame, key_rows: pd.DataFrame,
target_group: str, group_names: Sequence[str],
target_group_label: str, group_labels: Sequence[str],
values: np.ndarray, sums: np.ndarray,
sumsq: np.ndarray,
counts: np.ndarray,
) -> None: ) -> None:
if key_rows.empty or values.size == 0: if key_rows.empty or sums.size == 0:
return return
table = key_rows.reset_index(drop=True).copy() rows = key_rows.reset_index(drop=True)
grouped = table.groupby(EXTRA_KEY_COLUMNS, dropna=False, sort=False) for row_idx, row in rows.iterrows():
for key, group in grouped: base_key = tuple(row[column] for column in EXTRA_KEY_COLUMNS)
if not isinstance(key, tuple): for group_idx, (group, label) in enumerate(zip(group_names, group_labels)):
key = (key,) count_row = counts[int(row_idx), int(group_idx)]
full_key = (*key, str(target_group), str(target_group_label)) n_add = float(np.nanmax(count_row)) if count_row.size else 0.0
idx = group.index.to_numpy(dtype=np.int64) if n_add <= 0:
vals_2d = values[idx].reshape(-1, values.shape[-1]) continue
acc = summary.setdefault( full_key = (*base_key, str(group), str(label))
full_key, acc = summary.setdefault(
{ full_key,
"n": 0.0, {
**{f"count__{col}": 0.0 for col in DISEASE_PARAMETER_COLUMNS}, "n": 0.0,
**{f"sum__{col}": 0.0 for col in DISEASE_PARAMETER_COLUMNS}, **{f"count__{col}": 0.0 for col in DISEASE_PARAMETER_COLUMNS},
**{f"sumsq__{col}": 0.0 for col in DISEASE_PARAMETER_COLUMNS}, **{f"sum__{col}": 0.0 for col in DISEASE_PARAMETER_COLUMNS},
}, **{f"sumsq__{col}": 0.0 for col in DISEASE_PARAMETER_COLUMNS},
) },
acc["n"] += float(vals_2d.shape[0]) )
for col_idx, column in enumerate(DISEASE_PARAMETER_COLUMNS): acc["n"] += n_add
vals = finite_float64(vals_2d[:, col_idx]) for col_idx, column in enumerate(DISEASE_PARAMETER_COLUMNS):
acc[f"count__{column}"] += float(vals.size) count = float(counts[int(row_idx), int(group_idx), int(col_idx)])
acc[f"sum__{column}"] += float(vals.sum()) if count <= 0:
acc[f"sumsq__{column}"] += float(np.square(vals).sum()) continue
acc[f"count__{column}"] += count
acc[f"sum__{column}"] += float(sums[int(row_idx), int(group_idx), int(col_idx)])
acc[f"sumsq__{column}"] += float(sumsq[int(row_idx), int(group_idx), int(col_idx)])
def write_death_summary_csv( def write_death_summary_csv(
@@ -624,6 +650,8 @@ def main() -> None:
"all_modeled_diseases": "All modeled diseases", "all_modeled_diseases": "All modeled diseases",
**organ_labels, **organ_labels,
} }
group_names = list(risk_groups.keys())
group_labels = [str(risk_group_labels[group]) for group in group_names]
state_dict = load_checkpoint_state_dict(checkpoint_path, map_location="cpu") state_dict = load_checkpoint_state_dict(checkpoint_path, map_location="cpu")
dist_mode = resolve_dist_mode_for_checkpoint(str(cfg.get("dist_mode", "exponential")), state_dict) dist_mode = resolve_dist_mode_for_checkpoint(str(cfg.get("dist_mode", "exponential")), state_dict)
@@ -635,6 +663,20 @@ def main() -> None:
load_model_state(model, state_dict) load_model_state(model, state_dict)
model.eval() model.eval()
group_token_mask = torch.zeros(
(len(group_names), int(dataset.vocab_size)),
dtype=torch.float32,
device=device,
)
for group_idx, group in enumerate(group_names):
valid_tokens = [
int(token)
for token in risk_groups[group]
if 0 <= int(token) < int(dataset.vocab_size) and int(token) != death_idx
]
if valid_tokens:
group_token_mask[group_idx, torch.as_tensor(valid_tokens, dtype=torch.long, device=device)] = 1.0
batch_size = int(cfg_get(args, cfg, "batch_size", 128)) batch_size = int(cfg_get(args, cfg, "batch_size", 128))
attribution_batch_size = int( attribution_batch_size = int(
cfg_get(args, cfg, "attribution_batch_size", max(batch_size * 32, 4096)) cfg_get(args, cfg, "attribution_batch_size", max(batch_size * 32, 4096))
@@ -675,6 +717,9 @@ def main() -> None:
death_summary: dict[tuple[Any, ...], dict[str, float]] = {} death_summary: dict[tuple[Any, ...], dict[str, float]] = {}
disease_parameter_summary: dict[tuple[Any, ...], dict[str, float]] = {} disease_parameter_summary: dict[tuple[Any, ...], dict[str, float]] = {}
death_key_chunks: list[pd.DataFrame] = []
death_value_chunks: list[np.ndarray] = []
disease_stat_chunks: list[tuple[pd.DataFrame, np.ndarray, np.ndarray, np.ndarray]] = []
for batch in tqdm(loader, desc="Extra-info attribution", dynamic_ncols=True): for batch in tqdm(loader, desc="Extra-info attribution", dynamic_ncols=True):
batch_dev = { batch_dev = {
@@ -697,19 +742,6 @@ def main() -> None:
) )
original_logits = model.calc_risk(hidden) original_logits = model.calc_risk(hidden)
original_rho = model.calc_weibull_rho(hidden) if dist_mode == "weibull" else None original_rho = model.calc_weibull_rho(hidden) if dist_mode == "weibull" else None
original_disease_params_by_group = {}
disease_distribution_name_by_group = {}
for group, tokens in risk_groups.items():
disease_distribution, original_disease_params = disease_distribution_parameters(
model,
hidden,
token_ids=tokens,
dist_mode=dist_mode,
logits=original_logits,
rho=original_rho,
)
disease_distribution_name_by_group[group] = disease_distribution
original_disease_params_by_group[group] = original_disease_params
for ablated_batch, type_ids, local_rows in iter_extra_info_ablated_batches( for ablated_batch, type_ids, local_rows in iter_extra_info_ablated_batches(
batch_dev, batch_dev,
@@ -751,38 +783,40 @@ def main() -> None:
original_death_params[row_tensor], original_death_params[row_tensor],
ablated_death_params, ablated_death_params,
).detach().cpu().numpy() ).detach().cpu().numpy()
update_death_summary( death_key_chunks.append(key_table)
death_summary, death_value_chunks.append(value_block)
key_rows=key_table,
values=value_block,
)
for group, tokens in risk_groups.items(): disease_values = all_disease_parameter_pair_block(
disease_distribution, ablated_disease_params = disease_distribution_parameters( original_logits=original_logits[row_tensor],
model, ablated_logits=ablated_logits,
ablated_hidden, dist_mode=dist_mode,
token_ids=tokens, original_rho=None if original_rho is None else original_rho[row_tensor],
dist_mode=dist_mode, ablated_rho=ablated_rho,
logits=ablated_logits, )
rho=ablated_rho, sums, sumsq, counts = grouped_parameter_stats(
) disease_values,
if disease_distribution_name_by_group[group] != disease_distribution: group_token_mask,
raise RuntimeError( )
"Disease distribution changed between original and ablated passes " disease_stat_chunks.append((key_table, sums, sumsq, counts))
f"for group {group!r}: {disease_distribution_name_by_group[group]!r} "
f"vs {disease_distribution!r}" death_summary.clear()
) for key_rows, values in zip(death_key_chunks, death_value_chunks):
value_block = disease_parameter_pair_block( update_death_summary(
original_disease_params_by_group[group][row_tensor], death_summary,
ablated_disease_params, key_rows=key_rows,
).detach().cpu().numpy() values=values,
update_disease_parameter_summary( )
disease_parameter_summary,
key_rows=key_table, for key_rows, sums, sumsq, counts in disease_stat_chunks:
target_group=group, update_disease_parameter_summary_from_group_stats(
target_group_label=risk_group_labels[group], disease_parameter_summary,
values=value_block, key_rows=key_rows,
) group_names=group_names,
group_labels=group_labels,
sums=sums,
sumsq=sumsq,
counts=counts,
)
death_summary_path = output_dir / "summary_extra_info_death_parameters.csv" death_summary_path = output_dir / "summary_extra_info_death_parameters.csv"
disease_summary_path = output_dir / "summary_extra_info_disease_parameters.csv" disease_summary_path = output_dir / "summary_extra_info_disease_parameters.csv"