Spaces:
Sleeping
Sleeping
File size: 4,185 Bytes
ce7c1f0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 | """Unified explainability helpers for the demo app and README §4.
This module intentionally sits above the task-specific `src.classify.explain`
implementation so the app can import one stable surface area for:
* LR token-level explanations
* lightweight BERT attention visualizations
"""
from __future__ import annotations
import json
import pickle
from pathlib import Path
from typing import Any
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from src.classify.explain import explain_lr
from src.utils import MODELS_DIR, RESULTS_DIR, get_logger
logger = get_logger(__name__)
def load_pickle_model(path: str | Path):
"""Load a pickled sklearn pipeline."""
with Path(path).open("rb") as fh:
return pickle.load(fh)
def explain_lr_model(
*,
pipeline_path: str | Path,
texts: list[str],
label_names: list[str],
task: str,
top_k: int = 15,
output_dir: str | Path | None = None,
) -> Path:
"""Render a SHAP explanation plot for the supplied texts."""
pipe = load_pickle_model(pipeline_path)
return explain_lr(
pipe,
np.asarray(texts, dtype=object),
label_names,
task=task,
top_k=top_k,
output_dir=Path(output_dir) if output_dir else None,
)
def render_attention_heatmap(
tokens: list[str],
attention: np.ndarray,
*,
title: str,
output_path: str | Path,
max_tokens: int = 20,
) -> Path:
"""Persist a compact attention heatmap for one BERT example."""
output = Path(output_path)
output.parent.mkdir(parents=True, exist_ok=True)
usable_tokens = tokens[:max_tokens]
usable_attention = np.asarray(attention, dtype=float)[: len(usable_tokens), : len(usable_tokens)]
fig, ax = plt.subplots(figsize=(max(6, len(usable_tokens) * 0.45), max(5, len(usable_tokens) * 0.45)))
im = ax.imshow(usable_attention, cmap="magma")
ax.set_xticks(range(len(usable_tokens)))
ax.set_yticks(range(len(usable_tokens)))
ax.set_xticklabels(usable_tokens, rotation=60, ha="right", fontsize=8)
ax.set_yticklabels(usable_tokens, fontsize=8)
ax.set_title(title)
fig.colorbar(im, ax=ax, fraction=0.04, pad=0.03)
fig.tight_layout()
fig.savefig(output, dpi=180)
plt.close(fig)
logger.info("Saved BERT attention heatmap → %s", output)
return output
def export_bert_attention_example(
*,
model_dir: str | Path,
text: str,
output_path: str | Path,
max_length: int = 128,
) -> dict[str, Any]:
"""Compute a representative attention map for one text and save it as an image."""
try:
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer
except ImportError as exc:
raise ImportError("Install transformers and torch for BERT attention export.") from exc
model_path = Path(model_dir)
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForSequenceClassification.from_pretrained(model_path, output_attentions=True)
model.eval()
inputs = tokenizer(
text,
return_tensors="pt",
truncation=True,
max_length=max_length,
)
with torch.no_grad():
outputs = model(**inputs)
tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
attentions = outputs.attentions
if not attentions:
raise RuntimeError("Model did not return attentions.")
last_layer = attentions[-1][0].mean(dim=0).cpu().numpy()
image_path = render_attention_heatmap(
tokens,
last_layer,
title="BERT attention example",
output_path=output_path,
)
return {
"image_path": str(image_path),
"tokens": tokens[:20],
}
def load_label_names(model_dir: str | Path) -> list[str]:
"""Read label names for a saved BERT model directory."""
path = Path(model_dir) / "label_names.json"
return json.loads(path.read_text(encoding="utf-8"))
__all__ = [
"RESULTS_DIR",
"MODELS_DIR",
"explain_lr_model",
"export_bert_attention_example",
"load_label_names",
"load_pickle_model",
"render_attention_heatmap",
]
|