#!/usr/bin/env python3 """MAS case study: visualize token-perturbation faithfulness for attribution methods. This script matches the faithfulness evaluation logic implemented in: - evaluations/faithfulness.py - llm_attr_eval.LLMAttributionEvaluator.faithfulness_test() For a single example and a selected attribution method, we: 1) Compute token-level attributions (Seq / Row / Recursive) over prompt tokens. 2) Rank prompt tokens by attribution mass. 3) Iteratively perturb the prompt by replacing one token at a time with PAD tokens. 4) Score the model as sum log p(generation + EOS | prompt) under the chat template. 5) Compute RISE / MAS / RISE+AP (AUCs) and visualize the perturbation impact as token heatmaps. Outputs JSON + HTML to exp/case_study/out/. """ from __future__ import annotations import argparse import json import os import sys import types from importlib.machinery import ModuleSpec from pathlib import Path from typing import Any, Dict, List, Optional, Sequence, Tuple import numpy as np def _early_set_cuda_visible_devices() -> None: """Set CUDA_VISIBLE_DEVICES before importing torch/transformers. Note: CUDA device indices are re-mapped inside the process after applying the mask. """ parser = argparse.ArgumentParser(add_help=False) parser.add_argument("--cuda", type=str, default=None) args, _ = parser.parse_known_args(sys.argv[1:]) cuda = args.cuda.strip() if isinstance(args.cuda, str) else "" if cuda and "," in cuda: os.environ["CUDA_VISIBLE_DEVICES"] = cuda if __name__ == "__main__": _early_set_cuda_visible_devices() import torch REPO_ROOT = Path(__file__).resolve().parents[2] if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) # Avoid optional vision deps when importing transformers. os.environ.setdefault("TRANSFORMERS_NO_TORCHVISION", "1") os.environ.setdefault("DISABLE_TRANSFORMERS_IMAGE_TRANSFORMS", "1") def _stub_torchvision() -> None: """Provide minimal torchvision stubs so transformers imports succeed without torchvision.""" if "torchvision" in sys.modules: return def _mk(name: str) -> types.ModuleType: mod = types.ModuleType(name) mod.__spec__ = ModuleSpec(name, loader=None) return mod tv = _mk("torchvision") tv.__dict__["__path__"] = [] submods = ["transforms", "_meta_registrations", "datasets", "io", "models", "ops", "utils"] for name in submods: mod = _mk(f"torchvision.{name}") sys.modules[f"torchvision.{name}"] = mod setattr(tv, name, mod) class _InterpolationMode: NEAREST = 0 NEAREST_EXACT = 0 BILINEAR = 1 BICUBIC = 2 LANCZOS = 3 BOX = 4 HAMMING = 5 sys.modules["torchvision.transforms"].InterpolationMode = _InterpolationMode sys.modules["torchvision.transforms"].__all__ = ["InterpolationMode"] ops_mod = sys.modules.get("torchvision.ops") or _mk("torchvision.ops") sys.modules["torchvision.ops"] = ops_mod setattr(tv, "ops", ops_mod) misc_mod = _mk("torchvision.ops.misc") sys.modules["torchvision.ops.misc"] = misc_mod setattr(ops_mod, "misc", misc_mod) class _FrozenBatchNorm2d: def __init__(self, *args, **kwargs): pass misc_mod.FrozenBatchNorm2d = _FrozenBatchNorm2d sys.modules["torchvision"] = tv def _stub_timm() -> None: """Provide minimal timm stubs to avoid optional vision deps.""" if "timm" in sys.modules: return def _mk(name: str) -> types.ModuleType: mod = types.ModuleType(name) mod.__spec__ = ModuleSpec(name, loader=None) return mod timm = _mk("timm") timm.__dict__["__path__"] = [] sys.modules["timm"] = timm data_mod = _mk("timm.data") sys.modules["timm.data"] = data_mod timm.data = data_mod class _ImageNetInfo: pass def _infer_imagenet_subset(*args, **kwargs): return None data_mod.ImageNetInfo = _ImageNetInfo data_mod.infer_imagenet_subset = _infer_imagenet_subset layers_mod = _mk("timm.layers") sys.modules["timm.layers"] = layers_mod timm.layers = layers_mod create_norm_mod = _mk("timm.layers.create_norm") sys.modules["timm.layers.create_norm"] = create_norm_mod layers_mod.create_norm = create_norm_mod def _get_norm_layer(*args, **kwargs): return None create_norm_mod.get_norm_layer = _get_norm_layer classifier_mod = _mk("timm.layers.classifier") sys.modules["timm.layers.classifier"] = classifier_mod layers_mod.classifier = classifier_mod def _stub_gemma3n() -> None: """Stub Gemma3n config module if transformers tries to import it.""" if "transformers.models.gemma3n.configuration_gemma3n" in sys.modules: return gemma_pkg = types.ModuleType("transformers.models.gemma3n") gemma_pkg.__spec__ = ModuleSpec("transformers.models.gemma3n", loader=None, is_package=True) sys.modules["transformers.models.gemma3n"] = gemma_pkg gemma_conf = types.ModuleType("transformers.models.gemma3n.configuration_gemma3n") gemma_conf.__spec__ = ModuleSpec("transformers.models.gemma3n.configuration_gemma3n", loader=None) class Gemma3nConfig: def __init__(self, *args, **kwargs): self.model_type = "gemma3n" class Gemma3nTextConfig(Gemma3nConfig): pass gemma_conf.Gemma3nConfig = Gemma3nConfig gemma_conf.Gemma3nTextConfig = Gemma3nTextConfig gemma_conf.__all__ = ["Gemma3nConfig", "Gemma3nTextConfig"] sys.modules["transformers.models.gemma3n.configuration_gemma3n"] = gemma_conf setattr(gemma_pkg, "configuration_gemma3n", gemma_conf) _stub_torchvision() _stub_timm() _stub_gemma3n() import transformers # noqa: E402 # Provide light stubs if Longformer classes are unavailable; we don't use them here. if not hasattr(transformers, "LongformerTokenizer"): class _DummyLongformerTokenizer: def __init__(self, *args, **kwargs): raise ImportError("LongformerTokenizer stubbed; install full transformers if needed.") transformers.LongformerTokenizer = _DummyLongformerTokenizer if not hasattr(transformers, "LongformerForMaskedLM"): class _DummyLongformerForMaskedLM: def __init__(self, *args, **kwargs): raise ImportError("LongformerForMaskedLM stubbed; install full transformers if needed.") transformers.LongformerForMaskedLM = _DummyLongformerForMaskedLM from exp.case_study import viz # noqa: E402 from exp.exp2 import dataset_utils as ds_utils # noqa: E402 from shared_utils import DEFAULT_PROMPT_TEMPLATE # noqa: E402 import llm_attr # noqa: E402 from evaluations.attribution_recovery import load_model # noqa: E402 def resolve_device(cuda: Optional[str], cuda_num: int) -> str: if cuda and isinstance(cuda, str) and "," in cuda: os.environ["CUDA_VISIBLE_DEVICES"] = cuda return "auto" if cuda and isinstance(cuda, str) and cuda.strip(): try: idx = int(cuda) except Exception: idx = 0 return f"cuda:{idx}" if torch.cuda.is_available() else "cpu" return f"cuda:{cuda_num}" if torch.cuda.is_available() else "cpu" def load_example(dataset: str, index: int, data_root: Path) -> Tuple[ds_utils.CachedExample, str]: ds_path = Path(dataset) if ds_path.exists(): examples = ds_utils.read_cached_jsonl(ds_path) dataset_name = ds_path.name else: loader = ds_utils.DatasetLoader(data_root=data_root) examples = loader.load(dataset) dataset_name = dataset if not examples: raise ValueError(f"No examples found for dataset={dataset}") if index < 0: index = len(examples) + index if not (0 <= index < len(examples)): raise IndexError(f"index {index} out of range for dataset with {len(examples)} examples") return examples[index], dataset_name def make_output_stem(dataset_name: str, index: int, method: str) -> str: safe_name = dataset_name.replace("/", "_").replace(" ", "_") return f"mas_case_{method}_{safe_name}_idx{index}" def format_prompt(tokenizer: Any, prompt: str) -> str: modified_prompt = DEFAULT_PROMPT_TEMPLATE.format(context=prompt, query="") formatted_prompt = [{"role": "user", "content": modified_prompt}] return tokenizer.apply_chat_template( formatted_prompt, tokenize=False, add_generation_prompt=True, enable_thinking=False, ) @torch.inference_mode() def compute_logprob_response_given_prompt(model: Any, prompt_ids: torch.Tensor, response_ids: torch.Tensor) -> torch.Tensor: """Compute log-probabilities of response_ids given prompt_ids. Shapes: prompt_ids: [B, N] response_ids: [B, M] returns: [B, M] """ input_ids = torch.cat([prompt_ids, response_ids], dim=1) attention_mask = torch.ones_like(input_ids) logits = model(input_ids=input_ids, attention_mask=attention_mask).logits # [B, N+M, V] log_probs = torch.nn.functional.log_softmax(logits, dim=-1) response_start = int(prompt_ids.shape[1]) logits_for_response = log_probs[:, response_start - 1 : -1, :] # [B, M, V] gathered = logits_for_response.gather(2, response_ids.unsqueeze(-1)) return gathered.squeeze(-1) @torch.inference_mode() def score_prompt_ids_with_generation(model: Any, *, prompt_ids: torch.Tensor, generation_ids: torch.Tensor) -> float: return float(compute_logprob_response_given_prompt(model, prompt_ids, generation_ids).sum().detach().cpu().item()) @torch.inference_mode() def _ensure_pad_token_id(tokenizer: Any) -> int: if tokenizer.pad_token_id is None: if tokenizer.eos_token_id is None: raise RuntimeError("tokenizer has neither pad_token_id nor eos_token_id; cannot define baseline token.") tokenizer.pad_token = tokenizer.eos_token return int(tokenizer.pad_token_id) def _find_subsequence_start(haystack: torch.Tensor, needle: torch.Tensor) -> Optional[int]: if haystack.ndim != 1 or needle.ndim != 1: raise ValueError("Expected 1D tensors for subsequence matching.") if needle.numel() == 0: return 0 hay_len = int(haystack.numel()) needle_len = int(needle.numel()) if needle_len > hay_len: return None for i in range(hay_len - needle_len + 1): if torch.equal(haystack[i : i + needle_len], needle): return i return None def decode_text_into_tokens(tokenizer: Any, text: str) -> List[str]: encoding = tokenizer(text, return_offsets_mapping=True, add_special_tokens=False) offsets = list(encoding["offset_mapping"]) tokens: List[str] = [] for start, end in offsets: tokens.append(text[start:end]) return tokens def auc(arr: np.ndarray) -> float: return float((arr.sum() - arr[0] / 2 - arr[-1] / 2) / (arr.shape[0] - 1)) def mas_trace( model: Any, tokenizer: Any, *, attribution: torch.Tensor, prompt: str, generation: str, user_prompt_indices: Optional[Sequence[int]] = None, keep_prompt_token_indices: Optional[Sequence[int]] = None, k: int = 20, ) -> Dict[str, Any]: """Return a token-level faithfulness trace (RISE/MAS/RISE+AP) plus per-token deltas.""" pad_token_id = _ensure_pad_token_id(tokenizer) user_prompt = " " + prompt formatted = format_prompt(tokenizer, user_prompt) formatted_ids = tokenizer(formatted, return_tensors="pt", add_special_tokens=False).input_ids user_ids = tokenizer(user_prompt, return_tensors="pt", add_special_tokens=False).input_ids prompt_ids = formatted_ids.to(model.device) prompt_ids_perturbed = prompt_ids.clone() gen_ids = tokenizer( generation + (tokenizer.eos_token or ""), return_tensors="pt", add_special_tokens=False, ).input_ids.to(model.device) attr_cpu = attribution.detach().cpu() w = attr_cpu.sum(0) P = int(w.numel()) if keep_prompt_token_indices is None: keep = list(range(P)) else: keep = [] seen: set[int] = set() for raw in keep_prompt_token_indices: try: idx = int(raw) except Exception: continue if 0 <= idx < P and idx not in seen: keep.append(idx) seen.add(idx) keep.sort() K = len(keep) if K: w_keep = w.index_select(0, torch.as_tensor(keep, dtype=torch.long)) sorted_local = torch.argsort(w_keep, descending=True) sorted_attr_indices = torch.as_tensor([keep[int(i.item())] for i in sorted_local], dtype=torch.long) attr_sum = float(w_keep.sum().item()) else: sorted_attr_indices = torch.zeros((0,), dtype=torch.long) attr_sum = 0.0 if int(user_ids.shape[1]) != P: raise ValueError( "Prompt-side attribution length does not match tokenized user prompt length: " f"attr P={P}, user_prompt P={int(user_ids.shape[1])}." ) prompt_positions: List[int] if user_prompt_indices is not None: prompt_positions = [int(x) for x in user_prompt_indices] if len(prompt_positions) != P: raise ValueError( "user_prompt_indices length does not match prompt-side attribution length: " f"indices P={len(prompt_positions)}, attr P={P}." ) if P and max(prompt_positions) >= int(prompt_ids_perturbed.shape[1]): raise ValueError("user_prompt_indices contains an out-of-bounds index for formatted prompt ids.") else: user_start = _find_subsequence_start(formatted_ids[0], user_ids[0]) if user_start is None: raise RuntimeError("Failed to locate user prompt token span inside formatted chat prompt.") prompt_positions = [int(user_start) + j for j in range(P)] if K > 0: steps = int(k) if k is not None else 0 if steps <= 0: steps = 1 steps = min(steps, K) else: steps = 0 scores = np.zeros(steps + 1, dtype=np.float64) density = np.zeros(steps + 1, dtype=np.float64) scores[0] = score_prompt_ids_with_generation(model, prompt_ids=prompt_ids_perturbed, generation_ids=gen_ids) density[0] = 1.0 if K == 0: return { "num_tokens": P, "sorted_attr_indices": [], "scores_raw": scores.tolist(), "density": density.tolist(), "normalized_model_response": scores.tolist(), "alignment_penalty": np.zeros_like(scores).tolist(), "corrected_scores": scores.tolist(), "token_deltas_raw": np.zeros(P, dtype=np.float64).tolist(), "attr_weights": np.zeros(P, dtype=np.float64).tolist(), "metrics": {"RISE": 0.0, "MAS": 0.0, "RISE+AP": 0.0}, } if attr_sum <= 0: density = np.linspace(1.0, 0.0, steps + 1) per_token_delta = np.zeros(P, dtype=np.float64) base = K // steps remainder = K % steps start = 0 for step in range(steps): size = base + (1 if step < remainder else 0) group = sorted_attr_indices[start : start + size] start += size for idx_t in group: idx = int(idx_t.item()) abs_pos = int(prompt_positions[idx]) prompt_ids_perturbed[0, abs_pos] = pad_token_id scores[step + 1] = score_prompt_ids_with_generation(model, prompt_ids=prompt_ids_perturbed, generation_ids=gen_ids) if attr_sum > 0: dec = float(w.index_select(0, group).sum().item()) / attr_sum density[step + 1] = density[step] - dec delta = scores[step] - scores[step + 1] for idx_t in group: idx = int(idx_t.item()) per_token_delta[idx] = delta min_normalized_pred = 1.0 normalized_model_response = scores.copy() for i in range(len(scores)): normalized_pred = (normalized_model_response[i] - scores[-1]) / (abs(scores[0] - scores[-1])) normalized_pred = np.clip(normalized_pred, 0.0, 1.0) min_normalized_pred = min(min_normalized_pred, float(normalized_pred)) normalized_model_response[i] = min_normalized_pred alignment_penalty = np.abs(normalized_model_response - density) corrected_scores = normalized_model_response + alignment_penalty corrected_scores = corrected_scores.clip(0, 1) corrected_scores = (corrected_scores - np.min(corrected_scores)) / (np.max(corrected_scores) - np.min(corrected_scores)) if np.isnan(corrected_scores).any(): corrected_scores = np.linspace(1, 0, len(scores)) rise = auc(normalized_model_response) mas = auc(corrected_scores) rise_ap = auc(normalized_model_response + alignment_penalty) if attr_sum > 0: attr_weights = np.zeros(P, dtype=np.float64) for idx in keep: attr_weights[idx] = float(w[idx].item()) / (attr_sum + 1e-12) else: attr_weights = np.zeros(P, dtype=np.float64) return { "num_tokens": P, "sorted_attr_indices": [int(i.item()) for i in sorted_attr_indices], "scores_raw": scores.tolist(), "density": density.tolist(), "normalized_model_response": normalized_model_response.tolist(), "alignment_penalty": alignment_penalty.tolist(), "corrected_scores": corrected_scores.tolist(), "token_deltas_raw": per_token_delta.tolist(), "attr_weights": attr_weights.tolist(), "metrics": {"RISE": rise, "MAS": mas, "RISE+AP": rise_ap}, } def compute_method_attribution( method: str, example: ds_utils.CachedExample, model: Any, tokenizer: Any, *, n_hops: int, sink_span: Optional[Tuple[int, int]], thinking_span: Optional[Tuple[int, int]], chunk_tokens: int, sink_chunk_tokens: int, attnlrp_neg_handling: str, attnlrp_norm_mode: str, ) -> Tuple[str, Any, llm_attr.LLMAttributionResult]: prompt = example.prompt target = example.target if method == "ifr": if sink_span is None: raise ValueError("IFR requires sink_span (use dataset sink_span or pass --sink_span).") attributor = llm_attr.LLMIFRAttribution(model, tokenizer, chunk_tokens=chunk_tokens, sink_chunk_tokens=sink_chunk_tokens) result = attributor.calculate_ifr_span(prompt, target=target, span=sink_span) return "IFR (ifr_span)", attributor, result if method == "ifr_all_positions_output_only": if sink_span is None: raise ValueError( "ifr_all_positions_output_only requires sink_span (use dataset sink_span or pass --sink_span)." ) attributor = llm_attr.LLMIFRAttribution(model, tokenizer, chunk_tokens=chunk_tokens, sink_chunk_tokens=sink_chunk_tokens) result = attributor.calculate_ifr_for_all_positions_output_only( prompt, target=target, sink_span=sink_span, ) return "IFR (ifr_all_positions_output_only)", attributor, result if method in ("ft", "ft_ifr"): attributor = llm_attr.LLMIFRAttribution(model, tokenizer, chunk_tokens=chunk_tokens, sink_chunk_tokens=sink_chunk_tokens) result = attributor.calculate_ifr_multi_hop( prompt, target=target, sink_span=sink_span, thinking_span=thinking_span, n_hops=int(n_hops), ) return "FT-IFR (ifr_multi_hop)", attributor, result if method in ("ft_improve", "ft_ifr_improve"): import ft_ifr_improve attributor = ft_ifr_improve.LLMIFRAttributionImproved( model, tokenizer, chunk_tokens=chunk_tokens, sink_chunk_tokens=sink_chunk_tokens, ) result = attributor.calculate_ifr_multi_hop_stop_words( prompt, target=target, sink_span=sink_span, thinking_span=thinking_span, n_hops=int(n_hops), ) return "FT-IFR (ifr_multi_hop_stop_words)", attributor, result if method == "ft_split_hop": import ft_ifr_improve attributor = ft_ifr_improve.LLMIFRAttributionSplitHop( model, tokenizer, chunk_tokens=chunk_tokens, sink_chunk_tokens=sink_chunk_tokens, ) result = attributor.calculate_ifr_multi_hop_split_hop( prompt, target=target, sink_span=sink_span, thinking_span=thinking_span, n_hops=int(n_hops), ) return "FT-IFR (ifr_multi_hop_split_hop)", attributor, result if method == "attnlrp": attributor = llm_attr.LLMLRPAttribution(model, tokenizer) result = attributor.calculate_attnlrp_ft_hop0( prompt, target=target, sink_span=sink_span, thinking_span=thinking_span, neg_handling=attnlrp_neg_handling, norm_mode=attnlrp_norm_mode, ) return "AttnLRP (ft_attnlrp hop0)", attributor, result if method == "ft_attnlrp": attributor = llm_attr.LLMLRPAttribution(model, tokenizer) result = attributor.calculate_attnlrp_aggregated_multi_hop( prompt, target=target, sink_span=sink_span, thinking_span=thinking_span, n_hops=int(n_hops), neg_handling=attnlrp_neg_handling, norm_mode=attnlrp_norm_mode, ) return "FT-AttnLRP (attnlrp_aggregated_multi_hop)", attributor, result raise ValueError(f"Unsupported method={method!r}") def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser("MAS case study (faithfulness perturbation visualization)") parser.add_argument("--dataset", type=str, default="exp/exp2/data/morehopqa.jsonl", help="Dataset name or JSONL path.") parser.add_argument("--data_root", type=str, default="exp/exp2/data", help="Cache root for dataset names.") parser.add_argument("--index", type=int, default=0, help="Sample index (supports negative for reverse).") parser.add_argument( "--method", type=str, choices=[ "ifr", "ifr_all_positions_output_only", "ft", "ft_ifr", "ft_improve", "ft_ifr_improve", "ft_split_hop", "attnlrp", "ft_attnlrp", ], default="ft", ) parser.add_argument("--model", type=str, default="qwen-8B", help="HF repo id (ignored if --model_path set).") parser.add_argument("--model_path", type=str, default=None, help="Local model path to override --model.") parser.add_argument("--cuda", type=str, default=None, help="CUDA spec (e.g., '0' or '0,1').") parser.add_argument("--cuda_num", type=int, default=0, help="Fallback GPU index when --cuda unset.") parser.add_argument("--n_hops", type=int, default=1, help="Number of hops for multi-hop methods.") parser.add_argument("--sink_span", type=int, nargs=2, default=None, help="Optional sink span over generation tokens.") parser.add_argument("--thinking_span", type=int, nargs=2, default=None, help="Optional thinking span over generation tokens.") parser.add_argument("--chunk_tokens", type=int, default=128, help="IFR chunk size.") parser.add_argument("--sink_chunk_tokens", type=int, default=32, help="IFR sink chunk size.") parser.add_argument( "--attnlrp_neg_handling", type=str, choices=["drop", "abs"], default="drop", help="FT-AttnLRP: how to handle negative values after each hop (drop=clamp>=0, abs=absolute value).", ) parser.add_argument( "--attnlrp_norm_mode", type=str, choices=["norm", "no_norm"], default="norm", help="FT-AttnLRP: norm enables per-hop global+thinking normalization + ratios; no_norm disables all three.", ) parser.add_argument("--output_dir", type=str, default="exp/case_study/out", help="Where to write HTML/JSON artifacts.") return parser.parse_args() def main() -> None: args = parse_args() device = resolve_device(args.cuda, args.cuda_num) if torch.cuda.is_available(): visible = os.environ.get("CUDA_VISIBLE_DEVICES") print(f"[info] CUDA_VISIBLE_DEVICES={visible!r} torch.cuda.device_count()={torch.cuda.device_count()} device={device}") if args.method == "ft_ifr": method_key = "ft" elif args.method == "ft_ifr_improve": method_key = "ft_improve" else: method_key = args.method model_name = args.model_path if args.model_path is not None else args.model model, tokenizer = load_model(model_name, device) example, ds_name = load_example(args.dataset, args.index, Path(args.data_root)) sink_span = tuple(args.sink_span) if args.sink_span is not None else tuple(example.sink_span) if example.sink_span else None thinking_span = ( tuple(args.thinking_span) if args.thinking_span is not None else tuple(example.thinking_span) if example.thinking_span else None ) method_label, attributor, attr_result = compute_method_attribution( method_key, example, model, tokenizer, n_hops=args.n_hops, sink_span=sink_span, thinking_span=thinking_span, chunk_tokens=args.chunk_tokens, sink_chunk_tokens=args.sink_chunk_tokens, attnlrp_neg_handling=args.attnlrp_neg_handling, attnlrp_norm_mode=args.attnlrp_norm_mode, ) indices_to_explain = example.indices_to_explain or example.sink_span if not (isinstance(indices_to_explain, list) and len(indices_to_explain) == 2): raise ValueError("MAS case study requires token-span indices_to_explain=[start_tok,end_tok] (e.g. sink_span).") seq_attr, row_attr, rec_attr = attr_result.get_all_token_attrs(indices_to_explain) prompt_tokens = decode_text_into_tokens(tokenizer, " " + example.prompt) generation_text = example.target if example.target is not None else (getattr(attributor, "generation", None) or "") variant_specs = [ ("seq", "Seq attribution", seq_attr), ("row", "Row attribution", row_attr), ("recursive", "Recursive attribution", rec_attr), ] formatted = format_prompt(tokenizer, " " + example.prompt) prompt_ids = tokenizer(formatted, return_tensors="pt", add_special_tokens=False).input_ids.to(model.device) gen_ids = tokenizer( generation_text + (tokenizer.eos_token or ""), return_tensors="pt", add_special_tokens=False, ).input_ids.to(model.device) base_score = score_prompt_ids_with_generation(model, prompt_ids=prompt_ids, generation_ids=gen_ids) panels_raw: List[Dict[str, Any]] = [] panels_display: List[Dict[str, Any]] = [] for variant_key, variant_label, variant_attr in variant_specs: prompt_len = int(seq_attr.shape[1] - seq_attr.shape[0]) # cols=(P+G), rows=G attr_prompt = variant_attr[:, :prompt_len] keep_prompt_token_indices = None if method_key == "ft_improve": import ft_ifr_improve keep_prompt_token_indices = ft_ifr_improve.keep_token_indices(list(getattr(attributor, "user_prompt_tokens", []) or [])) trace = mas_trace( model, tokenizer, attribution=attr_prompt.to(device="cpu"), prompt=example.prompt, generation=generation_text, user_prompt_indices=getattr(attributor, "user_prompt_indices", None), keep_prompt_token_indices=keep_prompt_token_indices, ) trace["variant"] = variant_key trace["variant_label"] = variant_label panel_raw = { "variant": variant_key, "variant_label": variant_label, "metrics": trace.get("metrics"), "sorted_attr_indices": trace.get("sorted_attr_indices"), "attr_weights": trace.get("attr_weights"), "token_deltas_raw": trace.get("token_deltas_raw"), "mas_trace": trace, } panels_raw.append(panel_raw) panel_display = { "variant": variant_key, "variant_label": variant_label, "metrics": trace.get("metrics"), "sorted_attr_indices": trace.get("sorted_attr_indices"), "attr_weights": trace.get("attr_weights"), "token_deltas_raw": trace.get("token_deltas_raw"), } panels_display.append(panel_display) case_meta: Dict[str, Any] = { "dataset": ds_name, "index": args.index, "mode": "mas", "attr_method": method_key, "attr_method_label": method_label, "sink_span": sink_span, "thinking_span": thinking_span, "n_hops": int(args.n_hops), "attnlrp_neg_handling": args.attnlrp_neg_handling if method_key in ("attnlrp", "ft_attnlrp") else None, "attnlrp_norm_mode": args.attnlrp_norm_mode if method_key in ("attnlrp", "ft_attnlrp") else None, "attnlrp_ratio_enabled": (args.attnlrp_norm_mode == "norm") if method_key in ("attnlrp", "ft_attnlrp") else None, "base_score": float(base_score), } record = { "meta": case_meta, "prompt": example.prompt, "target": example.target, "generation": generation_text, "prompt_tokens": prompt_tokens, "panels": panels_raw, } out_dir = Path(args.output_dir) out_dir.mkdir(parents=True, exist_ok=True) stem = make_output_stem(ds_name, args.index, method_key) json_path = out_dir / f"{stem}.json" html_path = out_dir / f"{stem}.html" with json_path.open("w", encoding="utf-8") as f: json.dump(record, f, ensure_ascii=False, indent=2) html = viz.render_mas_token_html( case_meta, prompt_tokens=prompt_tokens, panels=panels_display, generation=generation_text, ) html_path.write_text(html, encoding="utf-8") print(f"[done] wrote {json_path}") print(f"[done] wrote {html_path}") if __name__ == "__main__": main()