diff --git a/.gitignore b/.gitignore index 641786cb8..f04fb7a8f 100644 --- a/.gitignore +++ b/.gitignore @@ -200,4 +200,4 @@ docs/source/generated/ docs/source/gen_modules/ # Local debug / scratch scripts (not public) -scripts/ +scripts/ \ No newline at end of file diff --git a/examples/language/plot_text_imputer.py b/examples/language/plot_text_imputer.py new file mode 100644 index 000000000..a36466e5f --- /dev/null +++ b/examples/language/plot_text_imputer.py @@ -0,0 +1,130 @@ +""" +Explaining Text with the TextImputer +==================================== + +This example shows how to explain a sentiment classifier with the +:class:`~shapiq.imputer.text.imputer.TextImputer`. Each word in the input text +becomes a player in a cooperative game, and the model prediction is evaluated +under masked word coalitions. + +The first run downloads the Hugging Face model. The word-level player strategy +uses NLTK tokenization; if needed, install the resource once with +``uv run python -m nltk.downloader punkt_tab``. +""" + +from __future__ import annotations + +from itertools import combinations + +try: + from transformers import AutoModelForSequenceClassification, AutoTokenizer +except ImportError as err: + from shapiq.imputer.text._error import _text_import_error + + raise _text_import_error from err + +import nltk + + +def ensure_nltk_resource(resource: str, download_name: str) -> None: + try: + nltk.data.find(resource) + + except LookupError: + nltk.download(download_name, quiet=True) + + +ensure_nltk_resource("tokenizers/punkt_tab", "punkt_tab") + +from shapiq.approximator import KernelSHAPIQ +from shapiq.imputer.text.imputer import TextImputer + +MODEL_NAME = "distilbert-base-uncased-finetuned-sst-2-english" +TEXT = "The movie is not bad." +EXPLANATION_BUDGET = 64 + +# %% +# Load the Classifier +# ------------------- +# We use a DistilBERT sentiment classifier from Hugging Face and keep the +# example on CPU so that it runs consistently across machines. + +device = "cpu" +tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) +model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME) +model.to(device).eval() + +# %% +# Build the Text Imputer +# ---------------------- +# The imputer masks absent words and returns the normalized positive-class +# probability. With normalization, the empty coalition becomes the baseline +# and the grand coalition is the full-text prediction minus that baseline. + +imputer = TextImputer( + model=model, + tokenizer=tokenizer, + text=TEXT, + model_type="encoder_classifier", + player_level="word", + perturbation_type="mask", + class_idx=1, + output_type="probability", + device=device, +) + +words = imputer.player_strategy.get_players() +normalized_score = float(imputer(imputer.grand_coalition)[0]) +print(f"Input text : {TEXT}") +print(f"Word players : {words}") +print(f"Normalized score : {normalized_score:.4f}") +print("Normalized score = full prediction - empty prediction") + +# %% +# Estimate Interactions +# --------------------- +# We estimate first- and second-order interactions with +# :class:`~shapiq.approximator.KernelSHAPIQ`. + +interaction_values = KernelSHAPIQ( + n=imputer.n_features, + index="k-SII", + max_order=2, + random_state=42, +).approximate( + budget=EXPLANATION_BUDGET, + game=imputer, +) + +# %% +# First-Order Effects +# ------------------- +# The first-order values show how each individual word contributes to the +# normalized sentiment prediction. + +print("First-order interaction values") +print(f"{'idx':>3} {'word':<18} {'value':>12}") +print("-" * 40) +for idx, word in enumerate(words): + print(f"{idx:>3} {word:<18} {interaction_values[(idx,)]:>+12.4f}") + +# %% +# Pairwise Effects +# ---------------- +# The second-order values show how pairs of words interact. For example, +# negations often interact strongly with the words they modify. + +print("Second-order interaction values") +print(f"{'pair':<31} {'value':>12}") +print("-" * 47) +for i, j in combinations(range(imputer.n_features), 2): + pair_name = f"{words[i]} x {words[j]}" + print(f"{pair_name:<31} {interaction_values[(i, j)]:>+12.4f}") + +# %% +# Force Plot +# ---------- +# Finally, a force plot visualizes how the word-level effects move the +# prediction away from the baseline. + +interaction_values.plot_force(feature_names=words) diff --git a/project_text_shapiq.md b/project_text_shapiq.md new file mode 100644 index 000000000..3cb972b3f --- /dev/null +++ b/project_text_shapiq.md @@ -0,0 +1,113 @@ +# Project: Shapley Interactions on LLMs — Text Imputers & Cool Use Cases + +**Type:** Pull Request(s) + Demo + +## Overview + +Language models are everywhere — encoder classifiers, causal LLMs for chat and code, seq2seq models for translation and summarization — and for every one of them, there is fascinating behavior that nobody fully understands. Where does a jailbreak actually "act"? Which few-shot demonstration did an in-context answer rely on? Which retrieved chunk is the RAG answer really grounded in? Which words in a prompt *interact* to flip a sentiment prediction? Shapley values and **Shapley interactions** are uniquely well-suited to answering this class of question, and shapiq has the full game-theoretic machinery — any interaction order, many indices (SII, k-SII, STII, FSII, FBII, BV, BII), many approximators — to go beyond what every other SHAP library offers. + +This project has two sides: **building a text imputer** for shapiq (the PR side) and **exploring interesting use cases** that showcase what Shapley interactions can reveal about language model behavior (the Demo side). The core question driving both is: + +> *What can shapiq show about LLM behavior that nothing else can — and what's the best way to enable it?* + +The PR deliverable is a **text imputer** — a clean, tested component that makes text/LLM explanations a first-class capability in shapiq. Once you have that, you'll use it to build polished demos on interesting use cases (jailbreaks, RAG attribution, in-context learning, and more). Your team has freedom in choosing **which use cases to explore** and how deep to go — pick the demos that excite you most. + +The starting points are shapiq's existing text handling: a hard-coded sentiment-analysis benchmark game (`shapiq_games.benchmark.local_xai.SentimentAnalysis`) and a sentence visualization utility (`src/shapiq/plot/sentence.py`). + +## Tasks + +### Task 1: Build a text imputer for shapiq (the PR) + +This is the engineering heart of the project. Build a **text imputer** — a subclass of `shapiq.imputer.base.Imputer` that lets you mask player-defined spans of a prompt and call an arbitrary LLM on the masked versions. This is what you need to build: + +- A flexible **player definition** — at minimum token-level and word-level players; span-level players (sentences, retrieved chunks, few-shot demonstrations) for more advanced use cases. +- Pluggable **masking strategies** — `[MASK]` replacement, `[PAD]` replacement, token removal, attention masking, and ideally at least one novel strategy (e.g. MLM-infill, neutral replacement). +- A flexible **target callable** — classification logits for encoder models, next-token / target-continuation log-likelihood for causal LLMs, and optionally perplexity, contrastive log-odds, etc. +- **Batched model calls** — LLMs are expensive; the value function must evaluate batches of coalitions efficiently. +- Integration with **HuggingFace `transformers`** (torch backend is sufficient; JAX/Flax is a nice-to-have). + +Of course, if you want to go beyond the text imputer, you are welcome to contribute additional infrastructure as well — for example: + +- **Text/LLM games:** New game classes (subclassing `shapiq.game.Game`, or as new `shapiq_games` benchmark games) tailored to LLM explanation scenarios — e.g. a prompt-attribution game, a RAG-grounding game, a jailbreak-detection game. Look at the existing `SentimentAnalysis` benchmark game for the pattern. +- **Visualization utilities:** Extensions to `src/shapiq/plot/sentence.py` — token-interaction heatmaps, interaction-graph overlays on prompts, side-by-side comparison plots for different indices or models. +- **Anything else** that would make text/LLM explanations in shapiq better for future users. + +But the text imputer is the core deliverable. Whatever you build must include tests, docstrings, and pass pre-commit. The existing `SentimentAnalysis` benchmark game (`src/shapiq_games/benchmark/local_xai/benchmark_language.py`) and `src/shapiq/plot/sentence.py` are your starting points. + +### Task 2: The demos — showcase your text imputer on interesting use cases + +> **Note on scaffolding:** You will likely need some quick-and-dirty scaffolding code to unblock your demo exploration before the PR code is polished. That's expected — build the minimum viable version first, get to the interesting demos, and then decide what's clean enough to promote to the PR. Work with HuggingFace `transformers` (torch backend); prefer real SOTA models — smaller open-weight models (1–8B range) are fine if larger ones don't fit your hardware. + +Now that you have a text imputer, use it to build **at least three** polished, self-contained demos that showcase what Shapley interactions can reveal (or cannot reveal) about language model behavior. Here are some directions we find interesting — you are welcome to replace these with better ideas: + +- **Prompt-injection and jailbreaks.** Feed a model a benign prompt plus a jailbreak payload. Use Shapley interactions to show which token groups *interact* to cause the jailbreak — and which don't. Does the interaction structure distinguish real jailbreaks from innocuous phrasing? +- **In-context learning attribution.** For a few-shot prompt, attribute the answer back to individual demonstrations (treat each demonstration as one player). Which example did the model actually rely on? Are there interaction effects between demonstrations? +- **RAG / retrieval attribution.** For a retrieval-augmented answer, treat each retrieved chunk (or each sentence in each chunk) as a player. Identify the *source of groundedness* and flag unsupported answers where no chunk has a meaningful attribution. +- **Chain-of-thought attribution.** Attribute a final answer back to the model's own reasoning steps. Which step was load-bearing? Which was filler? +- **Contrastive / counterfactual explanations.** Given two almost-identical prompts with very different outputs (e.g. minor negation flip, pronoun swap), show the interactions responsible for the divergence. +- **Agentic tool-use explanations.** For a tool-calling agent, attribute the decision to call (or not call) a specific tool back to tokens in the user request and system prompt. +- **Multilinguality & robustness.** Compare attributions across translations or paraphrases of the same prompt — are explanations stable? Are the same interactions present? +- **Word-level interactions in classification.** On sentiment / NLI / toxicity tasks, use higher-order interactions to show phenomena first-order SVs miss: negation + adjective, subject-verb agreement, multi-word named entities, sarcasm cues. + +For each demo: + +1. **Frame the question clearly.** What are you trying to show? Why is it interesting? +2. **Pick a concrete SOTA model** (or a small set) and a concrete input (real jailbreak payloads from public datasets, real RAG traces, real few-shot prompts, etc.). Use genuinely interesting examples, not toy inputs. +3. **Use the right interaction index.** Shapley values alone are fine for simple cases, but the project's unique angle is interactions — use **k-SII, STII, or FSII** where pairwise or higher-order structure matters. Make a deliberate choice and explain it. +4. **Visualize the result.** Extend or reuse `src/shapiq/plot/sentence.py`; add heatmaps, interaction graphs, side-by-side comparisons — whatever makes the findings clearest. Visual quality matters. +5. **Draw a conclusion.** What did you learn? What's surprising? Where did the method break down? An honest "this didn't work and here's why" is a legitimate demo result. + +Format the demos as a collection of notebooks, a Gradio / Streamlit app, or a Hugging Face Space — whatever serves the content best. They should be fully reproducible (fixed seeds, pinned HF model revisions, clear install instructions) and runnable by anyone with a reasonable GPU. + +### Task 3: Comparison with existing libraries + +A demo that only shows what shapiq can do isn't enough — we also want to know how it compares. Include at least one comparison section (a notebook, a page, a dashboard) where you run shapiq head-to-head against existing text-explanation libraries on a shared input + model: + +- [`shap.Explainer` with `shap.maskers.Text`](https://shap.readthedocs.io/en/latest/generated/shap.maskers.Text.html) — the current de-facto Shapley-on-text baseline. +- [captum's `ShapleyValueSampling`](https://captum.ai/api/shapley_value_sampling.html) applied at the token-embedding level. +- [**Inseq**](https://inseq.org/) — a dedicated sequence-attribution library for generation; especially relevant for causal LLM comparisons. + +Report: (i) runtime and memory, (ii) agreement of attributions — do the same tokens come out as important? if not, why? (iii) API ergonomics, (iv) **what shapiq can uniquely do**: any-order interactions, many indices — show concrete examples where this buys something the baselines cannot offer. + +Honest discussion beats a cherry-picked win: where shapiq is slower or clunkier, say so. + +### Task 4: Additional PRs (optional) + +If your demo exploration surfaces additional pieces of reusable code beyond your main PR (Task 1), you are encouraged to upstream them as additional PRs. This is a bonus, not a requirement — but well-motivated additions are always welcome. Each PR should meet shapiq's normal bar: tests, docstrings, pre-commit passes, and a clear motivation in the PR description. + +## Relevant Existing Code + +| Path | Description | +|------|-------------| +| `src/shapiq_games/benchmark/local_xai/benchmark_language.py` | Existing hard-coded `SentimentAnalysis` benchmark game (DistilBERT + `[MASK]` / removal) | +| `src/shapiq/plot/sentence.py` | Sentence-level visualization — primary starting point for token / interaction overlays | +| `docs/source/auto_examples/language/plot_sentiment_analysis.py` | Existing sentiment-analysis example (KernelSHAP + KernelSHAPIQ) | +| `src/shapiq/imputer/base.py` | `Imputer` base class — if you build scaffolding, subclass this | +| `src/shapiq/imputer/marginal_imputer.py` / `baseline_imputer.py` | Imputer references for shape and API conventions | +| `src/shapiq/explainer/tabular.py` | How an imputer plugs into an explainer | +| `src/shapiq/game_theory/exact.py` | `ExactComputer` — ground truth for any correctness test on small inputs | +| `src/shapiq/game_theory/indices.py` | Interaction indices available in shapiq (SV, SII, k-SII, STII, FSII, FBII, BV, BII) | + +## References + +- **SHAP text maskers:** *SHAP documentation on text maskers*. [shap.maskers.Text](https://shap.readthedocs.io/en/latest/generated/shap.maskers.Text.html) — the de-facto standard for Shapley on text. +- **Inseq:** Sarti et al., *Inseq: An Interpretability Toolkit for Sequence Generation Models*, ACL 2023. [arXiv:2302.13942](https://arxiv.org/abs/2302.13942). +- **Captum:** Kokhlikyan et al., *Captum: A unified and generic model interpretability library for PyTorch*, 2020. [arXiv:2009.07896](https://arxiv.org/abs/2009.07896). +- **Ferret:** Attanasio et al., *ferret: a Framework for Benchmarking Explainers on Transformers*, EACL 2023. [arXiv:2208.01575](https://arxiv.org/abs/2208.01575) — methodology template for evaluating and comparing text-explanation methods. +- **Shapley attributions for LLMs — recent work:** survey the 2024–2026 arXiv literature on prompt attribution, jailbreak explanation, RAG attribution, and in-context-learning attribution. Bring a reading list to your first group meeting. +- **shapiq paper:** Muschalik et al., *shapiq: Shapley Interactions for Machine Learning*, NeurIPS 2024 — for the architecture and indices you are building on. + +## Expected Deliverables + +**PR(s):** + +- A clean, well-tested PR contributing a text imputer to shapiq (see Task 1). Additional infrastructure contributions (games, visualization utilities) are welcome but not required. +- Tests, docstrings, and passing pre-commit for all PR code. +- All existing tests and pre-commit checks must continue to pass (`uv run pre-commit run --all-files`, `uv run pytest tests/shapiq`). + +**Demo:** + +- A polished, reproducible demo covering **at least three** distinct use cases of Shapley values / Shapley interactions on LLMs, on real SOTA models and real inputs. +- A dedicated comparison section against `shap`, `captum`, and `Inseq` on a shared input + model, with honest runtime / agreement / ergonomics reporting. +- Clear visualizations that make the findings legible to non-expert viewers. +- Fully reproducible: pinned dependencies, pinned HF model revisions, fixed seeds, explicit install and run instructions. diff --git a/pyproject.toml b/pyproject.toml index 93144b1d3..d49a79ecb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,6 +22,7 @@ dependencies = [ "colour", "pillow", ] + authors = [ {name = "Maximilian Muschalik", email = "Maximilian.Muschalik@lmu.de"}, {name = "Santo M. A. R. Thies", email = "Santo.Thies@lmu.de"}, @@ -82,6 +83,13 @@ shapleig = [ "botorch>=0.14.0", "linear_operator>=0.6", # also a transitive gpytorch dependency, but imported directly ] +text = [ + # required by shapiq.imputer.TextImputer and text explanation utilities + "torch", + "transformers", + "nltk>=3.9.4", +] + benchmark = [ # optional model backends for shapiq_benchmark, imported lazily "optuna", @@ -117,6 +125,9 @@ testpaths = [ "tests/shapiq_games", "tests/shapiq_benchmark" ] +markers = [ + "slow: marks tests that use real external model checkpoints", +] pythonpath = ["src"] minversion = "8.0" diff --git a/src/shapiq/imputer/__init__.py b/src/shapiq/imputer/__init__.py index 0420bafbd..268c91135 100644 --- a/src/shapiq/imputer/__init__.py +++ b/src/shapiq/imputer/__init__.py @@ -20,4 +20,15 @@ "TabPFNImputer", "GaussianImputer", "GaussianCopulaImputer", + "TextImputer", ] + + +def __getattr__(name: str) -> type: + if name == "TextImputer": + from .text import TextImputer + + return TextImputer + + msg = f"module {__name__!r} has no attribute {name!r}" + raise AttributeError(msg) diff --git a/src/shapiq/imputer/text/__init__.py b/src/shapiq/imputer/text/__init__.py new file mode 100644 index 000000000..d4c980e16 --- /dev/null +++ b/src/shapiq/imputer/text/__init__.py @@ -0,0 +1,11 @@ +"""Imputation strategies for handling missing feature coalitions. + +All imputers inherit from :class:`~shapiq.imputer.Imputer` and convert a model + +prediction function into a cooperative game by imputing unobserved feature values. + +""" + +from .imputer import TextImputer + +__all__ = ["TextImputer"] diff --git a/src/shapiq/imputer/text/_error.py b/src/shapiq/imputer/text/_error.py new file mode 100644 index 000000000..1070a538b --- /dev/null +++ b/src/shapiq/imputer/text/_error.py @@ -0,0 +1,12 @@ +"""Import error handling for the text module.""" + +from __future__ import annotations + +_text_msg = ( + "The text explanation module requires the optional dependencies " + "torch, transformers, and nltk, but they are not installed.\n" + "Install them with:\n\n" + " pip install 'shapiq[text]'" +) + +_text_import_error = ImportError(_text_msg) diff --git a/src/shapiq/imputer/text/callables.py b/src/shapiq/imputer/text/callables.py new file mode 100644 index 000000000..041dff111 --- /dev/null +++ b/src/shapiq/imputer/text/callables.py @@ -0,0 +1,483 @@ +"""Model callables for coalition-based text explanations.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING + +import numpy as np + +try: + import torch +except ImportError as err: + from ._error import _text_import_error + + raise _text_import_error from err + +if TYPE_CHECKING: + from transformers import PreTrainedModel, PreTrainedTokenizerBase + from transformers.modeling_outputs import BaseModelOutput + +# ============================================================================= +# TARGET CALLABLES +# ============================================================================= + + +class BaseTargetCallable(ABC): + """Abstract interface for model-specific scoring.""" + + def __init__( + self, + model: PreTrainedModel, + tokenizer: PreTrainedTokenizerBase, + device: str, + ) -> None: + """Abstract interface for model-specific scoring.""" + self.model = model + self.tokenizer = tokenizer + self.device = device + + @abstractmethod + def predict( + self, + texts: list[str], + ) -> np.ndarray: + """Return scalar scores.""" + + def predict_from_inputs( + self, + inputs: list[dict[str, torch.Tensor]], + ) -> np.ndarray: + """Return scalar scores from pre-tokenized model inputs.""" + msg = f"{self.__class__.__name__} does not support pre-tokenized inputs." + raise NotImplementedError(msg) + + +# ============================================================================= +# Encoder only support +# ============================================================================= + + +class EncoderClassifierCallable(BaseTargetCallable): + """Score text with an encoder-only sequence-classification model. + + The callable returns either the selected class logit or its softmax probability. + It is suitable for models such as BERT, RoBERTa, or DistilBERT with a sequence-classification head. + """ + + def __init__( + self, + model: PreTrainedModel, + tokenizer: PreTrainedTokenizerBase, + device: str, + class_idx: int = 1, + output_type: str = "logit", + ) -> None: + """Encoder classifier scoring.""" + super().__init__(model, tokenizer, device) + + self.class_idx = class_idx + self.output_type = output_type + + if output_type not in {"logit", "probability"}: + msg = "output_type must be 'logit' or 'probability'" + raise ValueError(msg) + + def predict( + self, + texts: list[str], + ) -> np.ndarray: + """Return one score per text from the configured classifier output. + + Texts are tokenized as a padded batch and evaluated in inference mode. + Depending on ``output_type``, the returned score is either the raw logit or + the softmax probability for ``class_idx``. + """ + encoded = self.tokenizer(texts, padding=True, truncation=True, return_tensors="pt") + + encoded = {key: value.to(self.device) for key, value in encoded.items()} + + with torch.no_grad(): + outputs = self.model(**encoded) + logits = outputs.logits + + if self.output_type == "logit": + scores = logits[:, self.class_idx] + + else: + probs = torch.softmax(logits, dim=-1) + scores = probs[:, self.class_idx] + + return scores.detach().cpu().numpy() + + def predict_from_inputs( + self, + inputs: list[dict[str, torch.Tensor]], + ) -> np.ndarray: + """Run encoder classifier inference from pre-tokenized inputs.""" + input_ids = torch.cat( + [item["input_ids"] for item in inputs], + dim=0, + ).to(self.device) + + attention_mask = torch.cat( + [item["attention_mask"] for item in inputs], + dim=0, + ).to(self.device) + + model_inputs = { + "input_ids": input_ids, + "attention_mask": attention_mask, + } + + if "token_type_ids" in inputs[0]: + model_inputs["token_type_ids"] = torch.cat( + [item["token_type_ids"] for item in inputs], + dim=0, + ).to(self.device) + + with torch.no_grad(): + outputs = self.model(**model_inputs) + logits = outputs.logits + + if self.output_type == "logit": + scores = logits[:, self.class_idx] + else: + probs = torch.softmax(logits, dim=-1) + scores = probs[:, self.class_idx] + + return scores.detach().cpu().numpy() + + +# ============================================================================= +# Causal LM support +# ============================================================================= + + +class CausalLMCallable(BaseTargetCallable): + """Score text using the log-probability of a target causal-LM continuation. + + Each input text is inserted into ``prompt_template``. The score is the + autoregressive log-probability of ``target_label`` after that prompt. + Multi-token target labels are scored token by token, conditioned on the + prompt and all preceding target tokens. + + This supports decoder-only models such as Gemma, Llama, GPT, and Qwen, + provided their tokenizer defines either a padding token or an EOS token. + """ + + def __init__( + self, + model: PreTrainedModel, + tokenizer: PreTrainedTokenizerBase, + device: str, + target_label: str = "good", + prompt_template: str = ("Review: {text}\n\nSentiment:"), + ) -> None: + """Causal language model scoring.""" + super().__init__(model, tokenizer, device) + self.prompt_template = prompt_template + self.target_token_ids = tokenizer.encode(target_label, add_special_tokens=False) + + if len(self.target_token_ids) == 0: + msg = f"Target label '{target_label}' produced no tokens." + raise ValueError(msg) + + if tokenizer.pad_token_id is None: + if tokenizer.eos_token_id is None: + msg = "Tokenizer must define either a pad token or eos token." + raise ValueError(msg) + + tokenizer.pad_token = tokenizer.eos_token + + tokenizer.padding_side = "left" + + def _build_prompt( + self, + text: str, + ) -> str: + """Construct a prompt for causal LM scoring.""" + return self.prompt_template.format(text=text) + + def _score_target_sequence( + self, + prompt: str, + ) -> float: + """Compute the autoregressive log-probability of the target label. + + For each target token, the model receives the prompt followed by preceding target tokens. + The final-position distribution is then used to score the next target token. + """ + prompt_ids = self.tokenizer.encode(prompt, add_special_tokens=False) + target_ids = self.target_token_ids + total_log_prob = 0.0 + + for i in range(len(target_ids)): + prefix_ids = target_ids[:i] + input_ids = prompt_ids + prefix_ids + encoded = {"input_ids": torch.tensor([input_ids], device=self.device)} + + with torch.no_grad(): + outputs = self.model(**encoded) + + logits = outputs.logits + next_token_logits = logits[:, -1, :] + log_probs = torch.log_softmax(next_token_logits, dim=-1) + + total_log_prob += log_probs[0, target_ids[i]].item() + + return total_log_prob + + def predict( + self, + texts: list[str], + ) -> np.ndarray: + """Score texts using target-sequence log probabilities.""" + scores = [] + for text in texts: + prompt = self._build_prompt(text) + score = self._score_target_sequence(prompt) + scores.append(score) + + return np.asarray(scores, dtype=np.float32) + + def predict_from_inputs( + self, + inputs: list[dict[str, torch.Tensor]], + ) -> np.ndarray: + """Score target sequence from pre-tokenized causal-LM prompt inputs.""" + scores = [] + + for item in inputs: + prompt_input_ids = item["input_ids"].to(self.device) + prompt_attention_mask = item["attention_mask"].to(self.device) + + total_log_prob = 0.0 + + for i, target_token_id in enumerate(self.target_token_ids): + prefix_ids = self.target_token_ids[:i] + + if len(prefix_ids) > 0: + prefix_tensor = torch.tensor( + [prefix_ids], + dtype=torch.long, + device=self.device, + ) + + input_ids = torch.cat( + [prompt_input_ids, prefix_tensor], + dim=1, + ) + + prefix_attention_mask = torch.ones_like(prefix_tensor) + + attention_mask = torch.cat( + [prompt_attention_mask, prefix_attention_mask], + dim=1, + ) + else: + input_ids = prompt_input_ids + attention_mask = prompt_attention_mask + + with torch.no_grad(): + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + ) + + logits = outputs.logits + next_token_logits = logits[:, -1, :] + log_probs = torch.log_softmax(next_token_logits, dim=-1) + + total_log_prob += log_probs[0, target_token_id].item() + + scores.append(total_log_prob) + + return np.asarray(scores, dtype=np.float32) + + +# ============================================================================= +# seq2seq support +# ============================================================================= + + +class Seq2SeqCallable(BaseTargetCallable): + """Score a fixed target sequence with an encoder-decoder model. + + For each input text, this callable computes the conditional log-probability + of generating ``target_label`` using teacher forcing. A multi-token target + is scored token by token. By default, the final score is the mean token + log-probability, making targets of different lengths more comparable. + + Args: + model: Encoder-decoder model used for target-sequence scoring. + tokenizer: Tokenizer associated with the model. + device: Device on which model inference is performed. + target_label: Target sequence whose conditional log-probability is scored. + prompt_template: Template used to format each input text. The template must contain a ``{text}`` placeholder. + normalize: Whether to average the total target log-probability over the number of target tokens. + """ + + def __init__( + self, + model: PreTrainedModel, + tokenizer: PreTrainedTokenizerBase, + device: str, + target_label: str = "positive", + prompt_template: str = "{text}", + *, + normalize: bool = True, + ) -> None: + """Initialize seq2seq target-sequence scoring.""" + super().__init__(model, tokenizer, device) + + if not getattr(model.config, "is_encoder_decoder", False): + msg = ( + "Seq2SeqCallable requires an encoder-decoder model with " + "model.config.is_encoder_decoder=True." + ) + raise ValueError(msg) + + self.target_label = target_label + self.prompt_template = prompt_template + self.normalize = normalize + + self.target_token_ids: list[int] = tokenizer.encode( + target_label, + add_special_tokens=False, + ) + if not self.target_token_ids: + msg_0 = f"Target label {target_label!r} produced no tokens after encoding." + raise ValueError(msg_0) + + decoder_start_token_id = model.config.decoder_start_token_id + if decoder_start_token_id is None: + decoder_start_token_id = tokenizer.pad_token_id + + if decoder_start_token_id is None: + msg_1 = ( + "Cannot determine decoder_start_token_id: neither " + "model.config.decoder_start_token_id nor tokenizer.pad_token_id " + "is available." + ) + raise ValueError(msg_1) + + self.decoder_start_token_id = decoder_start_token_id + + def _build_prompt(self, text: str) -> str: + """Wrap the original text into a prompt template.""" + return self.prompt_template.format(text=text) + + def _encode_inputs(self, texts: list[str]) -> dict[str, torch.Tensor]: + """Encode a list of texts into encoder input tensors.""" + encoded = self.tokenizer( + texts, + padding=True, + truncation=True, + return_tensors="pt", + ) + return {key: value.to(self.device) for key, value in encoded.items()} + + def _compute_log_prob_for_target( + self, + encoder_outputs: BaseModelOutput, + attention_mask: torch.Tensor, + batch_size: int, + ) -> np.ndarray: + """Compute the log-probability of the decoder generating the target token sequence.""" + total_log_probs = torch.zeros(batch_size, device=self.device) + + decoder_input_ids = torch.full( + (batch_size, 1), + self.decoder_start_token_id, + dtype=torch.long, + device=self.device, + ) + + for target_token_id in self.target_token_ids: + with torch.no_grad(): + outputs = self.model( + encoder_outputs=encoder_outputs, + attention_mask=attention_mask, + decoder_input_ids=decoder_input_ids, + ) + + log_probs = torch.log_softmax(outputs.logits[:, -1, :], dim=-1) + total_log_probs += log_probs[:, target_token_id] + + next_token = torch.full( + (batch_size, 1), + target_token_id, + dtype=torch.long, + device=self.device, + ) + decoder_input_ids = torch.cat([decoder_input_ids, next_token], dim=1) + + if self.normalize: + total_log_probs /= len(self.target_token_ids) + + return total_log_probs.cpu().numpy() + + def predict(self, texts: list[str]) -> np.ndarray: + """Compute target-sequence log-probability scores. + + Args: + texts: Input texts to score. + + Returns: + One target-sequence log-probability score per input text. + """ + prompts = [self._build_prompt(text) for text in texts] + encoder_inputs = self._encode_inputs(prompts) + + encoder = self.model.get_encoder() + with torch.no_grad(): + encoder_outputs = encoder( + input_ids=encoder_inputs["input_ids"], + attention_mask=encoder_inputs["attention_mask"], + return_dict=True, + ) + + return self._compute_log_prob_for_target( + encoder_outputs=encoder_outputs, + attention_mask=encoder_inputs["attention_mask"], + batch_size=len(prompts), + ) + + def predict_from_inputs( + self, + inputs: list[dict[str, torch.Tensor]], + ) -> np.ndarray: + """Score target sequences from pre-tokenized encoder inputs. + + Args: + inputs: Pre-tokenized encoder inputs containing ``input_ids`` and ``attention_mask``. + + Returns: + One target-sequence log-probability score per input. + + """ + input_ids = torch.cat( + [item["input_ids"] for item in inputs], + dim=0, + ).to(self.device) + + attention_mask = torch.cat( + [item["attention_mask"] for item in inputs], + dim=0, + ).to(self.device) + + encoder = self.model.get_encoder() + + with torch.no_grad(): + encoder_outputs = encoder( + input_ids=input_ids, + attention_mask=attention_mask, + return_dict=True, + ) + + return self._compute_log_prob_for_target( + encoder_outputs=encoder_outputs, + attention_mask=attention_mask, + batch_size=input_ids.shape[0], + ) diff --git a/src/shapiq/imputer/text/imputer.py b/src/shapiq/imputer/text/imputer.py new file mode 100644 index 000000000..3d6a233aa --- /dev/null +++ b/src/shapiq/imputer/text/imputer.py @@ -0,0 +1,422 @@ +"""Text imputer for coalition-based text explanations.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Literal, cast + +import numpy as np + +try: + import torch +except ImportError as err: + from ._error import _text_import_error + + raise _text_import_error from err + +from shapiq.imputer.base import Imputer + +from .callables import ( + CausalLMCallable, + EncoderClassifierCallable, + Seq2SeqCallable, +) +from .perturbations import ( + BasePerturbationStrategy, + MLMInfillingPerturbation, + create_perturbation_strategy, +) +from .players import ( + BasePlayerStrategy, + create_player_strategy, +) +from .tensor_perturbation import ( + TENSOR_PERTURBATION_STRATEGIES, + BaseTensorPerturbationStrategy, + create_tensor_perturbation_strategy, +) + +if TYPE_CHECKING: + from transformers import PreTrainedModel, PreTrainedTokenizerBase + +# ============================================================================= +# TEXT IMPUTER +# ============================================================================= + + +class TextImputer(Imputer): + """Coalition-based text imputer for model-agnostic Shapley explanations. + + ``TextImputer`` combines three independent components: + + - a player strategy, which chooses the text features to explain; + - either a text perturbation strategy, which creates perturbed strings, + or a tensor perturbation strategy, which creates model-ready inputs; + - a target callable, which maps the perturbed representation to a scalar + model score. + + The resulting object is callable with a coalition matrix and can therefore + be passed directly to shapiq approximators. A coalition entry of ``1`` + keeps a player; ``0`` marks it as missing. + + For ordinary perturbation strategies, each coalition is evaluated once. + For ``MLMInfillingPerturbation``, the imputer evaluates multiple sampled + infillings and returns their average score, approximating ``E[f(X) | X_S]``. + + Parameters + ---------- + model + Hugging Face model whose output is explained. + tokenizer + Tokenizer associated with ``model``. + text + Original text instance to explain. + player_level + Player granularity. Available levels are ``"subword"``, ``"word"``, + ``"named_entity"``, ``"chunk"``, and ``"sentence"``. + perturbation_type + Missing-player strategy. Text perturbations include ``"mask"``, ``"pad"``, + ``"removal"``, ``"neutral"``, ``"wordnet_neutral"``, and ``"mlm_infilling"``. + Tensor perturbations include ``"attention_mask"``. + model_type + Target-model interface. Available model types are ``"encoder_classifier"``, + ``"causal_lm"``, and ``"seq2seq"``. + """ + + def __init__( + self, + model: PreTrainedModel, + tokenizer: PreTrainedTokenizerBase, + text: str, + *, + batch_size: int = 16, + device: str | None = None, + # --------------------------------------------------------------------- + # encoder classifier settings + # --------------------------------------------------------------------- + class_idx: int = 1, + output_type: str = "logit", + # --------------------------------------------------------------------- + # causal LM settings + # --------------------------------------------------------------------- + target_label: str = "good", + prompt_template: str = ("Review: {text}\n\nSentiment:"), + # --------------------------------------------------------------------- + # Seq2Seq settings + # --------------------------------------------------------------------- + normalize_target_logprob: bool = True, + # --------------------------------------------------------------------- + # architecture selection + # --------------------------------------------------------------------- + player_level: str = "word", + perturbation_type: str = "mask", + player_strategy: BasePlayerStrategy | None = None, + perturbation_strategy: BasePerturbationStrategy | None = None, + tensor_perturbation_strategy: BaseTensorPerturbationStrategy | None = None, + # --------------------------------------------------------------------- + # MLM infilling settings + # --------------------------------------------------------------------- + mlm_model_name: str = "bert-base-uncased", + mlm_num_samples: int = 100, + # --------------------------------------------------------------------- + # Generalize target callable support. + # --------------------------------------------------------------------- + model_type: str = "encoder_classifier", + ) -> None: + """Initialize the Text Imputer.""" + self.model = model + self.tokenizer = tokenizer + self.text = text + self.batch_size = batch_size + + if device is None: + if torch.cuda.is_available(): + device = "cuda" + + elif torch.backends.mps.is_available(): + device = "mps" + + else: + device = "cpu" + self.device = device + + if not hasattr(self.model, "hf_device_map"): + self.model = self.model.to(self.device) + + self.model.eval() + + # ============================================================================= + # PLAYER STRATEGY + # ============================================================================= + + if player_strategy is None: + player_strategy = create_player_strategy( + level=player_level, text=text, tokenizer=tokenizer + ) + super().__init__( + model=model, + data=np.empty((1, player_strategy.n_players)), + ) + + self.player_level = player_level + self.player_strategy = player_strategy + self.model_type = model_type + + # ============================================================================= + # PERTURBATION STRATEGY + # ============================================================================= + + if perturbation_strategy is not None and tensor_perturbation_strategy is not None: + msg = ( + "Only one of perturbation_strategy and tensor_perturbation_strategy " + "can be provided." + ) + raise ValueError(msg) + + self.perturbation_type = perturbation_type + self.perturbation_mode: Literal["text", "tensor"] + + if perturbation_type in TENSOR_PERTURBATION_STRATEGIES: + if perturbation_strategy is not None: + msg = ( + f"perturbation_type={perturbation_type!r} is a tensor perturbation, " + "so perturbation_strategy must be None. " + "Use tensor_perturbation_strategy instead." + ) + raise ValueError(msg) + + if tensor_perturbation_strategy is None: + tensor_perturbation_strategy = create_tensor_perturbation_strategy( + strategy=perturbation_type, + tokenizer=tokenizer, + ) + + self.perturbation_mode = "tensor" + self.perturbation_strategy = None + self.tensor_perturbation_strategy = tensor_perturbation_strategy + + else: + if tensor_perturbation_strategy is not None: + msg = ( + f"perturbation_type={perturbation_type!r} is a text perturbation, " + "so tensor_perturbation_strategy must be None. " + "Use perturbation_strategy instead." + ) + raise ValueError(msg) + + if perturbation_strategy is None: + perturbation_strategy = create_perturbation_strategy( + strategy=perturbation_type, + tokenizer=tokenizer, + mlm_model_name=mlm_model_name, + mlm_num_samples=mlm_num_samples, + device=self.device, + ) + + self.perturbation_mode = "text" + self.perturbation_strategy = perturbation_strategy + self.tensor_perturbation_strategy = None + + # MLM infilling currently supports only word, named-entity, and chunk players. + + if ( + self.perturbation_mode == "text" + and isinstance(self.perturbation_strategy, MLMInfillingPerturbation) + and self.player_level in {"sentence", "subword"} + ): + msg = "MLMInfillingPerturbation currently supports only word, named-entity, and chunk players." + raise ValueError(msg) + + # ============================================================================= + # TARGET CALLABLE + # ============================================================================= + + if model_type == "encoder_classifier": + self.target_callable = EncoderClassifierCallable( + model=model, + tokenizer=tokenizer, + device=device, + class_idx=class_idx, + output_type=output_type, + ) + + elif model_type == "seq2seq": + self.target_callable = Seq2SeqCallable( + model=model, + tokenizer=tokenizer, + device=self.device, + target_label=target_label, + prompt_template=prompt_template, + normalize=normalize_target_logprob, + ) + + elif model_type == "causal_lm": + self.target_callable = CausalLMCallable( + model=model, + tokenizer=tokenizer, + device=self.device, + target_label=target_label, + prompt_template=prompt_template, + ) + + else: + msg = "model_type must be one of:\n- 'encoder_classifier'\n- 'causal_lm'\n- 'seq2seq'" + raise ValueError(msg) + + self._compute_reference_predictions() + + def coalition_to_text( + self, + coalition: np.ndarray, + ) -> str: + """Convert one coalition into a perturbed text. + + This method is only valid for text perturbation strategies. Tensor + perturbations build model-ready inputs directly and must not be routed + through this string-based path. + """ + if self.perturbation_mode != "text": + msg = ( + "coalition_to_text() can only be used with text perturbation strategies. " + f"Got perturbation_mode={self.perturbation_mode!r}." + ) + raise RuntimeError(msg) + + if self.perturbation_strategy is None: + msg = "perturbation_strategy is required in text perturbation mode." + raise RuntimeError(msg) + + return self.player_strategy.coalition_to_text( + coalition, + self.perturbation_strategy, + ) + + def _coalitions_to_texts( + self, + coalitions: np.ndarray, + ) -> list[str]: + """Convert coalition matrix into perturbed texts.""" + if self.perturbation_mode != "text": + msg = ( + "_coalitions_to_texts() can only be used with text perturbation strategies. " + f"Got perturbation_mode={self.perturbation_mode!r}." + ) + raise RuntimeError(msg) + + return [self.coalition_to_text(coalition) for coalition in coalitions] + + def _predict_batch( + self, + texts: list[str], + ) -> np.ndarray: + """Run model-family-specific inference.""" + return self.target_callable.predict(texts) + + def _batched_predict( + self, + texts: list[str], + ) -> np.ndarray: + """Predict in batches.""" + all_scores = [] + + for start in range(0, len(texts), self.batch_size): + batch = texts[start : start + self.batch_size] + batch_scores = self._predict_batch(batch) + all_scores.append(batch_scores) + + return np.concatenate(all_scores) + + def _batched_predict_from_inputs( + self, + inputs: list[dict[str, torch.Tensor]], + ) -> np.ndarray: + """Predict from pre-tokenized model inputs in batches.""" + all_scores = [] + + for start in range(0, len(inputs), self.batch_size): + batch = inputs[start : start + self.batch_size] + batch_scores = self.target_callable.predict_from_inputs(batch) + all_scores.append(batch_scores) + + return np.concatenate(all_scores) + + def _evaluate_coalitions( + self, + coalitions: np.ndarray, + ) -> np.ndarray: + if self.perturbation_mode == "text" and isinstance( + self.perturbation_strategy, MLMInfillingPerturbation + ): + num_samples = self.perturbation_strategy.num_samples + all_scores = [] + + # Stored for debugging and demonstrations of sampled MLM infillings. + self._last_generated_texts = [] + + for _ in range(num_samples): + self.perturbation_strategy.clear_cache() + texts = self._coalitions_to_texts(coalitions) + self._last_generated_texts.extend(texts) + scores = self._batched_predict(texts) + all_scores.append(scores) + + all_scores = np.stack(all_scores, axis=0) + return np.mean(all_scores, axis=0) + if self.perturbation_mode == "tensor": + if self.tensor_perturbation_strategy is None: + msg = "tensor_perturbation_strategy is required in tensor perturbation mode." + raise RuntimeError(msg) + players = self.player_strategy.get_players() + + masked_inputs = self.tensor_perturbation_strategy.evaluate( + players=players, + coalitions=coalitions, + model_type=self.model_type, + prompt_template=cast( + "str | None", + getattr(self.target_callable, "prompt_template", None), + ), + player_separator="" if self.player_level == "subword" else " ", + ) + + return self._batched_predict_from_inputs(masked_inputs) + + texts = self._coalitions_to_texts(coalitions) + return self._batched_predict(texts) + + def value_function( + self, + coalitions: np.ndarray, + ) -> np.ndarray: + """Evaluate one or more coalitions. + + For text perturbations, each coalition is converted to one perturbed text + and scored once. For MLM infilling, this process is repeated + ``mlm_num_samples`` times with fresh sampled replacements, and the returned + value is the mean score across samples. + + For tensor perturbations, coalitions are converted directly into model-ready + inputs and scored through the target callable's tensor-input interface. + """ + coalitions = np.asarray(coalitions) + + if coalitions.ndim == 1: + coalitions = coalitions.reshape(1, -1) + + if coalitions.shape[1] != self.n_features: + msg = f"Expected coalition width {self.n_features}, got {coalitions.shape[1]}" + raise ValueError(msg) + + scores = self._evaluate_coalitions(coalitions) + empty_mask = ~np.any(coalitions, axis=1) + scores[empty_mask] = self.empty_prediction + return scores + + def _compute_reference_predictions(self) -> None: + self.full_prediction = float( + self._evaluate_coalitions(self.grand_coalition.reshape(1, -1))[0] + ) + + self.empty_prediction = float( + self._evaluate_coalitions(self.empty_coalition.reshape(1, -1))[0] + ) + self.normalization_value = self.empty_prediction diff --git a/src/shapiq/imputer/text/perturbations.py b/src/shapiq/imputer/text/perturbations.py new file mode 100644 index 000000000..9c9a80e61 --- /dev/null +++ b/src/shapiq/imputer/text/perturbations.py @@ -0,0 +1,508 @@ +"""Perturbation strategies used by the TextImputer.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, cast + +import numpy as np + +try: + import nltk + import torch + from nltk.corpus import wordnet as wn + from transformers import AutoModelForMaskedLM, AutoTokenizer +except ImportError as err: + from ._error import _text_import_error + + raise _text_import_error from err + +if TYPE_CHECKING: + from transformers import PreTrainedModel, PreTrainedTokenizerBase + + +def _require_nltk_resource(resource_path: str, download_name: str) -> None: + """Raise a helpful error if an NLTK resource is not installed.""" + try: + nltk.data.find(resource_path) + except LookupError as error: + try: + nltk.data.find(f"{resource_path}.zip") + except LookupError: + pass + else: + return + + msg = ( + f"Missing NLTK resource '{download_name}'. " + "Install it once with:\n\n" + " import nltk\n" + f" nltk.download('{download_name}')\n" + ) + raise LookupError(msg) from error + + +# ============================================================================= +# PERTURBATION STRATEGIES +# ============================================================================= + + +class BasePerturbationStrategy(ABC): + """Abstract interface for representing missing text players. + + A perturbation strategy receives one missing player and returns replacement text. + Simple strategies can ignore coalition-level context, whereas MLM infilling + requires it to generate replacements consistently. + """ + + @abstractmethod + def perturb( + self, + player: str, + *, + context: dict | None = None, + ) -> str: + """Return a replacement for one missing player. + + Parameters: + player: + Original text player that is absent from the coalition. + context: + Optional coalition-level information. When provided, it may contain + ``players``, ``coalition``, and ``mask_index``. MLM infilling requires + this context; simple replacement strategies ignore it. + + Returns: + str: + Replacement text. Returning an empty string removes the player. + """ + + +# ============================================================================= +# [MASK] replacement +# ============================================================================= + + +class MaskTokenPerturbation(BasePerturbationStrategy): + """Replace a missing player with the tokenizer's mask token.""" + + def __init__( + self, + tokenizer: PreTrainedTokenizerBase, + ) -> None: + """Initialize mask-token replacement from the provided tokenizer.""" + self.mask_token = tokenizer.mask_token + + if self.mask_token is None: + msg = ( + "Tokenizer does not define a mask token. " + "MaskTokenPerturbation requires a masked language model tokenizer." + ) + raise ValueError(msg) + + def perturb( + self, + player: str, + *, + context: dict | None = None, + ) -> str: + """Replace missing words with [MASK].""" + del player + del context + return self.mask_token + + +# ============================================================================= +# [PAD] replacement +# ============================================================================= + + +class PadTokenPerturbation(BasePerturbationStrategy): + """Replace missing players with the tokenizer's PAD token.""" + + def __init__( + self, + tokenizer: PreTrainedTokenizerBase, + ) -> None: + """Initialize PAD replacement strategy.""" + self.pad_token = tokenizer.pad_token + + if self.pad_token is None: + msg = f"{tokenizer.__class__.__name__} does not define a pad token." + raise ValueError(msg) + + def perturb( + self, + player: str, + *, + context: dict | None = None, + ) -> str: + """Return the PAD token.""" + del player + del context + return self.pad_token + + +# ============================================================================= +# REMOVAL PERTURBATION +# ============================================================================= + + +class RemovalPerturbation(BasePerturbationStrategy): + """Remove a player by replacing it with an empty string.""" + + def perturb( + self, + player: str, + *, + context: dict | None = None, + ) -> str: + """Remove a player by replacing it with an empty string.""" + del player + del context + return "" + + +# ============================================================================= +# NEUTRAL PERTURBATION +# ============================================================================= + + +class NeutralPerturbation(BasePerturbationStrategy): + """Replace missing players with neutral placeholder text.""" + + def __init__( + self, + neutral_text: str = "something", + ) -> None: + """Replace missing players with neutral placeholder text.""" + self.neutral_text = neutral_text + + def perturb( + self, + player: str, + *, + context: dict | None = None, + ) -> str: + """Return neutral replacement text.""" + del player + del context + return self.neutral_text + + +# ============================================================================= +# WORDNET NEUTRAL PERTURBATION +# ============================================================================= + + +def _penn_to_wn(tag: str) -> str | None: + """Map a Penn Treebank POS tag to a WordNet POS tag when available.""" + if tag.startswith("N"): + return wn.NOUN + + if tag.startswith("V"): + return wn.VERB + + if tag.startswith("J"): + return wn.ADJ + + if tag.startswith("R"): + return wn.ADV + + return None + + +def _get_neutral_replacement(word: str, pos_tag: str) -> str: + """Return a broad WordNet hypernym or a neutral fallback. + + The first WordNet synset and its first hypernym are used as a lightweight, + deterministic approximation of a semantically broader replacement. + If no suitable mapping exists, ``"something"`` is returned. + """ + wn_pos = _penn_to_wn(pos_tag) + + if wn_pos is None: + return "something" + + synsets = wn.synsets(word, pos=wn_pos) + + if not synsets: + return "something" + + hypernyms = synsets[0].hypernyms() + + if not hypernyms: + return "something" + + replacement = hypernyms[0].lemma_names()[0] + + return replacement.replace("_", " ").split()[0] + + +class WordNetNeutralPerturbation(BasePerturbationStrategy): + """WordNet-based semantic neutral replacement.""" + + def perturb( + self, + player: str, + *, + context: dict | None = None, + ) -> str: + """Replace a player with a semantic hypernym.""" + del context + _require_nltk_resource("corpora/wordnet", "wordnet") + _require_nltk_resource( + "taggers/averaged_perceptron_tagger_eng", "averaged_perceptron_tagger_eng" + ) + tag = nltk.pos_tag([player])[0][1] + + return _get_neutral_replacement(player, tag) + + +# ============================================================================= +# MLM Infilling Perturbation +# support: +# - word-level players +# - named-entity-level players +# - chunk-level players +# ============================================================================= + + +class MLMInfillingPerturbation(BasePerturbationStrategy): + """Replace missing player spans with samples from a masked language model. + + For a coalition, all missing players are masked simultaneously and an MLM predicts replacements conditioned on the remaining players. + Replacements are cached per coalition so that all missing players in the same coalition are generated from one MLM forward pass. + This strategy currently supports player strategies whose players represent contiguous text spans: + - ``WordPlayerStrategy`` + - ``NamedEntityPlayerStrategy`` + - ``ChunkPlayerStrategy`` + + It does not support subword or sentence players in the current version. Subword masking can produce invalid token fragments, + while sentence-level masking is not a suitable input representation for the current MLM setup. + + Parameters: + model_name: + Hugging Face masked-language-model checkpoint used for infilling. + device: + Device on which the MLM is evaluated. + num_samples: + Number of independently sampled infillings used by ``TextImputer`` to estimate a coalition value by Monte Carlo averaging. + + """ + + def __init__( + self, + model_name: str = "bert-base-uncased", + device: str = "cpu", + num_samples: int = 100, + ) -> None: + """Initialize MLM-based infilling strategy.""" + self.tokenizer = cast( + "PreTrainedTokenizerBase", + AutoTokenizer.from_pretrained(model_name), + ) + self.model_name = model_name + self.device = device + model = AutoModelForMaskedLM.from_pretrained(model_name) + + model = model.to(device) + + self.model = cast( + "PreTrainedModel", + model, + ) + self.model.eval() + + self.mask_token = self.tokenizer.mask_token + + if self.mask_token is None: + msg = f"{model_name} does not define a mask token." + raise ValueError(msg) + + self._cache: dict[tuple, dict[int, str]] = {} + + self.num_samples = num_samples + + def clear_cache(self) -> None: + """Discard cached replacements before generating a new MLM sample. + + ``TextImputer`` calls this once per Monte Carlo sample.Within one sample, + the cache ensures that all missing players of the same coalitionuse replacements from the same MLM forward pass. + + """ + self._cache.clear() + + def _build_cache_key( + self, + players: list[str], + coalition: np.ndarray, + ) -> tuple: + """Create a stable cache key for one player sequence and coalition.""" + return (tuple(players), tuple(coalition.tolist())) + + def _predict_masks( + self, + players: list[str], + coalition: np.ndarray, + ) -> dict[int, str]: + """Predict one replacement for every missing player in a coalition. + + All absent players are replaced by the MLM mask token at once. Candidate + tokens that decode to special tokens, empty strings, or WordPiece continuations + are rejected. If no valid candidate is sampled after a fixed number of attempts, + the neutral fallback ``"something"`` is used. + + """ + masked_players = [] + + for keep, player in zip(coalition, players, strict=False): + if keep: + masked_players.append(player) + + else: + masked_players.append(self.mask_token) + + text = " ".join(masked_players) + + encoded = self.tokenizer(text, return_tensors="pt") + + encoded = {k: v.to(self.device) for k, v in encoded.items()} + + with torch.no_grad(): + outputs = self.model(**encoded) + + logits = outputs.logits + + mask_positions = (encoded["input_ids"][0] == self.tokenizer.mask_token_id).nonzero( + as_tuple=True + )[0] + + replacements = {} + + for player_idx, token_pos in zip(np.where(coalition == 0)[0], mask_positions, strict=False): + probs = torch.softmax(logits[0, token_pos], dim=-1) + + replacement = "something" + max_sampling_attempts = 100 + + for _ in range(max_sampling_attempts): + candidate_id = int( + torch.multinomial( + probs, + num_samples=1, + ).item() + ) + + token = cast( + "str", + self.tokenizer.decode( + [candidate_id], + skip_special_tokens=True, + ), + ).strip() + + if token == "": + continue + + if token in { + self.tokenizer.cls_token, + self.tokenizer.sep_token, + self.tokenizer.pad_token, + self.tokenizer.mask_token, + }: + continue + + if token.startswith("##"): + continue + + replacement = token + break + + replacements[player_idx] = replacement + + return replacements + + def perturb( + self, + player: str, + *, + context: dict | None = None, + ) -> str: + """Return the cached or newly generated MLM replacement for one player. + + The ``context`` dictionary must contain the complete player list, the + coalition, and the index of the currently missing player. This allows one + coalition-level MLM prediction to be shared across all missing players. + + """ + if context is None: + msg = "MLMInfillingPerturbation requires context." + raise ValueError(msg) + + players = context["players"] + coalition = np.asarray(context["coalition"]) + mask_index = context["mask_index"] + + cache_key = self._build_cache_key(players, coalition) + + if cache_key not in self._cache: + self._cache[cache_key] = self._predict_masks(players, coalition) + + replacements = self._cache[cache_key] + + return replacements.get(mask_index, player) + + +# ============================================================================= +# PERTURBATION DICTIONARY AND FACTORY +# ============================================================================= + +PERTURBATION_STRATEGIES = { + "mask": MaskTokenPerturbation, + "pad": PadTokenPerturbation, + "removal": RemovalPerturbation, + "neutral": NeutralPerturbation, + "wordnet_neutral": WordNetNeutralPerturbation, + "mlm_infilling": MLMInfillingPerturbation, +} + + +def create_perturbation_strategy( + strategy: str, + tokenizer: PreTrainedTokenizerBase, + mlm_model_name: str = "bert-base-uncased", + mlm_num_samples: int = 100, + device: str = "cpu", +) -> BasePerturbationStrategy: + """Create a perturbation strategy from a string identifier.""" + if strategy not in PERTURBATION_STRATEGIES: + msg = ( + f"Unknown perturbation strategy: {strategy}. " + f"Available strategies: {list(PERTURBATION_STRATEGIES)}" + ) + + raise ValueError(msg) + if strategy == "mask": + return MaskTokenPerturbation(tokenizer) + + if strategy == "pad": + return PadTokenPerturbation(tokenizer) + + if strategy == "removal": + return RemovalPerturbation() + + if strategy == "neutral": + return NeutralPerturbation() + + if strategy == "wordnet_neutral": + return WordNetNeutralPerturbation() + + if strategy == "mlm_infilling": + return MLMInfillingPerturbation( + model_name=mlm_model_name, + device=device, + num_samples=mlm_num_samples, + ) + msg_0 = f"Unhandled perturbation strategy: {strategy}" + raise RuntimeError(msg_0) diff --git a/src/shapiq/imputer/text/players.py b/src/shapiq/imputer/text/players.py new file mode 100644 index 000000000..bc7c94307 --- /dev/null +++ b/src/shapiq/imputer/text/players.py @@ -0,0 +1,394 @@ +"""Player strategies used by the TextImputer.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING + +try: + import nltk + from nltk.tree import Tree +except ImportError as err: + from ._error import _text_import_error + + raise _text_import_error from err + +if TYPE_CHECKING: + import numpy as np + from transformers import PreTrainedTokenizerBase + + from .perturbations import BasePerturbationStrategy + + +def _require_nltk_resource(resource_path: str, download_name: str) -> None: + """Raise a helpful error if an NLTK resource is not installed.""" + try: + nltk.data.find(resource_path) + except LookupError as error: + try: + nltk.data.find(f"{resource_path}.zip") + except LookupError: + pass + else: + return + + msg = ( + f"Missing NLTK resource '{download_name}'. " + "Install it once with:\n\n" + " import nltk\n" + f" nltk.download('{download_name}')\n" + ) + raise LookupError(msg) from error + + +# ============================================================================= +# PLAYER STRATEGIES +# ============================================================================= + + +class BasePlayerStrategy(ABC): + """Abstract interface for converting text into Shapley players. + + A player strategy defines the feature granularity used for attribution. + For example, a text can be represented by subwords, words, named entities, syntactic chunks, or sentences. + + Implementations must provide: + - ``get_players`` to expose the extracted text units; + - ``n_players`` to report the number of units; + - ``coalition_to_text`` to reconstruct a perturbed text for a coalition. + + A coalition uses ``1`` for a kept player and ``0`` for a missing player. + Missing players are replaced by the supplied perturbation strategy. + """ + + _passes_context = True + + @abstractmethod + def get_players(self) -> list[str]: + """Return player list.""" + + def coalition_to_text( + self, + coalition: np.ndarray, + perturbation_strategy: BasePerturbationStrategy, + ) -> str: + """Reconstruct text for a coalition using a perturbation strategy. + + Parameters: + coalition: + One-dimensional binary vector. A value of ``1`` keeps the corresponding player; a value of ``0`` replaces it using ``perturbation_strategy``. + perturbation_strategy: + Strategy that determines how missing players are represented. + + Returns: + str: The perturbed text corresponding to the coalition. + """ + if len(coalition) != self.n_players: + msg = f"Coalition length {len(coalition)} does not match n_players={self.n_players}" + raise ValueError(msg) + + players = self.get_players() + output_parts: list[str] = [] + + for idx, (keep, player) in enumerate(zip(coalition, players, strict=False)): + if keep: + output_parts.append(player) + else: + context = ( + { + "players": players, + "coalition": coalition, + "mask_index": idx, + } + if self._passes_context + else None + ) + + replacement = perturbation_strategy.perturb( + player, + context=context, + ) + + if replacement != "": + output_parts.append(replacement) + return self._join(output_parts) + + @abstractmethod + def _join( + self, + parts: list[str], + ) -> str: + """Join perturbed players into the final text.""" + + @property + @abstractmethod + def n_players(self) -> int: + """Return number of players.""" + + +# ============================================================================= +# SUBWORD-LEVEL PLAYER STRATEGY +# ============================================================================= + + +class SubwordPlayerStrategy(BasePlayerStrategy): + """Tokenizer/subword-level player strategy. + + Uses the provided HuggingFace tokenizer to define players as tokenizer tokens (WordPiece/BPE/SentencePiece, etc.). + """ + + def __init__( + self, + text: str, + tokenizer: PreTrainedTokenizerBase, + ) -> None: + """Initialize subword-level player strategy.""" + self.text = text + self.tokenizer = tokenizer + + self.subwords = tokenizer.tokenize(text) + + def get_players(self) -> list[str]: + """Return subword players.""" + return self.subwords + + @property + def n_players(self) -> int: + """Return number of subword players.""" + return len(self.subwords) + + def _join( + self, + parts: list[str], + ) -> str: + """Join subword tokens into text.""" + return self.tokenizer.convert_tokens_to_string(parts) + + +# ============================================================================= +# WORD-LEVEL PLAYER STRATEGY +# ============================================================================= + + +class WordPlayerStrategy(BasePlayerStrategy): + """Word-level player strategy.""" + + def __init__( + self, + text: str, + ) -> None: + """Initialize word-level player strategy.""" + self.text = text + _require_nltk_resource("tokenizers/punkt_tab", "punkt_tab") + self.words = nltk.word_tokenize(text) + + def get_players(self) -> list[str]: + """Return word players.""" + return self.words + + @property + def n_players(self) -> int: + """Return number of word players.""" + return len(self.words) + + def _join( + self, + parts: list[str], + ) -> str: + """Join words into text.""" + return " ".join(parts) + + +# ============================================================================= +# NAMED-ENTITY PLAYER STRATEGY +# ============================================================================= + + +class NamedEntityPlayerStrategy(BasePlayerStrategy): + """Named-entity-level player strategy using NLTK NER.""" + + def __init__( + self, + text: str, + ) -> None: + """Initialize named-entity player strategy.""" + self.text = text + + _require_nltk_resource("tokenizers/punkt_tab", "punkt_tab") + _require_nltk_resource( + "taggers/averaged_perceptron_tagger_eng", "averaged_perceptron_tagger_eng" + ) + _require_nltk_resource("chunkers/maxent_ne_chunker_tab", "maxent_ne_chunker_tab") + _require_nltk_resource("corpora/words", "words") + + tokens = nltk.word_tokenize(text) + pos_tags = nltk.pos_tag(tokens) + ner_tree = nltk.ne_chunk(pos_tags) + + self.players: list[str] = [] + + for node in ner_tree: + if isinstance(node, Tree): + entity = " ".join(word for word, _tag in node.leaves()) + + self.players.append(entity) + + else: + word, _tag = node + self.players.append(word) + + def get_players(self) -> list[str]: + """Return entity-aware players.""" + return self.players + + @property + def n_players(self) -> int: + """Return number of players.""" + return len(self.players) + + def _join( + self, + parts: list[str], + ) -> str: + """Join entity into text.""" + return " ".join(parts) + + +# ============================================================================= +# CHUNK-LEVEL PLAYER STRATEGY +# ============================================================================= + + +class ChunkPlayerStrategy(BasePlayerStrategy): + """Chunk-level player strategy using POS-based phrase chunking.""" + + def __init__( + self, + text: str, + ) -> None: + """Initialize chunk-level player strategy.""" + self.text = text + _require_nltk_resource("tokenizers/punkt_tab", "punkt_tab") + _require_nltk_resource( + "taggers/averaged_perceptron_tagger_eng", "averaged_perceptron_tagger_eng" + ) + tokens = nltk.word_tokenize(text) + pos_tags = nltk.pos_tag(tokens) + + grammar = r""" + NP: {
?*+} + VP: {*} + """ + + chunker = nltk.RegexpParser(grammar) + tree = chunker.parse(pos_tags) + + self.chunks: list[str] = [] + + for node in tree: + if isinstance(node, Tree): + phrase = " ".join(word for word, _tag in node.leaves()) + self.chunks.append(phrase) + + else: + word, _tag = node + self.chunks.append(word) + + def get_players(self) -> list[str]: + """Return chunk players.""" + return self.chunks + + @property + def n_players(self) -> int: + """Return number of chunk players.""" + return len(self.chunks) + + def _join( + self, + parts: list[str], + ) -> str: + """Join chunks into text.""" + return " ".join(parts) + + +# ============================================================================= +# SENTENCE-LEVEL PLAYER STRATEGY +# ============================================================================= + + +class SentencePlayerStrategy(BasePlayerStrategy): + """Sentence-level player strategy using NLTK sentence splitting.""" + + # Sentence perturbations do not use context. + _passes_context = False + + def __init__( + self, + text: str, + ) -> None: + """Sentence-level player strategy using NLTK sentence splitting.""" + self.text = text + _require_nltk_resource("tokenizers/punkt_tab", "punkt_tab") + self.sentences = nltk.sent_tokenize(text) + + def get_players(self) -> list[str]: + """Return sentence players.""" + return self.sentences + + @property + def n_players(self) -> int: + """Return number of sentence players.""" + return len(self.sentences) + + def _join( + self, + parts: list[str], + ) -> str: + """Join sentences into text.""" + return " ".join(parts) + + +# ============================================================================= +# PLAYER DICTIONARY AND FACTORY +# ============================================================================= + +PLAYER_STRATEGIES = { + "subword": SubwordPlayerStrategy, + "word": WordPlayerStrategy, + "named_entity": NamedEntityPlayerStrategy, + "chunk": ChunkPlayerStrategy, + "sentence": SentencePlayerStrategy, +} + + +def create_player_strategy( + level: str, + text: str, + tokenizer: PreTrainedTokenizerBase, +) -> BasePlayerStrategy: + """Create a player strategy from a string identifier.""" + if level == "subword": + return SubwordPlayerStrategy( + text=text, + tokenizer=tokenizer, + ) + + if level == "word": + return WordPlayerStrategy(text=text) + + if level == "named_entity": + return NamedEntityPlayerStrategy(text=text) + + if level == "chunk": + return ChunkPlayerStrategy(text=text) + + if level == "sentence": + return SentencePlayerStrategy(text=text) + + msg = ( + f"Unknown player level: {level}. " + "Available levels: " + "['subword', 'word', 'named_entity', 'chunk', 'sentence']" + ) + + raise ValueError(msg) diff --git a/src/shapiq/imputer/text/tensor_perturbation.py b/src/shapiq/imputer/text/tensor_perturbation.py new file mode 100644 index 000000000..1e9eb490a --- /dev/null +++ b/src/shapiq/imputer/text/tensor_perturbation.py @@ -0,0 +1,378 @@ +"""Tensor perturbation strategies used by the TextImputer. + +Tensor perturbations do not create perturbed strings. Instead, they build +model-ready tensor inputs, such as ``input_ids`` and ``attention_mask``. +They are intentionally separated from text perturbation strategies because +they use a different interface and should not be routed through +``_coalitions_to_texts``. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING + +import numpy as np + +try: + import torch +except ImportError as err: + from ._error import _text_import_error + + raise _text_import_error from err + +if TYPE_CHECKING: + from transformers import PreTrainedTokenizerBase + +# ============================================================================= +# Tensor PERTURBATION STRATEGIES +# ============================================================================= + + +class BaseTensorPerturbationStrategy(ABC): + """Base class for perturbations that produce model-ready tensor inputs. + + Unlike text perturbation strategies, tensor perturbations do not implement + a string-in/string-out ``perturb`` method. They directly build model inputs + for coalitions and are consumed by tensor-based prediction paths. + """ + + @abstractmethod + def evaluate( + self, + players: list[str], + coalitions: np.ndarray, + *, + model_type: str, + prompt_template: str | None = None, + player_separator: str = "", + ) -> list[dict[str, torch.Tensor]]: + """Build model-ready inputs for coalitions using attention masking. + + For encoder classifiers, players are tokenized as one input sequence. For + causal LM and seq2seq models, players are inserted into ``prompt_template`` and + only player tokens inside ``"{text}"`` are maskable. + """ + + +class AttentionMaskPerturbation(BaseTensorPerturbationStrategy): + """Build model inputs by masking missing players in the attention mask. + + This perturbation does not create perturbed strings. Instead, it maps + players to token spans and sets the corresponding attention mask entries + of missing players to 0. + """ + + def __init__( + self, + tokenizer: PreTrainedTokenizerBase, + ) -> None: + """Initialize attention-mask perturbation.""" + self.tokenizer = tokenizer + + @staticmethod + def build_attention_mask_for_coalition( + base_attention_mask: torch.Tensor, + player_spans: list[tuple[int, int]], + coalition: np.ndarray, + ) -> torch.Tensor: + """Build an attention mask for one coalition.""" + coalition = np.asarray(coalition, dtype=bool) + + if len(coalition) != len(player_spans): + msg = ( + f"Coalition length {len(coalition)} does not match " + f"number of player spans {len(player_spans)}." + ) + raise ValueError(msg) + + attention_mask = base_attention_mask.clone() + + for keep, (start, end) in zip(coalition, player_spans, strict=False): + if not keep: + attention_mask[..., start:end] = 0 + + return attention_mask + + @staticmethod + def build_tokenized_players( + players: list[str], + tokenizer: PreTrainedTokenizerBase, + player_separator: str = "", + ) -> tuple[dict[str, torch.Tensor], list[tuple[int, int]]]: + """Tokenize players into one sequence and record their token spans.""" + all_token_ids: list[int] = [] + player_spans: list[tuple[int, int]] = [] + + for idx, player in enumerate(players): + text_piece = player if idx == 0 else f"{player_separator}{player}" + + token_ids = tokenizer.encode( + text_piece, + add_special_tokens=False, + ) + + start = len(all_token_ids) + all_token_ids.extend(token_ids) + end = len(all_token_ids) + + player_spans.append((start, end)) + + input_ids = torch.tensor( + [all_token_ids], + dtype=torch.long, + ) + attention_mask = torch.ones_like(input_ids) + + return ( + { + "input_ids": input_ids, + "attention_mask": attention_mask, + }, + player_spans, + ) + + @classmethod + def build_inputs_for_coalitions( + cls, + tokenizer: PreTrainedTokenizerBase, + players: list[str], + coalitions: np.ndarray, + player_separator: str = "", + ) -> list[dict[str, torch.Tensor]]: + """Build masked model inputs using attention masking. + + Args: + tokenizer: HuggingFace tokenizer. + players: Text players to explain. + coalitions: Coalition matrix of shape ``(n_coalitions, n_players)``. + player_separator: String inserted between adjacent players before tokenization. + + Returns: + A list of model input dictionaries. Each dictionary contains: + - input_ids: Encoded text. + - attention_mask: Attention mask with missing-player tokens set to 0. + """ + coalitions = np.asarray(coalitions, dtype=bool) + + if coalitions.ndim == 1: + coalitions = coalitions.reshape(1, -1) + + encoded, player_spans = cls.build_tokenized_players( + players=players, + tokenizer=tokenizer, + player_separator=player_separator, + ) + + if coalitions.shape[1] != len(player_spans): + msg = f"Expected coalition width {len(player_spans)}, got {coalitions.shape[1]}." + raise ValueError(msg) + + masked_inputs: list[dict[str, torch.Tensor]] = [] + + for coalition in coalitions: + attention_mask = cls.build_attention_mask_for_coalition( + base_attention_mask=encoded["attention_mask"], + player_spans=player_spans, + coalition=coalition, + ) + + masked_inputs.append( + { + "input_ids": encoded["input_ids"], + "attention_mask": attention_mask, + }, + ) + + return masked_inputs + + @staticmethod + def build_tokenized_prompt_players( + players: list[str], + tokenizer: PreTrainedTokenizerBase, + prompt_template: str, + player_separator: str = "", + ) -> tuple[dict[str, torch.Tensor], list[tuple[int, int]]]: + """Tokenize prompt-wrapped players and record player token spans. + + This is for causal LM scoring. The prompt template must contain "{text}". + Only tokens corresponding to players inside "{text}" are maskable. + Prompt instruction tokens are always kept visible. + """ + if "{text}" not in prompt_template: + msg = "prompt_template must contain '{text}'." + raise ValueError(msg) + + prefix, suffix = prompt_template.split("{text}", maxsplit=1) + + prefix_ids = tokenizer.encode( + prefix, + add_special_tokens=False, + ) + + suffix_ids = tokenizer.encode( + suffix, + add_special_tokens=False, + ) + + all_token_ids: list[int] = [] + player_spans: list[tuple[int, int]] = [] + + all_token_ids.extend(prefix_ids) + + for idx, player in enumerate(players): + text_piece = player if idx == 0 else f"{player_separator}{player}" + + token_ids = tokenizer.encode( + text_piece, + add_special_tokens=False, + ) + + start = len(all_token_ids) + all_token_ids.extend(token_ids) + end = len(all_token_ids) + + player_spans.append((start, end)) + + all_token_ids.extend(suffix_ids) + + input_ids = torch.tensor( + [all_token_ids], + dtype=torch.long, + ) + attention_mask = torch.ones_like(input_ids) + + return ( + { + "input_ids": input_ids, + "attention_mask": attention_mask, + }, + player_spans, + ) + + @classmethod + def build_prompt_inputs_for_coalitions( + cls, + tokenizer: PreTrainedTokenizerBase, + players: list[str], + coalitions: np.ndarray, + prompt_template: str, + player_separator: str = "", + ) -> list[dict[str, torch.Tensor]]: + """Build causal-LM prompt inputs using attention masking. + + The generated inputs represent prompt_template.format(text=players_text), + but attention masking is applied only to player tokens inside "{text}". + """ + coalitions = np.asarray(coalitions, dtype=bool) + + if coalitions.ndim == 1: + coalitions = coalitions.reshape(1, -1) + + encoded, player_spans = cls.build_tokenized_prompt_players( + players=players, + tokenizer=tokenizer, + prompt_template=prompt_template, + player_separator=player_separator, + ) + + if coalitions.shape[1] != len(player_spans): + msg = f"Expected coalition width {len(player_spans)}, got {coalitions.shape[1]}." + raise ValueError(msg) + + masked_inputs: list[dict[str, torch.Tensor]] = [] + + for coalition in coalitions: + attention_mask = cls.build_attention_mask_for_coalition( + base_attention_mask=encoded["attention_mask"], + player_spans=player_spans, + coalition=coalition, + ) + + masked_inputs.append( + { + "input_ids": encoded["input_ids"], + "attention_mask": attention_mask, + }, + ) + + return masked_inputs + + def evaluate( + self, + players: list[str], + coalitions: np.ndarray, + *, + model_type: str = "encoder_classifier", + prompt_template: str | None = None, + player_separator: str = "", + ) -> list[dict[str, torch.Tensor]]: + """Build masked model inputs using attention masking.""" + if model_type == "causal_lm": + if prompt_template is None: + msg = "prompt_template is required for causal_lm attention masking." + raise ValueError(msg) + + return self.build_prompt_inputs_for_coalitions( + tokenizer=self.tokenizer, + players=players, + coalitions=coalitions, + prompt_template=prompt_template, + player_separator=player_separator, + ) + + if model_type == "encoder_classifier": + return self.build_inputs_for_coalitions( + tokenizer=self.tokenizer, + players=players, + coalitions=coalitions, + player_separator=player_separator, + ) + + if model_type == "seq2seq": + if prompt_template is None: + msg = "prompt_template is required for seq2seq attention masking." + raise ValueError(msg) + + return self.build_prompt_inputs_for_coalitions( + tokenizer=self.tokenizer, + players=players, + coalitions=coalitions, + prompt_template=prompt_template, + player_separator=player_separator, + ) + + msg = f"Unknown model_type for attention masking: {model_type}." + raise ValueError(msg) + + +TENSOR_PERTURBATION_STRATEGIES: dict[ + str, + type[BaseTensorPerturbationStrategy], +] = { + "attention_mask": AttentionMaskPerturbation, +} + + +def create_tensor_perturbation_strategy( + strategy: str, + *, + tokenizer: PreTrainedTokenizerBase, +) -> BaseTensorPerturbationStrategy: + """Create a tensor perturbation strategy from a string identifier. + + This factory is intentionally separate from the text perturbation factory + to avoid mixing string-returning and tensor-returning perturbations. + """ + if strategy not in TENSOR_PERTURBATION_STRATEGIES: + msg = ( + f"Unknown tensor perturbation strategy: {strategy}. " + f"Available strategies: {list(TENSOR_PERTURBATION_STRATEGIES)}" + ) + raise ValueError(msg) + + if strategy == "attention_mask": + return AttentionMaskPerturbation(tokenizer=tokenizer) + + msg = f"Unhandled tensor perturbation strategy: {strategy}" + raise RuntimeError(msg) diff --git a/src/shapiq/imputer/text_imputer.py b/src/shapiq/imputer/text_imputer.py new file mode 100644 index 000000000..e42f35bb3 --- /dev/null +++ b/src/shapiq/imputer/text_imputer.py @@ -0,0 +1,73 @@ +"""Backward-compatible imports for the text imputer.""" + +from __future__ import annotations + +from .text.callables import ( + BaseTargetCallable, + CausalLMCallable, + EncoderClassifierCallable, + Seq2SeqCallable, +) +from .text.imputer import TextImputer +from .text.perturbations import ( + PERTURBATION_STRATEGIES, + BasePerturbationStrategy, + MaskTokenPerturbation, + MLMInfillingPerturbation, + NeutralPerturbation, + PadTokenPerturbation, + RemovalPerturbation, + WordNetNeutralPerturbation, + _get_neutral_replacement, + _penn_to_wn, + _require_nltk_resource, + create_perturbation_strategy, +) +from .text.players import ( + PLAYER_STRATEGIES, + BasePlayerStrategy, + ChunkPlayerStrategy, + NamedEntityPlayerStrategy, + SentencePlayerStrategy, + SubwordPlayerStrategy, + WordPlayerStrategy, + create_player_strategy, +) +from .text.tensor_perturbation import ( + TENSOR_PERTURBATION_STRATEGIES, + AttentionMaskPerturbation, + BaseTensorPerturbationStrategy, + create_tensor_perturbation_strategy, +) + +__all__ = [ + "PERTURBATION_STRATEGIES", + "PLAYER_STRATEGIES", + "TENSOR_PERTURBATION_STRATEGIES", + "AttentionMaskPerturbation", + "BasePerturbationStrategy", + "BasePlayerStrategy", + "BaseTargetCallable", + "BaseTensorPerturbationStrategy", + "CausalLMCallable", + "ChunkPlayerStrategy", + "EncoderClassifierCallable", + "MLMInfillingPerturbation", + "MaskTokenPerturbation", + "NamedEntityPlayerStrategy", + "NeutralPerturbation", + "PadTokenPerturbation", + "RemovalPerturbation", + "SentencePlayerStrategy", + "Seq2SeqCallable", + "SubwordPlayerStrategy", + "TextImputer", + "WordNetNeutralPerturbation", + "WordPlayerStrategy", + "_get_neutral_replacement", + "_penn_to_wn", + "_require_nltk_resource", + "create_perturbation_strategy", + "create_player_strategy", + "create_tensor_perturbation_strategy", +] diff --git a/tests/shapiq/tests_unit/tests_imputer/test_text_imputer.py b/tests/shapiq/tests_unit/tests_imputer/test_text_imputer.py new file mode 100644 index 000000000..395ce8cf2 --- /dev/null +++ b/tests/shapiq/tests_unit/tests_imputer/test_text_imputer.py @@ -0,0 +1,1393 @@ +"""Fast mock-based unit tests for shapiq TextImputer.""" + +from __future__ import annotations + +import os +from types import SimpleNamespace +from unittest.mock import MagicMock, call, patch + +import numpy as np +import pytest +from nltk.tree import Tree + +torch = pytest.importorskip("torch") +pytest.importorskip("transformers") + +from transformers import AutoModelForSequenceClassification, AutoTokenizer # noqa: E402 + +from shapiq.imputer.text_imputer import ( # noqa: E402 + BaseTargetCallable, + CausalLMCallable, + ChunkPlayerStrategy, + EncoderClassifierCallable, + MaskTokenPerturbation, + MLMInfillingPerturbation, + NamedEntityPlayerStrategy, + NeutralPerturbation, + PadTokenPerturbation, + RemovalPerturbation, + SentencePlayerStrategy, + SubwordPlayerStrategy, + TextImputer, + WordNetNeutralPerturbation, + WordPlayerStrategy, + _get_neutral_replacement, + _penn_to_wn, + _require_nltk_resource, + create_perturbation_strategy, + create_player_strategy, +) + +MODULE = "shapiq.imputer.text_imputer" +PLAYERS_MODULE = "shapiq.imputer.text.players" +PERTURBATIONS_MODULE = "shapiq.imputer.text.perturbations" +CALLABLES_MODULE = "shapiq.imputer.text.callables" + + +class DummyTokenizer: + """Small tokenizer substitute used by fast unit tests.""" + + mask_token = "[MASK]" + mask_token_id = 99 + pad_token = "[PAD]" + pad_token_id = 0 + eos_token = "" + eos_token_id = 2 + cls_token = "[CLS]" + sep_token = "[SEP]" + padding_side = "right" + + def __init__(self) -> None: + self.encode_return_values: list[list[int]] = [] + + def tokenize(self, text: str) -> list[str]: + return ["un", "##happy"] + + def convert_tokens_to_string(self, tokens: list[str]) -> str: + return " ".join(tokens).replace(" ##", "") + + def encode( + self, + text: str, + *, + add_special_tokens: bool = False, + ) -> list[int]: + return self.encode_return_values.pop(0) + + def __call__( + self, + texts, + *, + padding: bool = False, + truncation: bool = False, + return_tensors: str = "pt", + ) -> dict[str, torch.Tensor]: + if isinstance(texts, str): + texts = [texts] + + return { + "input_ids": torch.tensor( + [[10, 99, 11] for _ in texts], + dtype=torch.long, + ), + "attention_mask": torch.ones( + (len(texts), 3), + dtype=torch.long, + ), + } + + def decode( + self, + token_ids: list[int], + *, + skip_special_tokens: bool = True, + ) -> str: + mapping = { + 0: "[PAD]", + 1: "[CLS]", + 2: "[SEP]", + 3: "[MASK]", + 4: "##suffix", + 5: "great", + 6: "thing", + } + return mapping[token_ids[0]] + + +@pytest.fixture +def tokenizer() -> DummyTokenizer: + return DummyTokenizer() + + +@pytest.fixture +def model() -> MagicMock: + model = MagicMock() + model.to.return_value = model + model.return_value = SimpleNamespace( + logits=torch.tensor([[1.0, 2.0]]), + ) + + return model + + +@pytest.fixture +def no_nltk_resource_check(): + """Avoid touching local NLTK data in tests.""" + with ( + patch("shapiq.imputer.text.players._require_nltk_resource"), + patch("shapiq.imputer.text.perturbations._require_nltk_resource"), + ): + yield + + +# ============================================================================ +# NLTK helper +# ============================================================================ + + +def test_require_nltk_resource_passes_when_resource_exists() -> None: + with patch(f"{PLAYERS_MODULE}.nltk.data.find") as find: + _require_nltk_resource("tokenizers/punkt_tab", "punkt_tab") + + find.assert_called_once_with("tokenizers/punkt_tab") + + +def test_require_nltk_resource_passes_when_zip_resource_exists() -> None: + with patch( + f"{PLAYERS_MODULE}.nltk.data.find", + side_effect=[LookupError("not installed"), None], + ) as find: + _require_nltk_resource("tokenizers/punkt_tab", "punkt_tab") + + assert find.call_count == 2 + + assert find.call_args_list == [ + call("tokenizers/punkt_tab"), + call("tokenizers/punkt_tab.zip"), + ] + + +def test_require_nltk_resource_has_helpful_error_when_missing() -> None: + with ( + patch( + f"{PLAYERS_MODULE}.nltk.data.find", + side_effect=LookupError("not installed"), + ), + pytest.raises(LookupError, match=r"nltk\.download\('punkt_tab'\)"), + ): + _require_nltk_resource("tokenizers/punkt_tab", "punkt_tab") + + +# ============================================================================ +# Player strategies +# ============================================================================ + + +def test_subword_player_strategy(tokenizer: DummyTokenizer) -> None: + strategy = SubwordPlayerStrategy("unhappy", tokenizer) + + assert strategy.get_players() == ["un", "##happy"] + assert strategy.n_players == 2 + + text = strategy.coalition_to_text( + np.array([1, 0]), + NeutralPerturbation("X"), + ) + + assert text == "un X" + + +def test_subword_player_strategy_rejects_wrong_coalition_length( + tokenizer: DummyTokenizer, +) -> None: + strategy = SubwordPlayerStrategy("unhappy", tokenizer) + + with pytest.raises(ValueError, match="does not match n_players=2"): + strategy.coalition_to_text( + np.array([1]), + NeutralPerturbation(), + ) + + +def test_word_player_strategy_with_mocked_nltk( + no_nltk_resource_check, +) -> None: + with patch( + f"{PLAYERS_MODULE}.nltk.word_tokenize", + return_value=["I", "love", "cats"], + ): + strategy = WordPlayerStrategy("I love cats") + + assert strategy.get_players() == ["I", "love", "cats"] + assert strategy.n_players == 3 + + assert ( + strategy.coalition_to_text( + np.array([1, 0, 1]), + MaskTokenPerturbation(DummyTokenizer()), + ) + == "I [MASK] cats" + ) + + assert ( + strategy.coalition_to_text( + np.array([1, 0, 1]), + RemovalPerturbation(), + ) + == "I cats" + ) + + +def test_word_player_strategy_passes_context_to_perturbation( + no_nltk_resource_check, +) -> None: + with patch( + f"{PLAYERS_MODULE}.nltk.word_tokenize", + return_value=["I", "love", "cats"], + ): + strategy = WordPlayerStrategy("I love cats") + + perturbation = MagicMock() + perturbation.perturb.return_value = "X" + + assert ( + strategy.coalition_to_text( + np.array([1, 0, 1]), + perturbation, + ) + == "I X cats" + ) + + perturbation.perturb.assert_called_once() + args, kwargs = perturbation.perturb.call_args + assert args == ("love",) + context = kwargs["context"] + assert context["players"] == ["I", "love", "cats"] + np.testing.assert_array_equal( + context["coalition"], + np.array([1, 0, 1]), + ) + + assert context["mask_index"] == 1 + + +def test_named_entity_player_strategy_groups_entities( + no_nltk_resource_check, +) -> None: + ner_tree = [ + Tree("PERSON", [("John", "NNP"), ("Smith", "NNP")]), + ("visited", "VBD"), + Tree("GPE", [("Berlin", "NNP")]), + ] + + with ( + patch(f"{PLAYERS_MODULE}.nltk.word_tokenize", return_value=["ignored"]), + patch(f"{PLAYERS_MODULE}.nltk.pos_tag", return_value=[("ignored", "NN")]), + patch(f"{PLAYERS_MODULE}.nltk.ne_chunk", return_value=ner_tree), + ): + strategy = NamedEntityPlayerStrategy("John Smith visited Berlin") + + assert strategy.get_players() == ["John Smith", "visited", "Berlin"] + assert ( + strategy.coalition_to_text( + np.array([1, 0, 1]), + NeutralPerturbation("something"), + ) + == "John Smith something Berlin" + ) + + +def test_chunk_player_strategy_groups_phrases( + no_nltk_resource_check, +) -> None: + parsed_tree = [ + Tree("NP", [("the", "DT"), ("movie", "NN")]), + ("was", "VBD"), + Tree("ADJP", [("very", "RB"), ("good", "JJ")]), + ] + + parser = MagicMock() + parser.parse.return_value = parsed_tree + + with ( + patch(f"{PLAYERS_MODULE}.nltk.word_tokenize", return_value=["ignored"]), + patch(f"{PLAYERS_MODULE}.nltk.pos_tag", return_value=[("ignored", "NN")]), + patch(f"{PLAYERS_MODULE}.nltk.RegexpParser", return_value=parser), + ): + strategy = ChunkPlayerStrategy("the movie was very good") + + assert strategy.get_players() == ["the movie", "was", "very good"] + assert ( + strategy.coalition_to_text( + np.array([0, 1, 1]), + NeutralPerturbation("something"), + ) + == "something was very good" + ) + + +def test_sentence_player_strategy_with_mocked_nltk( + no_nltk_resource_check, +) -> None: + with patch( + f"{PLAYERS_MODULE}.nltk.sent_tokenize", + return_value=["First.", "Second."], + ): + strategy = SentencePlayerStrategy("First. Second.") + + assert strategy.get_players() == ["First.", "Second."] + assert ( + strategy.coalition_to_text( + np.array([1, 0]), + PadTokenPerturbation(DummyTokenizer()), + ) + == "First. [PAD]" + ) + + +def test_player_factory_creates_correct_strategy( + tokenizer: DummyTokenizer, + no_nltk_resource_check, +) -> None: + assert isinstance( + create_player_strategy("subword", "unhappy", tokenizer), + SubwordPlayerStrategy, + ) + + with patch(f"{PLAYERS_MODULE}.nltk.word_tokenize", return_value=["hello"]): + assert isinstance( + create_player_strategy("word", "hello", tokenizer), + WordPlayerStrategy, + ) + + with ( + patch( + f"{PLAYERS_MODULE}.NamedEntityPlayerStrategy", + return_value=MagicMock(), + ) as named_entity_strategy, + patch( + f"{PLAYERS_MODULE}.ChunkPlayerStrategy", + return_value=MagicMock(), + ) as chunk_strategy, + patch( + f"{PLAYERS_MODULE}.SentencePlayerStrategy", + return_value=MagicMock(), + ) as sentence_strategy, + ): + create_player_strategy("named_entity", "hello", tokenizer) + create_player_strategy("chunk", "hello", tokenizer) + create_player_strategy("sentence", "hello", tokenizer) + + named_entity_strategy.assert_called_once_with(text="hello") + chunk_strategy.assert_called_once_with(text="hello") + sentence_strategy.assert_called_once_with(text="hello") + + +def test_player_factory_rejects_unknown_level(tokenizer: DummyTokenizer) -> None: + with pytest.raises(ValueError, match="Unknown player level"): + create_player_strategy("not_real", "hello", tokenizer) + + +# ============================================================================ +# Basic perturbations and WordNet perturbation +# ============================================================================ + + +def test_mask_and_pad_perturbations(tokenizer: DummyTokenizer) -> None: + assert MaskTokenPerturbation(tokenizer).perturb("word") == "[MASK]" + assert PadTokenPerturbation(tokenizer).perturb("word") == "[PAD]" + assert RemovalPerturbation().perturb("word") == "" + assert NeutralPerturbation("neutral").perturb("word") == "neutral" + + +def test_mask_and_pad_require_special_tokens() -> None: + tokenizer_without_mask = DummyTokenizer() + tokenizer_without_mask.mask_token = None + + with pytest.raises(ValueError, match="does not define a mask token"): + MaskTokenPerturbation(tokenizer_without_mask) + + tokenizer_without_pad = DummyTokenizer() + tokenizer_without_pad.pad_token = None + + with pytest.raises(ValueError, match="does not define a pad token"): + PadTokenPerturbation(tokenizer_without_pad) + + +@pytest.mark.parametrize( + ("tag", "expected"), + [ + ("NN", "n"), + ("VBZ", "v"), + ("JJ", "a"), + ("RB", "r"), + ("IN", None), + ], +) +def test_penn_to_wordnet_mapping(tag: str, expected: str | None) -> None: + fake_wn = SimpleNamespace( + NOUN="n", + VERB="v", + ADJ="a", + ADV="r", + ) + + with patch(f"{PERTURBATIONS_MODULE}.wn", fake_wn): + assert _penn_to_wn(tag) == expected + + +def test_get_neutral_replacement_uses_hypernym() -> None: + hypernym = MagicMock() + hypernym.lemma_names.return_value = ["living_thing"] + + synset = MagicMock() + synset.hypernyms.return_value = [hypernym] + + fake_wn = SimpleNamespace( + NOUN="n", + VERB="v", + ADJ="a", + ADV="r", + synsets=MagicMock(return_value=[synset]), + ) + with patch(f"{PERTURBATIONS_MODULE}.wn", fake_wn): + assert _get_neutral_replacement("cat", "NN") == "living" + + +@pytest.mark.parametrize("tag", ["IN", "NN"]) +def test_get_neutral_replacement_falls_back_to_something(tag: str) -> None: + if tag == "IN": + assert _get_neutral_replacement("of", tag) == "something" + else: + fake_wn = SimpleNamespace( + NOUN="n", + VERB="v", + ADJ="a", + ADV="r", + synsets=MagicMock(return_value=[]), + ) + with patch(f"{PERTURBATIONS_MODULE}.wn", fake_wn): + assert _get_neutral_replacement("unknown", tag) == "something" + + +def test_get_neutral_replacement_falls_back_when_no_hypernym() -> None: + synset = MagicMock() + synset.hypernyms.return_value = [] + + fake_wn = SimpleNamespace( + NOUN="n", + VERB="v", + ADJ="a", + ADV="r", + synsets=MagicMock(return_value=[synset]), + ) + + with patch(f"{PERTURBATIONS_MODULE}.wn", fake_wn): + assert _get_neutral_replacement("cat", "NN") == "something" + + +def test_wordnet_neutral_perturbation( + no_nltk_resource_check, +) -> None: + with ( + patch(f"{PERTURBATIONS_MODULE}.nltk.pos_tag", return_value=[("cat", "NN")]), + patch( + f"{PERTURBATIONS_MODULE}._get_neutral_replacement", + return_value="animal", + ), + ): + result = WordNetNeutralPerturbation().perturb("cat") + + assert result == "animal" + + +def test_perturbation_factory(tokenizer: DummyTokenizer) -> None: + assert isinstance( + create_perturbation_strategy("mask", tokenizer), + MaskTokenPerturbation, + ) + assert isinstance( + create_perturbation_strategy("pad", tokenizer), + PadTokenPerturbation, + ) + assert isinstance( + create_perturbation_strategy("removal", tokenizer), + RemovalPerturbation, + ) + assert isinstance( + create_perturbation_strategy("neutral", tokenizer), + NeutralPerturbation, + ) + assert isinstance( + create_perturbation_strategy("wordnet_neutral", tokenizer), + WordNetNeutralPerturbation, + ) + + with pytest.raises(ValueError, match="Unknown perturbation strategy"): + create_perturbation_strategy("not_real", tokenizer) + + +def test_perturbation_factory_creates_mlm_infilling( + tokenizer: DummyTokenizer, +) -> None: + with patch( + f"{PERTURBATIONS_MODULE}.MLMInfillingPerturbation", + ) as mlm_strategy: + result = create_perturbation_strategy( + "mlm_infilling", + tokenizer, + mlm_model_name="fake-mlm", + mlm_num_samples=5, + device="cpu", + ) + + mlm_strategy.assert_called_once_with( + model_name="fake-mlm", + device="cpu", + num_samples=5, + ) + assert result is mlm_strategy.return_value + + +# ============================================================================ +# MLM infilling: all model calls are mocked +# ============================================================================ + + +def test_mlm_infilling_initializes_model_and_tokenizer() -> None: + tokenizer = MagicMock() + tokenizer.mask_token = "[MASK]" + + model = MagicMock() + model.to.return_value = model + + with ( + patch( + f"{PERTURBATIONS_MODULE}.AutoTokenizer.from_pretrained", + return_value=tokenizer, + ) as tokenizer_loader, + patch( + f"{PERTURBATIONS_MODULE}.AutoModelForMaskedLM.from_pretrained", + return_value=model, + ) as model_loader, + ): + perturbation = MLMInfillingPerturbation( + model_name="fake-mlm", + device="cpu", + num_samples=5, + ) + + tokenizer_loader.assert_called_once_with("fake-mlm") + model_loader.assert_called_once_with("fake-mlm") + model.to.assert_called_once_with("cpu") + model.eval.assert_called_once() + + assert perturbation.tokenizer is tokenizer + assert perturbation.model is model + assert perturbation.model_name == "fake-mlm" + assert perturbation.device == "cpu" + assert perturbation.mask_token == "[MASK]" + assert perturbation._cache == {} + assert perturbation.num_samples == 5 + + +def test_mlm_infilling_rejects_tokenizer_without_mask_token() -> None: + tokenizer = MagicMock() + tokenizer.mask_token = None + + model = MagicMock() + model.to.return_value = model + + with ( + patch( + f"{PERTURBATIONS_MODULE}.AutoTokenizer.from_pretrained", + return_value=tokenizer, + ), + patch( + f"{PERTURBATIONS_MODULE}.AutoModelForMaskedLM.from_pretrained", + return_value=model, + ), + pytest.raises(ValueError, match="does not define a mask token"), + ): + MLMInfillingPerturbation( + model_name="fake-mlm", + device="cpu", + ) + + +def make_mlm_without_constructor( + tokenizer: DummyTokenizer, + model: MagicMock, +) -> MLMInfillingPerturbation: + """Create an MLM perturbation without downloading a Hugging Face model.""" + perturbation = object.__new__(MLMInfillingPerturbation) + perturbation.tokenizer = tokenizer + perturbation.model = model + perturbation.model_name = "fake-mlm" + perturbation.device = "cpu" + perturbation.mask_token = "[MASK]" + perturbation._cache = {} + perturbation.num_samples = 3 + return perturbation + + +def test_mlm_infilling_requires_context( + tokenizer: DummyTokenizer, + model: MagicMock, +) -> None: + perturbation = make_mlm_without_constructor(tokenizer, model) + + with pytest.raises(ValueError, match="requires context"): + perturbation.perturb("movie") + + +def test_mlm_infilling_caches_one_prediction_per_coalition( + tokenizer: DummyTokenizer, + model: MagicMock, +) -> None: + perturbation = make_mlm_without_constructor(tokenizer, model) + perturbation._predict_masks = MagicMock( + return_value={1: "great"}, + ) + + context = { + "players": ["This", "movie", "works"], + "coalition": np.array([1, 0, 1]), + "mask_index": 1, + } + + assert perturbation.perturb("movie", context=context) == "great" + assert perturbation.perturb("movie", context=context) == "great" + + perturbation._predict_masks.assert_called_once() + perturbation.clear_cache() + assert perturbation._cache == {} + + +def test_mlm_infilling_returns_original_player_if_index_missing( + tokenizer: DummyTokenizer, + model: MagicMock, +) -> None: + perturbation = make_mlm_without_constructor(tokenizer, model) + perturbation._predict_masks = MagicMock(return_value={}) + + assert ( + perturbation.perturb( + "movie", + context={ + "players": ["movie"], + "coalition": np.array([0]), + "mask_index": 0, + }, + ) + == "movie" + ) + + +def test_mlm_predict_masks_filters_invalid_tokens( + tokenizer: DummyTokenizer, + model: MagicMock, +) -> None: + perturbation = make_mlm_without_constructor(tokenizer, model) + + logits = torch.zeros((1, 3, 10)) + model.return_value = SimpleNamespace(logits=logits) + + # Invalid candidates: [PAD], [CLS], [SEP], [MASK], ##suffix. + # First valid candidate is "great". + sampled_ids = iter([0, 1, 2, 3, 4, 5]) + + with patch( + f"{PERTURBATIONS_MODULE}.torch.multinomial", + side_effect=lambda *args, **kwargs: torch.tensor([next(sampled_ids)]), + ): + replacements = perturbation._predict_masks( + players=["This", "movie", "works"], + coalition=np.array([1, 0, 1]), + ) + + assert replacements == {1: "great"} + model.assert_called_once() + + +def test_mlm_predict_masks_falls_back_after_failed_sampling( + tokenizer: DummyTokenizer, + model: MagicMock, +) -> None: + perturbation = make_mlm_without_constructor(tokenizer, model) + + logits = torch.zeros((1, 3, 10)) + model.return_value = SimpleNamespace(logits=logits) + + with patch( + f"{PERTURBATIONS_MODULE}.torch.multinomial", + return_value=torch.tensor([0]), # always [PAD] + ): + replacements = perturbation._predict_masks( + players=["This", "movie", "works"], + coalition=np.array([1, 0, 1]), + ) + + assert replacements == {1: "something"} + + +# ============================================================================ +# Target callables +# ============================================================================ + + +class DummyTargetCallable(BaseTargetCallable): + """Minimal target callable used to test the default tensor-input path.""" + + def predict(self, texts: list[str]) -> np.ndarray: + return np.zeros(len(texts)) + + +def test_base_target_callable_rejects_pre_tokenized_inputs( + tokenizer: DummyTokenizer, + model: MagicMock, +) -> None: + callable_ = DummyTargetCallable( + model=model, + tokenizer=tokenizer, + device="cpu", + ) + + with pytest.raises(NotImplementedError, match="does not support pre-tokenized inputs"): + callable_.predict_from_inputs([]) + + +def test_encoder_callable_returns_logits( + tokenizer: DummyTokenizer, + model: MagicMock, +) -> None: + model.return_value = SimpleNamespace( + logits=torch.tensor([[1.0, 2.0], [3.0, 4.0]]), + ) + + callable_ = EncoderClassifierCallable( + model=model, + tokenizer=tokenizer, + device="cpu", + class_idx=1, + output_type="logit", + ) + + np.testing.assert_allclose( + callable_.predict(["a", "b"]), + np.array([2.0, 4.0]), + ) + + +def test_encoder_callable_returns_probabilities( + tokenizer: DummyTokenizer, + model: MagicMock, +) -> None: + model.return_value = SimpleNamespace( + logits=torch.tensor([[0.0, 0.0]]), + ) + + callable_ = EncoderClassifierCallable( + model=model, + tokenizer=tokenizer, + device="cpu", + class_idx=1, + output_type="probability", + ) + + np.testing.assert_allclose( + callable_.predict(["a"]), + np.array([0.5]), + ) + + +def test_encoder_callable_predicts_probabilities_from_inputs_with_token_type_ids( + tokenizer: DummyTokenizer, + model: MagicMock, +) -> None: + model.return_value = SimpleNamespace( + logits=torch.tensor( + [ + [0.0, 0.0], + [1.0, 2.0], + ] + ), + ) + + callable_ = EncoderClassifierCallable( + model=model, + tokenizer=tokenizer, + device="cpu", + class_idx=1, + output_type="probability", + ) + + inputs = [ + { + "input_ids": torch.tensor([[10, 11]]), + "attention_mask": torch.tensor([[1, 1]]), + "token_type_ids": torch.tensor([[0, 0]]), + }, + { + "input_ids": torch.tensor([[12, 13]]), + "attention_mask": torch.tensor([[1, 1]]), + "token_type_ids": torch.tensor([[1, 1]]), + }, + ] + + scores = callable_.predict_from_inputs(inputs) + + expected = torch.softmax( + torch.tensor( + [ + [0.0, 0.0], + [1.0, 2.0], + ] + ), + dim=-1, + )[:, 1].numpy() + + np.testing.assert_allclose(scores, expected) + + model.assert_called_once() + model_inputs = model.call_args.kwargs + + assert "token_type_ids" in model_inputs + assert model_inputs["token_type_ids"].shape == (2, 2) + + +def test_encoder_callable_rejects_invalid_output_type( + tokenizer: DummyTokenizer, + model: MagicMock, +) -> None: + with pytest.raises(ValueError, match="output_type"): + EncoderClassifierCallable( + model=model, + tokenizer=tokenizer, + device="cpu", + output_type="not_real", + ) + + +def test_causal_callable_scores_multi_token_target( + tokenizer: DummyTokenizer, + model: MagicMock, +) -> None: + # First encode call is target label; following calls are prompt encodings. + tokenizer.encode_return_values = [ + [5, 6], # target label + [10, 11], # prompt + ] + + logits = torch.zeros((1, 2, 20)) + logits[0, -1, 5] = 3.0 + logits[0, -1, 6] = 4.0 + model.return_value = SimpleNamespace(logits=logits) + + callable_ = CausalLMCallable( + model=model, + tokenizer=tokenizer, + device="cpu", + target_label="very good", + ) + + scores = callable_.predict(["nice review"]) + + assert scores.shape == (1,) + assert model.call_count == 2 + assert callable_._build_prompt("nice") == ("Review: nice\n\nSentiment:") + + +def test_causal_callable_uses_eos_as_pad_when_pad_is_missing( + model: MagicMock, +) -> None: + tokenizer = DummyTokenizer() + tokenizer.pad_token_id = None + tokenizer.encode_return_values = [[5]] + + CausalLMCallable( + model=model, + tokenizer=tokenizer, + device="cpu", + ) + + assert tokenizer.pad_token == "" + assert tokenizer.padding_side == "left" + + +def test_causal_callable_rejects_tokenizer_without_pad_or_eos( + model: MagicMock, +) -> None: + tokenizer = DummyTokenizer() + tokenizer.pad_token_id = None + tokenizer.eos_token_id = None + tokenizer.encode_return_values = [[5]] + + with pytest.raises( + ValueError, + match="Tokenizer must define either a pad token or eos token", + ): + CausalLMCallable( + model=model, + tokenizer=tokenizer, + device="cpu", + ) + + +def test_causal_callable_predicts_from_inputs_with_multi_token_target( + tokenizer: DummyTokenizer, + model: MagicMock, +) -> None: + tokenizer.encode_return_values = [[5, 6]] + + logits = torch.zeros((1, 3, 20)) + logits[0, -1, 5] = 3.0 + logits[0, -1, 6] = 4.0 + model.return_value = SimpleNamespace(logits=logits) + + callable_ = CausalLMCallable( + model=model, + tokenizer=tokenizer, + device="cpu", + target_label="very good", + ) + + inputs = [ + { + "input_ids": torch.tensor([[10, 11]]), + "attention_mask": torch.tensor([[1, 1]]), + } + ] + + scores = callable_.predict_from_inputs(inputs) + + assert scores.shape == (1,) + assert np.isfinite(scores[0]) + assert model.call_count == 2 + + second_call_inputs = model.call_args_list[1].kwargs + assert second_call_inputs["input_ids"].shape == (1, 3) + assert second_call_inputs["attention_mask"].shape == (1, 3) + + +def test_causal_callable_rejects_empty_target( + tokenizer: DummyTokenizer, + model: MagicMock, +) -> None: + tokenizer.encode_return_values = [[]] + + with pytest.raises(ValueError, match="produced no tokens"): + CausalLMCallable( + model=model, + tokenizer=tokenizer, + device="cpu", + ) + + +# ============================================================================ +# TextImputer orchestration +# ============================================================================ + + +def make_player_strategy() -> MagicMock: + strategy = MagicMock() + strategy.n_players = 2 + strategy.coalition_to_text.side_effect = [ + "full-text", + "empty-text", + "text-1", + "text-2", + "text-3", + ] + return strategy + + +def test_text_imputer_creates_default_player_strategy( + tokenizer: DummyTokenizer, + model: MagicMock, +) -> None: + player_strategy = make_player_strategy() + + with patch( + "shapiq.imputer.text.imputer.create_player_strategy", + return_value=player_strategy, + ) as create_strategy: + TextImputer( + model=model, + tokenizer=tokenizer, + text="original", + perturbation_strategy=NeutralPerturbation(), + ) + + create_strategy.assert_called_once_with( + level="word", + text="original", + tokenizer=tokenizer, + ) + + +def test_text_imputer_creates_default_perturbation_strategy( + tokenizer: DummyTokenizer, + model: MagicMock, +) -> None: + player_strategy = make_player_strategy() + perturbation = NeutralPerturbation() + + with patch( + "shapiq.imputer.text.imputer.create_perturbation_strategy", + return_value=perturbation, + ) as create_strategy: + TextImputer( + model=model, + tokenizer=tokenizer, + text="original", + player_strategy=player_strategy, + perturbation_type="neutral", + mlm_model_name="test-mlm", + mlm_num_samples=5, + device="cpu", + ) + + create_strategy.assert_called_once_with( + strategy="neutral", + tokenizer=tokenizer, + mlm_model_name="test-mlm", + mlm_num_samples=5, + device="cpu", + ) + + +def test_text_imputer_rejects_both_perturbation_strategies( + tokenizer: DummyTokenizer, + model: MagicMock, +) -> None: + player_strategy = make_player_strategy() + + with pytest.raises( + ValueError, + match="Only one of perturbation_strategy and tensor_perturbation_strategy", + ): + TextImputer( + model=model, + tokenizer=tokenizer, + text="original", + player_strategy=player_strategy, + perturbation_strategy=NeutralPerturbation(), + tensor_perturbation_strategy=MagicMock(), + ) + + +def test_text_imputer_rejects_text_strategy_for_tensor_perturbation( + tokenizer: DummyTokenizer, + model: MagicMock, +) -> None: + player_strategy = make_player_strategy() + + with pytest.raises(ValueError, match="is a tensor perturbation"): + TextImputer( + model=model, + tokenizer=tokenizer, + text="original", + player_strategy=player_strategy, + perturbation_type="attention_mask", + perturbation_strategy=NeutralPerturbation(), + ) + + +def test_text_imputer_rejects_tensor_strategy_for_text_perturbation( + tokenizer: DummyTokenizer, + model: MagicMock, +) -> None: + player_strategy = make_player_strategy() + + with pytest.raises(ValueError, match="is a text perturbation"): + TextImputer( + model=model, + tokenizer=tokenizer, + text="original", + player_strategy=player_strategy, + perturbation_type="mask", + tensor_perturbation_strategy=MagicMock(), + ) + + +def test_text_imputer_rejects_coalition_to_text_in_tensor_mode( + tokenizer: DummyTokenizer, + model: MagicMock, +) -> None: + player_strategy = make_player_strategy() + tensor_perturbation_strategy = MagicMock() + tensor_perturbation_strategy.evaluate.return_value = [{"input_ids": torch.tensor([[1, 2]])}] + + with patch.object( + EncoderClassifierCallable, + "predict_from_inputs", + return_value=np.array([0.5]), + ): + imputer = TextImputer( + model=model, + tokenizer=tokenizer, + text="original", + player_strategy=player_strategy, + perturbation_type="attention_mask", + tensor_perturbation_strategy=tensor_perturbation_strategy, + ) + + with pytest.raises( + RuntimeError, + match=r"coalition_to_text\(\) can only be used with text perturbation strategies", + ): + imputer.coalition_to_text(np.array([1, 0])) + + +def test_text_imputer_rejects_coalitions_to_texts_in_tensor_mode( + tokenizer: DummyTokenizer, + model: MagicMock, +) -> None: + player_strategy = make_player_strategy() + tensor_perturbation_strategy = MagicMock() + tensor_perturbation_strategy.evaluate.return_value = [{"input_ids": torch.tensor([[1, 2]])}] + + with patch.object( + EncoderClassifierCallable, + "predict_from_inputs", + return_value=np.array([0.5]), + ): + imputer = TextImputer( + model=model, + tokenizer=tokenizer, + text="original", + player_strategy=player_strategy, + perturbation_type="attention_mask", + tensor_perturbation_strategy=tensor_perturbation_strategy, + ) + + with pytest.raises( + RuntimeError, + match=r"_coalitions_to_texts\(\) can only be used with text perturbation strategies", + ): + imputer._coalitions_to_texts(np.array([[1, 0]])) + + +def test_text_imputer_batches_and_returns_scores( + tokenizer: DummyTokenizer, + model: MagicMock, +) -> None: + player_strategy = make_player_strategy() + perturbation = NeutralPerturbation() + + imputer = TextImputer( + model=model, + tokenizer=tokenizer, + text="original", + batch_size=2, + player_strategy=player_strategy, + perturbation_strategy=perturbation, + ) + + imputer.target_callable = MagicMock() + imputer.target_callable.predict.side_effect = [ + np.array([0.1, 0.2]), + np.array([0.3]), + ] + + scores = imputer.value_function( + np.array([[1, 0], [0, 1], [0, 0]]), + ) + + np.testing.assert_allclose(scores, np.array([0.1, 0.2, imputer.empty_prediction])) + assert imputer.target_callable.predict.call_args_list == [ + call(["text-1", "text-2"]), + call(["text-3"]), + ] + + +def test_text_imputer_accepts_one_dimensional_coalition( + tokenizer: DummyTokenizer, + model: MagicMock, +) -> None: + player_strategy = make_player_strategy() + + imputer = TextImputer( + model=model, + tokenizer=tokenizer, + text="original", + player_strategy=player_strategy, + perturbation_strategy=NeutralPerturbation(), + ) + + imputer.target_callable = MagicMock() + imputer.target_callable.predict.return_value = np.array([0.7]) + + np.testing.assert_allclose( + imputer.value_function(np.array([1, 0])), + np.array([0.7]), + ) + + +def test_text_imputer_rejects_wrong_coalition_width( + tokenizer: DummyTokenizer, + model: MagicMock, +) -> None: + player_strategy = make_player_strategy() + + imputer = TextImputer( + model=model, + tokenizer=tokenizer, + text="original", + player_strategy=player_strategy, + perturbation_strategy=NeutralPerturbation(), + ) + + with pytest.raises(ValueError, match="Expected coalition width 2"): + imputer.value_function(np.array([[1, 0, 1]])) + + +def test_text_imputer_mlm_averages_multiple_samples( + tokenizer: DummyTokenizer, + model: MagicMock, +) -> None: + player_strategy = MagicMock() + player_strategy.n_players = 1 + player_strategy.coalition_to_text.side_effect = [ + # full prediction + "full-1", + "full-2", + "full-3", + # empty prediction + "empty-1", + "empty-2", + "empty-3", + # value_function + "sample-1", + "sample-2", + "sample-3", + ] + + mlm = make_mlm_without_constructor(tokenizer, model) + mlm.num_samples = 3 + mlm.clear_cache = MagicMock() + + imputer = TextImputer( + model=model, + tokenizer=tokenizer, + text="original", + player_strategy=player_strategy, + perturbation_strategy=mlm, + player_level="word", + ) + + imputer.target_callable = MagicMock() + imputer.target_callable.predict.side_effect = [ + np.array([1.0]), + np.array([2.0]), + np.array([3.0]), + ] + + np.testing.assert_allclose( + imputer.value_function(np.array([[0]])), + np.array([2.0]), + ) + + assert mlm.clear_cache.call_count == mlm.num_samples * 3 + assert imputer._last_generated_texts == [ + "sample-1", + "sample-2", + "sample-3", + ] + + +@pytest.mark.parametrize("player_level", ["subword", "sentence"]) +def test_text_imputer_rejects_unsupported_mlm_player_levels( + tokenizer: DummyTokenizer, + model: MagicMock, + player_level: str, +) -> None: + mlm = make_mlm_without_constructor(tokenizer, model) + player_strategy = MagicMock() + player_strategy.n_players = 1 + + with pytest.raises(ValueError, match="supports only word"): + TextImputer( + model=model, + tokenizer=tokenizer, + text="original", + player_level=player_level, + player_strategy=player_strategy, + perturbation_strategy=mlm, + ) + + +def test_text_imputer_rejects_unknown_model_type( + tokenizer: DummyTokenizer, + model: MagicMock, +) -> None: + player_strategy = MagicMock() + player_strategy.n_players = 1 + + with pytest.raises(ValueError, match="model_type must be one of"): + TextImputer( + model=model, + tokenizer=tokenizer, + text="original", + player_strategy=player_strategy, + perturbation_strategy=NeutralPerturbation(), + model_type="not_real", + ) + + +def test_text_imputer_full_prediction_and_call( + tokenizer: DummyTokenizer, + model: MagicMock, +) -> None: + player_strategy = MagicMock() + player_strategy.n_players = 1 + player_strategy.coalition_to_text.return_value = "perturbed" + + imputer = TextImputer( + model=model, + tokenizer=tokenizer, + text="original", + player_strategy=player_strategy, + perturbation_strategy=NeutralPerturbation(), + ) + + assert imputer.full_prediction == 2.0 + + +@pytest.mark.slow +@pytest.mark.skipif( + os.environ.get("RUN_SLOW_TESTS") != "1", + reason="Set RUN_SLOW_TESTS=1 to run slow end-to-end tests.", +) +def test_text_imputer_end_to_end_with_tiny_checkpoint() -> None: + """Run TextImputer end-to-end with a real tiny Hugging Face checkpoint.""" + model_name = "hf-internal-testing/tiny-random-bert" + + tokenizer = AutoTokenizer.from_pretrained(model_name) + model = AutoModelForSequenceClassification.from_pretrained(model_name) + + imputer = TextImputer( + model=model, + tokenizer=tokenizer, + text="This movie is surprisingly good.", + player_level="subword", + perturbation_type="mask", + model_type="encoder_classifier", + class_idx=1, + output_type="logit", + device="cpu", + ) + + coalitions = np.stack( + [ + imputer.empty_coalition, + imputer.grand_coalition, + ] + ) + + scores = imputer(coalitions) + + assert scores.shape == (2,) + assert np.all(np.isfinite(scores)) + assert scores[0] == pytest.approx(0.0) diff --git a/tests/shapiq/tests_unit/tests_imputer/test_text_imputer_attention.py b/tests/shapiq/tests_unit/tests_imputer/test_text_imputer_attention.py new file mode 100644 index 000000000..7127b4655 --- /dev/null +++ b/tests/shapiq/tests_unit/tests_imputer/test_text_imputer_attention.py @@ -0,0 +1,450 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import numpy as np +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("transformers") + +from shapiq.imputer.text_imputer import ( # noqa: E402 + AttentionMaskPerturbation, + TextImputer, + create_tensor_perturbation_strategy, +) + + +class TinyTokenizer: + """Tiny tokenizer for attention-mask unit tests. + + It avoids downloading HuggingFace models/tokenizers while still exposing the + methods and attributes used by TextImputer. + """ + + def __init__(self) -> None: + self.vocab: dict[str, int] = { + "": 0, + "": 1, + } + self.inv_vocab: dict[int, str] = { + 0: "", + 1: "", + } + self.eos_token = "" + self.eos_token_id = 0 + self.pad_token = "" + self.pad_token_id = 1 + self.mask_token = "[MASK]" + self.mask_token_id = 2 + self.padding_side = "right" + + def encode( + self, + text: str, + *, + add_special_tokens: bool = False, + ) -> list[int]: + """Encode text by whitespace tokens.""" + token_ids: list[int] = [] + + for token in text.split(): + if token not in self.vocab: + token_id = len(self.vocab) + self.vocab[token] = token_id + self.inv_vocab[token_id] = token + + token_ids.append(self.vocab[token]) + + return token_ids + + def tokenize(self, text: str) -> list[str]: + """Return whitespace tokens.""" + return text.split() + + def decode( + self, + token_ids: list[int] | torch.Tensor, + *, + skip_special_tokens: bool = False, + ) -> str: + """Decode token ids.""" + if isinstance(token_ids, torch.Tensor): + token_ids = token_ids.detach().cpu().tolist() + + return " ".join(self.inv_vocab[int(token_id)] for token_id in token_ids) + + def convert_ids_to_tokens( + self, + token_ids: list[int] | torch.Tensor, + ) -> list[str]: + """Convert token ids to token strings.""" + if isinstance(token_ids, torch.Tensor): + token_ids = token_ids.detach().cpu().tolist() + + return [self.inv_vocab[int(token_id)] for token_id in token_ids] + + +class StaticPlayerStrategy: + """Player strategy with predefined players.""" + + def __init__(self, players: list[str]) -> None: + self.players = players + + def get_players(self) -> list[str]: + return self.players + + @property + def n_players(self) -> int: + return len(self.players) + + def coalition_to_text(self, coalition: np.ndarray, perturbation_strategy) -> str: + output: list[str] = [] + + for keep, player in zip(coalition, self.players, strict=False): + if keep: + output.append(player) + else: + output.append(perturbation_strategy.perturb(player)) + + return " ".join(output) + + +class FakeEncoderClassifier(torch.nn.Module): + """Fake encoder classifier whose score depends on visible tokens.""" + + def forward( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + **_: object, + ) -> SimpleNamespace: + visible_count = attention_mask.float().sum(dim=1) + logits = torch.stack([-visible_count, visible_count], dim=1) + + return SimpleNamespace(logits=logits) + + +class FakeCausalLM(torch.nn.Module): + """Fake causal LM whose target score depends on visible prompt tokens.""" + + def __init__(self, target_token_id: int, vocab_size: int = 128) -> None: + super().__init__() + self.target_token_id = target_token_id + self.vocab_size = vocab_size + + def forward( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + **_: object, + ) -> SimpleNamespace: + batch_size, sequence_length = input_ids.shape + logits = torch.zeros( + batch_size, + sequence_length, + self.vocab_size, + device=input_ids.device, + ) + + visible_count = attention_mask.float().sum(dim=1) + logits[:, :, self.target_token_id] = visible_count[:, None] + + return SimpleNamespace(logits=logits) + + +def test_attention_mask_builds_inputs_for_one_coalition() -> None: + """Attention masking keeps input_ids fixed and masks missing player spans.""" + tokenizer = TinyTokenizer() + players = ["Paris", "is", "beautiful", "."] + + coalition = np.array([False, True, False, True], dtype=bool) + + masked_inputs = AttentionMaskPerturbation.build_inputs_for_coalitions( + tokenizer=tokenizer, + players=players, + coalitions=coalition, + player_separator=" ", + ) + + encoded = masked_inputs[0] + + input_ids = encoded["input_ids"].tolist()[0] + attention_mask = encoded["attention_mask"].tolist()[0] + + assert tokenizer.decode(input_ids) == "Paris is beautiful ." + assert attention_mask == [0, 1, 0, 1] + + +def test_attention_mask_builds_inputs_for_batch_coalitions() -> None: + """Attention masking supports a matrix of coalitions.""" + tokenizer = TinyTokenizer() + players = ["Paris", "is", "beautiful", "."] + + coalitions = np.array( + [ + [True, True, True, True], + [False, False, True, True], + ], + dtype=bool, + ) + + masked_inputs = AttentionMaskPerturbation.build_inputs_for_coalitions( + tokenizer=tokenizer, + players=players, + coalitions=coalitions, + player_separator=" ", + ) + + assert len(masked_inputs) == 2 + assert masked_inputs[0]["attention_mask"].tolist()[0] == [1, 1, 1, 1] + assert masked_inputs[1]["attention_mask"].tolist()[0] == [0, 0, 1, 1] + + # The token ids stay unchanged across attention-masked coalitions. + assert masked_inputs[0]["input_ids"].tolist() == masked_inputs[1]["input_ids"].tolist() + + +def test_attention_mask_prompt_keeps_prompt_tokens_visible() -> None: + """Prompt tokens should remain visible while player tokens are maskable.""" + tokenizer = TinyTokenizer() + players = ["Paris", "is", "beautiful"] + + coalition = np.array([False, True, True], dtype=bool) + + masked_inputs = AttentionMaskPerturbation.build_prompt_inputs_for_coalitions( + tokenizer=tokenizer, + players=players, + coalitions=coalition, + prompt_template="Question: {text} Answer:", + player_separator=" ", + ) + + encoded = masked_inputs[0] + + tokens = tokenizer.convert_ids_to_tokens(encoded["input_ids"][0]) + attention_mask = encoded["attention_mask"].tolist()[0] + + assert tokens == ["Question:", "Paris", "is", "beautiful", "Answer:"] + assert attention_mask == [1, 0, 1, 1, 1] + + +def test_attention_mask_wrong_coalition_width_raises() -> None: + """Coalition width must match the number of players.""" + tokenizer = TinyTokenizer() + players = ["Paris", "is", "beautiful"] + + wrong_coalition = np.array([[True, False]], dtype=bool) + + with pytest.raises(ValueError, match="Expected coalition width"): + AttentionMaskPerturbation.build_inputs_for_coalitions( + tokenizer=tokenizer, + players=players, + coalitions=wrong_coalition, + player_separator=" ", + ) + + +def test_text_imputer_attention_mask_encoder_path_returns_scores() -> None: + """TextImputer should score attention-masked encoder inputs.""" + tokenizer = TinyTokenizer() + model = FakeEncoderClassifier() + player_strategy = StaticPlayerStrategy(["Paris", "is", "beautiful", "."]) + + imputer = TextImputer( + model=model, + tokenizer=tokenizer, + text="Paris is beautiful.", + player_strategy=player_strategy, + perturbation_type="attention_mask", + model_type="encoder_classifier", + class_idx=1, + output_type="logit", + batch_size=2, + device="cpu", + ) + + coalitions = np.array( + [ + [True, True, True, True], + [False, False, True, True], + ], + dtype=bool, + ) + + scores = imputer(coalitions) + + assert scores.shape == (2,) + assert scores[0] > scores[1] + + +def test_text_imputer_attention_mask_causal_lm_path_returns_scores() -> None: + """TextImputer should score attention-masked causal-LM prompt inputs.""" + tokenizer = TinyTokenizer() + target_token_id = tokenizer.encode("yes")[0] + model = FakeCausalLM(target_token_id=target_token_id) + player_strategy = StaticPlayerStrategy(["Paris", "is", "beautiful"]) + + imputer = TextImputer( + model=model, + tokenizer=tokenizer, + text="Paris is beautiful.", + player_strategy=player_strategy, + perturbation_type="attention_mask", + model_type="causal_lm", + target_label="yes", + prompt_template="Question: {text} Answer:", + batch_size=2, + device="cpu", + ) + + coalitions = np.array( + [ + [True, True, True], + [False, False, True], + ], + dtype=bool, + ) + + scores = imputer(coalitions) + + assert scores.shape == (2,) + assert np.all(np.isfinite(scores)) + assert scores[0] > scores[1] + + +def test_attention_mask_full_prediction_uses_full_coalition() -> None: + """full_prediction should work for attention-mask perturbation.""" + tokenizer = TinyTokenizer() + model = FakeEncoderClassifier() + player_strategy = StaticPlayerStrategy(["Paris", "is", "beautiful", "."]) + + imputer = TextImputer( + model=model, + tokenizer=tokenizer, + text="Paris is beautiful.", + player_strategy=player_strategy, + perturbation_type="attention_mask", + model_type="encoder_classifier", + class_idx=1, + output_type="logit", + batch_size=2, + device="cpu", + ) + + score = imputer.full_prediction + + assert isinstance(score, float) + assert np.isfinite(score) + + +def test_attention_mask_rejects_mismatched_coalition_and_player_spans() -> None: + """Coalition length must match the number of player spans.""" + base_attention_mask = torch.ones((1, 3), dtype=torch.long) + player_spans = [(0, 1), (1, 2), (2, 3)] + coalition = np.array([True, False], dtype=bool) + + with pytest.raises(ValueError, match="does not match number of player spans"): + AttentionMaskPerturbation.build_attention_mask_for_coalition( + base_attention_mask=base_attention_mask, + player_spans=player_spans, + coalition=coalition, + ) + + +def test_attention_mask_prompt_template_requires_text_placeholder() -> None: + """Prompt templates must contain the text placeholder.""" + tokenizer = TinyTokenizer() + players = ["Paris", "is", "beautiful"] + coalition = np.array([True, False, True], dtype=bool) + + with pytest.raises(ValueError, match=r"prompt_template must contain '\{text\}'"): + AttentionMaskPerturbation.build_prompt_inputs_for_coalitions( + tokenizer=tokenizer, + players=players, + coalitions=coalition, + prompt_template="Question: Paris is beautiful Answer:", + player_separator=" ", + ) + + +def test_attention_mask_prompt_rejects_wrong_coalition_width() -> None: + """Prompt coalition width must match the number of players.""" + tokenizer = TinyTokenizer() + players = ["Paris", "is", "beautiful"] + wrong_coalition = np.array([[True, False]], dtype=bool) + + with pytest.raises(ValueError, match="Expected coalition width"): + AttentionMaskPerturbation.build_prompt_inputs_for_coalitions( + tokenizer=tokenizer, + players=players, + coalitions=wrong_coalition, + prompt_template="Question: {text} Answer:", + player_separator=" ", + ) + + +def test_attention_mask_causal_lm_requires_prompt_template() -> None: + """Causal-LM attention masking requires a prompt template.""" + tokenizer = TinyTokenizer() + perturbation = AttentionMaskPerturbation(tokenizer=tokenizer) + + with pytest.raises(ValueError, match="prompt_template is required for causal_lm"): + perturbation.evaluate( + players=["Paris", "is", "beautiful"], + coalitions=np.array([[True, False, True]], dtype=bool), + model_type="causal_lm", + ) + + +def test_attention_mask_seq2seq_requires_prompt_template() -> None: + """Seq2seq attention masking requires a prompt template.""" + tokenizer = TinyTokenizer() + perturbation = AttentionMaskPerturbation(tokenizer=tokenizer) + + with pytest.raises(ValueError, match="prompt_template is required for seq2seq"): + perturbation.evaluate( + players=["Paris", "is", "beautiful"], + coalitions=np.array([[True, False, True]], dtype=bool), + model_type="seq2seq", + ) + + +def test_attention_mask_seq2seq_builds_prompt_inputs() -> None: + """Seq2seq attention masking should build prompt-based masked inputs.""" + tokenizer = TinyTokenizer() + perturbation = AttentionMaskPerturbation(tokenizer=tokenizer) + + masked_inputs = perturbation.evaluate( + players=["Paris", "is", "beautiful"], + coalitions=np.array([[False, True, True]], dtype=bool), + model_type="seq2seq", + prompt_template="Question: {text} Answer:", + player_separator=" ", + ) + + assert len(masked_inputs) == 1 + assert masked_inputs[0]["attention_mask"].tolist()[0] == [1, 0, 1, 1, 1] + + +def test_attention_mask_rejects_unknown_model_type() -> None: + """Attention masking should reject unsupported model types.""" + tokenizer = TinyTokenizer() + perturbation = AttentionMaskPerturbation(tokenizer=tokenizer) + + with pytest.raises(ValueError, match="Unknown model_type for attention masking"): + perturbation.evaluate( + players=["Paris"], + coalitions=np.array([[True]], dtype=bool), + model_type="unsupported", + ) + + +def test_create_tensor_perturbation_strategy_rejects_unknown_strategy() -> None: + """Tensor perturbation factory should reject unknown strategies.""" + tokenizer = TinyTokenizer() + + with pytest.raises(ValueError, match="Unknown tensor perturbation strategy"): + create_tensor_perturbation_strategy( + "unsupported", + tokenizer=tokenizer, + ) diff --git a/tests/shapiq/tests_unit/tests_imputer/test_text_imputer_seq2seq.py b/tests/shapiq/tests_unit/tests_imputer/test_text_imputer_seq2seq.py new file mode 100644 index 000000000..f3c5dc7d1 --- /dev/null +++ b/tests/shapiq/tests_unit/tests_imputer/test_text_imputer_seq2seq.py @@ -0,0 +1,881 @@ +# ============================================================================ +# Pytest unit tests for Seq2SeqCallable +# Coverage: model type validation, single-token target, multi-token +# teacher forcing, normalisation, prompt template, end-to-end +# integration with TextImputer +# ============================================================================ +from __future__ import annotations + +import os + +os.environ["TOKENIZERS_PARALLELISM"] = "false" + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import numpy as np +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("transformers") + +from shapiq.imputer.text_imputer import ( # noqa: E402 + NeutralPerturbation, + Seq2SeqCallable, + TextImputer, +) + +MODULE = "shapiq.imputer.text_imputer" + + +# ============================================================================ +# Mock fixtures — no real model downloads; all model calls go through MagicMock +# ============================================================================ + + +def make_seq2seq_tokenizer() -> MagicMock: + """Return a minimal seq2seq tokenizer substitute. + + encode() return values are controlled via the encode_queue list: + each call to encode() pops from the front of the queue. + """ + tok = MagicMock() + tok.pad_token = "[PAD]" + tok.pad_token_id = 0 + tok.eos_token = "" + tok.eos_token_id = 2 + + tok.encode_queue = [] + tok.encode.side_effect = lambda text, **kwargs: tok.encode_queue.pop(0) + + tok.return_value = { + "input_ids": torch.tensor([[10, 11, 12]]), + "attention_mask": torch.ones((1, 3), dtype=torch.long), + } + return tok + + +def make_seq2seq_model(decoder_start_token_id: int = 0) -> MagicMock: + """Return a minimal seq2seq model substitute. + + config.is_encoder_decoder is set to True to mimic T5 / BART. + model.get_encoder() returns an encoder mock whose output is a + SimpleNamespace with a last_hidden_state tensor. + The return value of model(**kwargs) can be overridden per test. + """ + model = MagicMock() + model.to.return_value = model + + model.config.is_encoder_decoder = True + model.config.decoder_start_token_id = decoder_start_token_id + + encoder_mock = MagicMock() + encoder_mock.return_value = SimpleNamespace( + last_hidden_state=torch.zeros((1, 3, 16)), + ) + model.get_encoder.return_value = encoder_mock + + return model + + +@pytest.fixture +def seq2seq_tokenizer() -> MagicMock: + return make_seq2seq_tokenizer() + + +@pytest.fixture +def seq2seq_model() -> MagicMock: + return make_seq2seq_model() + + +# ============================================================================ +# TEST 1 — Model type validation +# ============================================================================ +# Seq2SeqCallable.__init__ reads model.config.is_encoder_decoder. +# If the flag is False or absent, a ValueError mentioning +# "is_encoder_decoder" must be raised. +# ============================================================================ + + +class TestModelTypeValidation: + def test_rejects_model_with_is_encoder_decoder_false( + self, + seq2seq_tokenizer: MagicMock, + ) -> None: + """ValueError must be raised when is_encoder_decoder=False.""" + seq2seq_tokenizer.encode_queue = [[1]] + + bad_model = make_seq2seq_model() + bad_model.config.is_encoder_decoder = False + + with pytest.raises(ValueError, match="is_encoder_decoder"): + Seq2SeqCallable( + model=bad_model, + tokenizer=seq2seq_tokenizer, + device="cpu", + ) + + def test_rejects_model_without_is_encoder_decoder_attribute( + self, + seq2seq_tokenizer: MagicMock, + ) -> None: + """When the config attribute is absent, getattr defaults to False and + the constructor must raise ValueError.""" + seq2seq_tokenizer.encode_queue = [[1]] + + bad_model = make_seq2seq_model() + del bad_model.config.is_encoder_decoder + + with pytest.raises(ValueError, match="is_encoder_decoder"): + Seq2SeqCallable( + model=bad_model, + tokenizer=seq2seq_tokenizer, + device="cpu", + ) + + def test_accepts_valid_seq2seq_model( + self, + seq2seq_model: MagicMock, + seq2seq_tokenizer: MagicMock, + ) -> None: + """No exception must be raised when is_encoder_decoder=True.""" + seq2seq_tokenizer.encode_queue = [[1]] + + Seq2SeqCallable( + model=seq2seq_model, + tokenizer=seq2seq_tokenizer, + device="cpu", + ) + + def test_rejects_empty_target_label( + self, + seq2seq_model: MagicMock, + seq2seq_tokenizer: MagicMock, + ) -> None: + """ValueError must be raised when the target label encodes to an empty list.""" + seq2seq_tokenizer.encode_queue = [[]] + + with pytest.raises(ValueError, match="produced no tokens"): + Seq2SeqCallable( + model=seq2seq_model, + tokenizer=seq2seq_tokenizer, + device="cpu", + target_label="", + ) + + def test_rejects_when_no_decoder_start_token_available( + self, + seq2seq_tokenizer: MagicMock, + ) -> None: + """ValueError must be raised when both decoder_start_token_id and + pad_token_id are None.""" + seq2seq_tokenizer.encode_queue = [[1]] + seq2seq_tokenizer.pad_token_id = None + + bad_model = make_seq2seq_model() + bad_model.config.decoder_start_token_id = None + + with pytest.raises(ValueError, match="decoder_start_token_id"): + Seq2SeqCallable( + model=bad_model, + tokenizer=seq2seq_tokenizer, + device="cpu", + ) + + def test_falls_back_to_pad_token_id_when_config_missing( + self, + seq2seq_tokenizer: MagicMock, + ) -> None: + """When config.decoder_start_token_id is None, tokenizer.pad_token_id + must be used as the fallback.""" + seq2seq_tokenizer.encode_queue = [[1]] + seq2seq_tokenizer.pad_token_id = 7 + + model = make_seq2seq_model() + model.config.decoder_start_token_id = None + + callable_obj = Seq2SeqCallable( + model=model, + tokenizer=seq2seq_tokenizer, + device="cpu", + ) + + assert callable_obj.decoder_start_token_id == 7 + + +# ============================================================================ +# TEST 2 — Single-token target: output shape and dtype +# ============================================================================ +# predict([text]) must return a numpy array of shape (1,) and dtype float32. +# The scalar value must be a finite negative number (log-probability). +# ============================================================================ + + +class TestSingleTokenTarget: + def _make_callable( + self, + seq2seq_model: MagicMock, + seq2seq_tokenizer: MagicMock, + token_id: int = 42, + logit_value: float = 2.0, + *, + normalize: bool = True, + ) -> Seq2SeqCallable: + """Build a Seq2SeqCallable with a single-token target. + + logit_value is placed at the target token position; all other + logits are zero. Because log_softmax(2.0) < 0, the score is + always negative. + """ + seq2seq_tokenizer.encode_queue = [[token_id]] + + vocab_size = 100 + logits = torch.zeros((1, 1, vocab_size)) + logits[0, 0, token_id] = logit_value + seq2seq_model.return_value = SimpleNamespace(logits=logits) + + return Seq2SeqCallable( + model=seq2seq_model, + tokenizer=seq2seq_tokenizer, + device="cpu", + target_label="positive", + normalize=normalize, + ) + + def test_output_is_numpy_array( + self, + seq2seq_model: MagicMock, + seq2seq_tokenizer: MagicMock, + ) -> None: + callable_obj = self._make_callable(seq2seq_model, seq2seq_tokenizer) + scores = callable_obj.predict(["text"]) + assert isinstance(scores, np.ndarray) + + def test_output_shape_is_one( + self, + seq2seq_model: MagicMock, + seq2seq_tokenizer: MagicMock, + ) -> None: + callable_obj = self._make_callable(seq2seq_model, seq2seq_tokenizer) + scores = callable_obj.predict(["text"]) + assert scores.shape == (1,) + + def test_output_dtype_is_float32( + self, + seq2seq_model: MagicMock, + seq2seq_tokenizer: MagicMock, + ) -> None: + callable_obj = self._make_callable(seq2seq_model, seq2seq_tokenizer) + scores = callable_obj.predict(["text"]) + assert scores.dtype == np.float32 + + def test_output_value_is_finite( + self, + seq2seq_model: MagicMock, + seq2seq_tokenizer: MagicMock, + ) -> None: + callable_obj = self._make_callable(seq2seq_model, seq2seq_tokenizer) + scores = callable_obj.predict(["text"]) + assert np.isfinite(scores[0]) + + def test_output_value_is_negative( + self, + seq2seq_model: MagicMock, + seq2seq_tokenizer: MagicMock, + ) -> None: + """A log-probability must always be non-positive.""" + callable_obj = self._make_callable(seq2seq_model, seq2seq_tokenizer) + scores = callable_obj.predict(["text"]) + assert scores[0] < 0 + + def test_batch_output_shape_matches_input_length( + self, + seq2seq_model: MagicMock, + seq2seq_tokenizer: MagicMock, + ) -> None: + """When predict receives N texts, the output shape must be (N,).""" + token_id = 42 + n_texts = 3 + vocab_size = 100 + + seq2seq_tokenizer.encode_queue = [[token_id]] + seq2seq_tokenizer.return_value = { + "input_ids": torch.tensor([[10, 11, 12]] * n_texts), + "attention_mask": torch.ones((n_texts, 3), dtype=torch.long), + } + + logits = torch.zeros((n_texts, 1, vocab_size)) + logits[:, 0, token_id] = 2.0 + seq2seq_model.return_value = SimpleNamespace(logits=logits) + + callable_obj = Seq2SeqCallable( + model=seq2seq_model, + tokenizer=seq2seq_tokenizer, + device="cpu", + target_label="positive", + ) + + scores = callable_obj.predict(["a", "b", "c"]) + assert scores.shape == (n_texts,) + + +# ============================================================================ +# TEST 3 — Multi-token target: teacher-forcing loop +# ============================================================================ +# When the target label contains N tokens, the decoder loop runs N times. +# At each step decoder_input_ids grows by one token, and the final score +# equals the sum of per-token log-probabilities. +# ============================================================================ + + +class TestMultiTokenTarget: + def _make_two_token_callable( + self, + seq2seq_model: MagicMock, + seq2seq_tokenizer: MagicMock, + token_ids: list[int], + logit_value: float = 3.0, + *, + normalize: bool = False, + ) -> Seq2SeqCallable: + """Build a callable with a two-token target. + + Every model() call returns the same logits tensor; + the target token positions are set to logit_value, the rest to zero. + """ + seq2seq_tokenizer.encode_queue = [token_ids] + + vocab_size = 100 + logits = torch.zeros((1, len(token_ids), vocab_size)) + for tid in token_ids: + logits[:, :, tid] = logit_value + seq2seq_model.return_value = SimpleNamespace(logits=logits) + + return Seq2SeqCallable( + model=seq2seq_model, + tokenizer=seq2seq_tokenizer, + device="cpu", + target_label="very positive", + normalize=normalize, + ) + + def test_model_called_once_per_target_token( + self, + seq2seq_model: MagicMock, + seq2seq_tokenizer: MagicMock, + ) -> None: + """model() must be called exactly N times for an N-token target.""" + token_ids = [10, 20] + callable_obj = self._make_two_token_callable(seq2seq_model, seq2seq_tokenizer, token_ids) + + callable_obj.predict(["text"]) + + assert seq2seq_model.call_count == len(token_ids) + + def test_scores_are_summed_across_tokens( + self, + seq2seq_model: MagicMock, + seq2seq_tokenizer: MagicMock, + ) -> None: + """With normalize=False the score must equal the sum of per-token log-probs.""" + token_ids = [10, 20] + logit_value = 3.0 + vocab_size = 100 + + seq2seq_tokenizer.encode_queue = [token_ids] + + def make_logits(tid: int) -> SimpleNamespace: + logits = torch.zeros((1, 1, vocab_size)) + logits[0, 0, tid] = logit_value + return SimpleNamespace(logits=logits) + + seq2seq_model.side_effect = [make_logits(10), make_logits(20)] + + callable_obj = Seq2SeqCallable( + model=seq2seq_model, + tokenizer=seq2seq_tokenizer, + device="cpu", + target_label="very positive", + normalize=False, + ) + + scores = callable_obj.predict(["text"]) + + ref = 0.0 + for tid in token_ids: + logits = torch.zeros(vocab_size) + logits[tid] = logit_value + ref += torch.log_softmax(logits, dim=-1)[tid].item() + + assert abs(float(scores[0]) - ref) < 1e-5 + + def test_decoder_input_ids_grow_by_one_per_step( + self, + seq2seq_model: MagicMock, + seq2seq_tokenizer: MagicMock, + ) -> None: + """The length of decoder_input_ids passed to model() must increase by + one at each teacher-forcing step.""" + token_ids = [10, 20, 30] + vocab_size = 100 + + seq2seq_tokenizer.encode_queue = [token_ids] + + captured_decoder_lengths: list[int] = [] + + def capture_call(**kwargs) -> SimpleNamespace: + length = kwargs["decoder_input_ids"].shape[1] + captured_decoder_lengths.append(length) + logits = torch.zeros((1, length, vocab_size)) + return SimpleNamespace(logits=logits) + + seq2seq_model.side_effect = capture_call + + callable_obj = Seq2SeqCallable( + model=seq2seq_model, + tokenizer=seq2seq_tokenizer, + device="cpu", + target_label="three tokens", + normalize=False, + ) + + callable_obj.predict(["text"]) + + # Step 1: decoder_input_ids = [start] → length 1 + # Step 2: decoder_input_ids = [start, t1] → length 2 + # Step 3: decoder_input_ids = [start, t1, t2] → length 3 + assert captured_decoder_lengths == list(range(1, len(token_ids) + 1)) + + +# ============================================================================ +# TEST 4 — Normalisation: normalize=True vs normalize=False +# ============================================================================ +# normalize=False → score = sum of per-token log-probs +# normalize=True → score = mean of per-token log-probs = sum / n_tokens +# ============================================================================ + + +class TestNormalization: + def _get_score( + self, + seq2seq_model: MagicMock, + seq2seq_tokenizer: MagicMock, + token_ids: list[int], + logit_value: float, + *, + normalize: bool, + ) -> float: + """Compute the callable score for a given token_ids target.""" + vocab_size = 100 + + seq2seq_tokenizer.encode_queue = [token_ids] + + def make_logits(**kwargs) -> SimpleNamespace: + dec_len = kwargs["decoder_input_ids"].shape[1] + logits = torch.zeros((1, dec_len, vocab_size)) + for tid in token_ids: + logits[:, :, tid] = logit_value + return SimpleNamespace(logits=logits) + + seq2seq_model.side_effect = make_logits + + callable_obj = Seq2SeqCallable( + model=seq2seq_model, + tokenizer=seq2seq_tokenizer, + device="cpu", + target_label="label", + normalize=normalize, + ) + + return float(callable_obj.predict(["text"])[0]) + + def test_single_token_normalize_equals_raw( + self, + seq2seq_model: MagicMock, + seq2seq_tokenizer: MagicMock, + ) -> None: + """For a single-token target, normalize=True and normalize=False + must produce the same score (dividing by 1 is a no-op).""" + raw = self._get_score(seq2seq_model, seq2seq_tokenizer, [5], 2.0, normalize=False) + + seq2seq_model.reset_mock() + norm = self._get_score(seq2seq_model, seq2seq_tokenizer, [5], 2.0, normalize=True) + + assert abs(raw - norm) < 1e-6 + + def test_multi_token_norm_equals_raw_divided_by_n( + self, + seq2seq_model: MagicMock, + seq2seq_tokenizer: MagicMock, + ) -> None: + """For a multi-token target, norm_score must equal raw_score / n_tokens.""" + token_ids = [5, 6, 7] + + raw = self._get_score(seq2seq_model, seq2seq_tokenizer, token_ids, 2.0, normalize=False) + + seq2seq_model.reset_mock() + norm = self._get_score(seq2seq_model, seq2seq_tokenizer, token_ids, 2.0, normalize=True) + + expected = raw / len(token_ids) + assert abs(norm - expected) < 1e-5 + + def test_multi_token_normalized_greater_than_raw( + self, + seq2seq_model: MagicMock, + seq2seq_tokenizer: MagicMock, + ) -> None: + """For a multi-token target, the normalised score must be greater than + the raw score, because dividing a negative number by n > 1 makes it + less negative.""" + token_ids = [5, 6] + + raw = self._get_score(seq2seq_model, seq2seq_tokenizer, token_ids, 2.0, normalize=False) + + seq2seq_model.reset_mock() + norm = self._get_score(seq2seq_model, seq2seq_tokenizer, token_ids, 2.0, normalize=True) + + assert norm > raw + + +# ============================================================================ +# TEST 5 — Prompt template +# ============================================================================ +# _build_prompt inserts the input text into the template string via .format(). +# Different templates change the text sent to the encoder, which must be +# reflected in what the tokenizer receives. +# ============================================================================ + + +class TestPromptTemplate: + def test_default_template_is_noop( + self, + seq2seq_model: MagicMock, + seq2seq_tokenizer: MagicMock, + ) -> None: + """With the default template '{text}', _build_prompt must return + the original text unchanged.""" + seq2seq_tokenizer.encode_queue = [[1]] + + callable_obj = Seq2SeqCallable( + model=seq2seq_model, + tokenizer=seq2seq_tokenizer, + device="cpu", + target_label="positive", + prompt_template="{text}", + ) + + assert callable_obj._build_prompt("hello world") == "hello world" + + @pytest.mark.parametrize( + "template,text,expected", + [ + ( + "sst2 sentence: {text}", + "great film", + "sst2 sentence: great film", + ), + ( + "Sentiment of '{text}':", + "great film", + "Sentiment of 'great film':", + ), + ( + "Q: {text}\nA:", + "great film", + "Q: great film\nA:", + ), + ], + ) + def test_build_prompt_formats_correctly( + self, + seq2seq_model: MagicMock, + seq2seq_tokenizer: MagicMock, + template: str, + text: str, + expected: str, + ) -> None: + """_build_prompt must return the correctly formatted string for + various template styles.""" + seq2seq_tokenizer.encode_queue = [[1]] + + callable_obj = Seq2SeqCallable( + model=seq2seq_model, + tokenizer=seq2seq_tokenizer, + device="cpu", + target_label="positive", + prompt_template=template, + ) + + assert callable_obj._build_prompt(text) == expected + + def test_prompt_is_passed_to_tokenizer( + self, + seq2seq_model: MagicMock, + seq2seq_tokenizer: MagicMock, + ) -> None: + """predict() must pass the rendered prompt — not the raw input text — + to the tokenizer.""" + seq2seq_tokenizer.encode_queue = [[1]] + + vocab_size = 100 + seq2seq_model.return_value = SimpleNamespace(logits=torch.zeros((1, 1, vocab_size))) + + template = "sst2 sentence: {text}" + input_text = "great film" + expected_prompt = "sst2 sentence: great film" + + callable_obj = Seq2SeqCallable( + model=seq2seq_model, + tokenizer=seq2seq_tokenizer, + device="cpu", + target_label="positive", + prompt_template=template, + ) + + callable_obj.predict([input_text]) + + call_args = seq2seq_tokenizer.call_args + actual_texts = call_args[0][0] + assert actual_texts == [expected_prompt] + + def test_different_templates_reach_encoder( + self, + seq2seq_model: MagicMock, + seq2seq_tokenizer: MagicMock, + ) -> None: + """Different prompt templates must cause the tokenizer to receive + different input texts.""" + vocab_size = 100 + seq2seq_model.return_value = SimpleNamespace(logits=torch.zeros((1, 1, vocab_size))) + + prompts_seen: list[list[str]] = [] + + def capture_tokenizer(texts, **kwargs): + prompts_seen.append(texts) + return { + "input_ids": torch.tensor([[10, 11, 12]]), + "attention_mask": torch.ones((1, 3), dtype=torch.long), + } + + seq2seq_tokenizer.side_effect = capture_tokenizer + + templates = ["{text}", "sst2 sentence: {text}"] + + for template in templates: + seq2seq_tokenizer.encode_queue = [[1]] + callable_obj = Seq2SeqCallable( + model=seq2seq_model, + tokenizer=seq2seq_tokenizer, + device="cpu", + target_label="positive", + prompt_template=template, + ) + callable_obj.predict(["great film"]) + + assert prompts_seen[0] != prompts_seen[1] + + +def test_seq2seq_callable_predicts_from_pre_tokenized_inputs( + seq2seq_model: MagicMock, + seq2seq_tokenizer: MagicMock, +) -> None: + seq2seq_tokenizer.encode_queue = [[5, 6]] + + vocab_size = 100 + seq2seq_model.return_value = SimpleNamespace( + logits=torch.zeros((2, 1, vocab_size)), + ) + + callable_obj = Seq2SeqCallable( + model=seq2seq_model, + tokenizer=seq2seq_tokenizer, + device="cpu", + target_label="positive", + ) + + inputs = [ + { + "input_ids": torch.tensor([[10, 11, 12]]), + "attention_mask": torch.tensor([[1, 1, 1]]), + }, + { + "input_ids": torch.tensor([[20, 21, 22]]), + "attention_mask": torch.tensor([[1, 1, 1]]), + }, + ] + + scores = callable_obj.predict_from_inputs(inputs) + + assert isinstance(scores, np.ndarray) + assert scores.shape == (2,) + assert np.all(np.isfinite(scores)) + + encoder = seq2seq_model.get_encoder.return_value + encoder.assert_called_once() + + encoder_inputs = encoder.call_args.kwargs + assert encoder_inputs["input_ids"].shape == (2, 3) + assert encoder_inputs["attention_mask"].shape == (2, 3) + assert encoder_inputs["return_dict"] is True + + +# ============================================================================ +# TEST 6 — TextImputer end-to-end integration +# ============================================================================ +# When TextImputer is initialised with model_type="seq2seq", the internal +# target_callable must be a Seq2SeqCallable instance, and both +# full_prediction() and value_function() must return finite values. +# ============================================================================ + + +class TestSeq2SeqTextImputer: + def _make_imputer( + self, + seq2seq_model: MagicMock, + seq2seq_tokenizer: MagicMock, + player_level: str = "word", + perturbation_type: str = "neutral", + target_label: str = "positive", + prompt_template: str = "{text}", + *, + normalize: bool = True, + ) -> TextImputer: + """Build a seq2seq TextImputer with minimal configuration.""" + seq2seq_tokenizer.encode_queue = [[1]] + + player_strategy = MagicMock() + player_strategy.n_players = 3 + player_strategy.coalition_to_text.return_value = "perturbed text" + + vocab_size = 100 + seq2seq_model.return_value = SimpleNamespace(logits=torch.zeros((1, 1, vocab_size))) + + return TextImputer( + model=seq2seq_model, + tokenizer=seq2seq_tokenizer, + text="original text", + model_type="seq2seq", + target_label=target_label, + prompt_template=prompt_template, + player_strategy=player_strategy, + perturbation_strategy=NeutralPerturbation(), + normalize_target_logprob=normalize, + device="cpu", + ) + + def test_target_callable_is_seq2seq_callable( + self, + seq2seq_model: MagicMock, + seq2seq_tokenizer: MagicMock, + ) -> None: + """With model_type='seq2seq', target_callable must be a + Seq2SeqCallable instance.""" + imputer = self._make_imputer(seq2seq_model, seq2seq_tokenizer) + assert isinstance(imputer.target_callable, Seq2SeqCallable) + + def test_target_label_forwarded_to_callable( + self, + seq2seq_model: MagicMock, + seq2seq_tokenizer: MagicMock, + ) -> None: + """target_label must be forwarded correctly to the internal callable.""" + imputer = self._make_imputer(seq2seq_model, seq2seq_tokenizer, target_label="negative") + assert imputer.target_callable.target_label == "negative" + + def test_prompt_template_forwarded_to_callable( + self, + seq2seq_model: MagicMock, + seq2seq_tokenizer: MagicMock, + ) -> None: + """prompt_template must be forwarded correctly to the internal callable.""" + template = "sst2 sentence: {text}" + imputer = self._make_imputer(seq2seq_model, seq2seq_tokenizer, prompt_template=template) + assert imputer.target_callable.prompt_template == template + + def test_normalize_forwarded_to_callable( + self, + seq2seq_model: MagicMock, + seq2seq_tokenizer: MagicMock, + ) -> None: + """normalize_target_logprob must be forwarded correctly to the + internal callable.""" + imputer = self._make_imputer(seq2seq_model, seq2seq_tokenizer, normalize=False) + assert imputer.target_callable.normalize is False + + def test_full_prediction_returns_finite_float( + self, + seq2seq_model: MagicMock, + seq2seq_tokenizer: MagicMock, + ) -> None: + """full_prediction must be a finite float.""" + imputer = self._make_imputer(seq2seq_model, seq2seq_tokenizer) + score = imputer.full_prediction + assert isinstance(score, float) + assert np.isfinite(score) + + def test_value_function_returns_correct_shape( + self, + seq2seq_model: MagicMock, + seq2seq_tokenizer: MagicMock, + ) -> None: + """value_function([[coalition]]) must return a numpy array of shape (1,).""" + imputer = self._make_imputer(seq2seq_model, seq2seq_tokenizer) + + imputer.target_callable = MagicMock() + imputer.target_callable.predict.return_value = np.array([-0.8], dtype=np.float32) + + coalition = np.array([[1, 0, 1]]) + scores = imputer.value_function(coalition) + + assert isinstance(scores, np.ndarray) + assert scores.shape == (1,) + assert np.isfinite(scores[0]) + + def test_value_function_returns_finite_scores_for_all_zero_coalition( + self, + seq2seq_model: MagicMock, + seq2seq_tokenizer: MagicMock, + ) -> None: + """An all-zero coalition (all players masked) must still produce a + finite score.""" + imputer = self._make_imputer(seq2seq_model, seq2seq_tokenizer) + + imputer.target_callable = MagicMock() + imputer.target_callable.predict.return_value = np.array([-2.0], dtype=np.float32) + + coalition = np.array([[0, 0, 0]]) + scores = imputer.value_function(coalition) + + assert np.isfinite(scores[0]) + + def test_encoder_reuse_across_target_tokens( + self, + seq2seq_model: MagicMock, + seq2seq_tokenizer: MagicMock, + ) -> None: + """For a batch of texts, the encoder must be called exactly once + regardless of the number of target tokens.""" + target_token_ids = [10, 20, 30] + vocab_size = 100 + + seq2seq_tokenizer.encode_queue = [target_token_ids] + + encoder_mock = MagicMock() + encoder_mock.return_value = SimpleNamespace(last_hidden_state=torch.zeros((1, 3, 16))) + seq2seq_model.get_encoder.return_value = encoder_mock + + seq2seq_model.return_value = SimpleNamespace(logits=torch.zeros((1, 1, vocab_size))) + + callable_obj = Seq2SeqCallable( + model=seq2seq_model, + tokenizer=seq2seq_tokenizer, + device="cpu", + target_label="three token label", + normalize=False, + ) + + callable_obj.predict(["text"]) + + # Encoder called once; model called once per target token + assert encoder_mock.call_count == 1 + assert seq2seq_model.call_count == len(target_token_ids) diff --git a/uv.lock b/uv.lock index 893d1fd69..65127cc84 100644 --- a/uv.lock +++ b/uv.lock @@ -1003,9 +1003,7 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5d/6e/4c37d51a2b7f82d2ff11bb6b5f7d766d9a011726624af255e843727627a3/greenlet-3.5.3-cp312-cp312-macosx_11_0_universal2.whl", hash = "sha256:719757059f5a53fd0dde23f78cffeafcdd97b21c850ddb7ca684a3c1a1f122e2", size = 288685, upload-time = "2026-06-26T18:22:08.977Z" }, { url = "https://files.pythonhosted.org/packages/7a/73/815dd90131c1b71ebdf53dbc7c276cafec2a1173b97559f97aba72724a87/greenlet-3.5.3-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:efa9f765dd09f9d0cdac651ffdf631ee59ec5dc6ee7a73e0c012ba9c52fbdf5b", size = 604761, upload-time = "2026-06-26T19:07:10.114Z" }, { url = "https://files.pythonhosted.org/packages/9f/57/079cfe76bcef36b153b25607ee91c6fcb58f17f8b23c86bbbeabe0c88d72/greenlet-3.5.3-cp312-cp312-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:7faba15ac005376e02a0384504e0243be3370ce010296a44a820feb342b505ab", size = 617044, upload-time = "2026-06-26T19:10:07.25Z" }, - { url = "https://files.pythonhosted.org/packages/fb/fb/d97dc261209c80744b7c8132693a30d70ec6e7315e632cb0a10b3fec94dd/greenlet-3.5.3-cp312-cp312-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:5795cd1101371140551c645f2d408b8d3c01a5a29cf8a9bce6e759c983682d23", size = 622351, upload-time = "2026-06-26T19:24:16.32Z" }, { url = "https://files.pythonhosted.org/packages/37/87/b4d095775a3fb1bcafbb483fc206b27ebb785724c83051447737085dc54e/greenlet-3.5.3-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:87142215824be6ac05e2e8e2786eec307ccbc27c36723c3881959df654af6861", size = 614244, upload-time = "2026-06-26T18:32:17.594Z" }, - { url = "https://files.pythonhosted.org/packages/8e/ac/e5fee13cbbd0e8de312d9a146584b8a51891c68847330ef9dc8b5109d23f/greenlet-3.5.3-cp312-cp312-manylinux_2_39_riscv64.whl", hash = "sha256:af4923b3096e26a36d7e9cf24ab88083a20f97d191e3b97f253731ce9b41b28c", size = 425395, upload-time = "2026-06-26T19:25:37.144Z" }, { url = "https://files.pythonhosted.org/packages/8a/70/7559b609683650fa2b95b8ab84b4ab0b26556a635d19675e12aa832d826d/greenlet-3.5.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:215275b1b49320987352e6c1b054acca0064f965a2c66992bed9a6f7d913f149", size = 1574210, upload-time = "2026-06-26T19:09:03.077Z" }, { url = "https://files.pythonhosted.org/packages/ae/73/be55392074c60fc37655ca40fa6022457bfbf6718e9e342a7b0b41f96dd2/greenlet-3.5.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:6b1b0eed82364b0e32c4ea0f221452d33e6bb17ae094d9f72aed9851812747ea", size = 1638627, upload-time = "2026-06-26T18:31:44.748Z" }, { url = "https://files.pythonhosted.org/packages/14/40/c57489acf8e37d74e2913d4eff63aa0dba17acccc4bdeef874dde2dbbec9/greenlet-3.5.3-cp312-cp312-win_amd64.whl", hash = "sha256:cde8adafa2365676f74a979744629589999093bc86e2484214f58e61df08902c", size = 239882, upload-time = "2026-06-26T18:23:27.518Z" }, @@ -1013,9 +1011,7 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/9b/ff/a620267401db30a50cc8450ee90730e2d4a85658c055c0e760d4ed47fb13/greenlet-3.5.3-cp313-cp313-macosx_11_0_universal2.whl", hash = "sha256:c8d87c2134d871df96ecdea9cec7cbaab286dadab0f56476e57aaf9e8ac11550", size = 287609, upload-time = "2026-06-26T18:21:14.724Z" }, { url = "https://files.pythonhosted.org/packages/d6/fa/5401ac78021c826a25b6dde0c705e0a8f29b617509f9185a31dac15fbe1b/greenlet-3.5.3-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a2d185dd1621757e70c3861cceffd5317ab4e7ed7eb09c82994828468527ade5", size = 607435, upload-time = "2026-06-26T19:07:11.412Z" }, { url = "https://files.pythonhosted.org/packages/e9/76/1dc144a2e56e65d36405078ed774224375ea520a1870a6e46e08bb4ac7bf/greenlet-3.5.3-cp313-cp313-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:1c514a468149bf8fbbab874188a3535cd8a48a3e353eb53a3d424296f8dbacd3", size = 619787, upload-time = "2026-06-26T19:10:08.396Z" }, - { url = "https://files.pythonhosted.org/packages/57/61/2f5b1adf256d039f5dab8005de8d3d7ad2b0070a3219c0e036b3fbfeb440/greenlet-3.5.3-cp313-cp313-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:9ad04dd75458c6300b047c61b8639092433d205a25a14e310d6582a480efcca1", size = 625580, upload-time = "2026-06-26T19:24:18.344Z" }, { url = "https://files.pythonhosted.org/packages/bf/87/c298cee62df1de4ad7fec32abda73526cff347fd143a6ed4ac369246668a/greenlet-3.5.3-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:915f887cf2682b66419b879423a2e072634aa7b7dce6f3ada4957cfced3f1e9a", size = 616786, upload-time = "2026-06-26T18:32:19.128Z" }, - { url = "https://files.pythonhosted.org/packages/3e/d9/ab7fc9e543e44d6879b0a6ef9a4b2188940fd180cc65d6f646883ddf7201/greenlet-3.5.3-cp313-cp313-manylinux_2_39_riscv64.whl", hash = "sha256:afaabdd554cd7ae9bbb3ca070b0d7fdfd207dbf1d16865f7233837709d354bda", size = 427933, upload-time = "2026-06-26T19:25:38.219Z" }, { url = "https://files.pythonhosted.org/packages/9e/2e/e6f009885ed0705ccf33fe0583c117cfd03cde77e31a596dd5785a30762b/greenlet-3.5.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:766cfd421c13e450feb340cd472a3ed9957d438727b7b4593ad7c76c5d2b0deb", size = 1574316, upload-time = "2026-06-26T19:09:04.273Z" }, { url = "https://files.pythonhosted.org/packages/ef/fe/43fd110b01e40da0adb7c90ac7ea744bef2d43dca00de5095fd2351c2a68/greenlet-3.5.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:2ecda9ec22edf38fa389369eaed8c3d37c05f3c54e69f69438dbb2cc1de1458b", size = 1638614, upload-time = "2026-06-26T18:31:46.297Z" }, { url = "https://files.pythonhosted.org/packages/0f/7c/062447147a61f8b4337b156fe70d32a165fcf2f89d7ca6255e572806705c/greenlet-3.5.3-cp313-cp313-win_amd64.whl", hash = "sha256:c82304750f057167ff60d188df1d0cc1764ce9567eadf03e6a7443bcedd0b30b", size = 239850, upload-time = "2026-06-26T18:21:54.613Z" }, @@ -1023,9 +1019,7 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c3/93/43e116ee114b28737ba7e12952a0d4e2f55944d0f84e42bc91ba7192a3c9/greenlet-3.5.3-cp314-cp314-macosx_11_0_universal2.whl", hash = "sha256:fd2e02fa07485778536a036222d616ab957b1d533f36b3ed98ce725d9c9d3117", size = 288202, upload-time = "2026-06-26T18:23:49.604Z" }, { url = "https://files.pythonhosted.org/packages/82/2f/146d218299046a43d1f029fd544b3d110d0f175a09c715c7e8da4a4a345d/greenlet-3.5.3-cp314-cp314-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:df0a0628d1597eb0897b62f55d1343f772405fd25f3b2a796c76874b0c2e22e8", size = 654096, upload-time = "2026-06-26T19:07:12.71Z" }, { url = "https://files.pythonhosted.org/packages/a0/cc/04738cafb3f45fa991ea44f9de94c47dcec964f5a972300988a6751f49d9/greenlet-3.5.3-cp314-cp314-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:ebd933a6adabc298bab47731a130fe6bfb888bd934eee37810f151159544540d", size = 666304, upload-time = "2026-06-26T19:10:09.503Z" }, - { url = "https://files.pythonhosted.org/packages/86/a9/73fa62893d5b84b4205544e6b673c654cc43aa5b9899bac00f04d64af73d/greenlet-3.5.3-cp314-cp314-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:8d19fe6c39ebff9259f07bcc685d3290f8fa4ea2278e51dd0008e4d6b0f2d814", size = 670657, upload-time = "2026-06-26T19:24:19.967Z" }, { url = "https://files.pythonhosted.org/packages/ce/aa/4e0dad5e605c270c784ab911c43da6adb136ccd4d81180f763ca429a723d/greenlet-3.5.3-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4b9d501b40e80b70e32323c799dd9b420a5577a9601469d362ae1ffb690f3a7c", size = 663635, upload-time = "2026-06-26T18:32:20.802Z" }, - { url = "https://files.pythonhosted.org/packages/29/7e/2ffce64929fb3cab7b65d5a0b20aaf9764e227681d731b041077fc9a525a/greenlet-3.5.3-cp314-cp314-manylinux_2_39_riscv64.whl", hash = "sha256:962c5df2db8cb446da51edf1ca5296c389d93b99c9d8aa2ee4c7d0d8f1218260", size = 473497, upload-time = "2026-06-26T19:25:39.421Z" }, { url = "https://files.pythonhosted.org/packages/d1/50/13efdbea246fe3d3b735e191fec08fb50809f53cd2383ebe123d0809e44b/greenlet-3.5.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:a1fad1d11e7d6aab184107baa8e4ece11ccba3ec9599cd7efa5ff4d70d43256a", size = 1621252, upload-time = "2026-06-26T19:09:05.647Z" }, { url = "https://files.pythonhosted.org/packages/f7/22/c0a336ae4a1410fd5f5121098e5bfbf1865f64c5ef80b4b5412886c4a332/greenlet-3.5.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:fad5aec764399f1b5cc347ad250a59660f20c8f8888ea6bae1f93b769cce1154", size = 1684824, upload-time = "2026-06-26T18:31:47.738Z" }, { url = "https://files.pythonhosted.org/packages/7a/94/91aec0030bea75c4b3244251d0de60a1f3432d1ecb53ab6c437fb5c3ba61/greenlet-3.5.3-cp314-cp314-win_amd64.whl", hash = "sha256:7669aa24cf2a1041d6f7899575b494a3ab4cf68bfcc8609b1dc0be7272db835e", size = 240754, upload-time = "2026-06-26T18:22:15.669Z" }, @@ -1033,18 +1027,14 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/91/95/3e161213d7f1d378d15aa9e792093e9bfe01844680d04b7fd6e0107c9098/greenlet-3.5.3-cp314-cp314t-macosx_11_0_universal2.whl", hash = "sha256:271a8ea7c1024e8a0d7dd2be66dd66dda8a07193f41a17b9e924f7600f5b62be", size = 296389, upload-time = "2026-06-26T18:22:20.657Z" }, { url = "https://files.pythonhosted.org/packages/00/92/715c44721abe2b4d1ae9abde4179411868a5bff312479f54e105d372f131/greenlet-3.5.3-cp314-cp314t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:19131729ae0ddc3c2e1ef85e650169b5e37ee32e400f215f78b94d7b0d567310", size = 653382, upload-time = "2026-06-26T19:07:14.209Z" }, { url = "https://files.pythonhosted.org/packages/a0/83/37a10372a1090a6624cca8e74c12df1a36c2dc36429ed0255b7fb1aeee23/greenlet-3.5.3-cp314-cp314t-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:1540dd8e5fc2a5aec40fbb98ef8e149fa47c89a4b4a1cf2575a14d3d1869d7a8", size = 659401, upload-time = "2026-06-26T19:10:10.876Z" }, - { url = "https://files.pythonhosted.org/packages/cb/73/8faec206b851c22b1733545fda900829a1f3f5b1c78ae7e0fb3dba57d9f4/greenlet-3.5.3-cp314-cp314t-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:b897d97759425953f69a9c0fac67f8fe333ec0ce7377ef186fb2b0c3ad5e354d", size = 659582, upload-time = "2026-06-26T19:24:21.357Z" }, { url = "https://files.pythonhosted.org/packages/db/e2/d1509cad4207da559cc42986ecdd8fc67ad0d1bba2bf03023c467fd5e0f3/greenlet-3.5.3-cp314-cp314t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e81fa194a1d20967877bdf9c7794db2bc99063e5be36aee710c08f04c5bb087f", size = 656969, upload-time = "2026-06-26T18:32:22.272Z" }, - { url = "https://files.pythonhosted.org/packages/b4/55/50c19e49f8045834ada71ef12f8ad048eba8517c6aa41161bed676328fae/greenlet-3.5.3-cp314-cp314t-manylinux_2_39_riscv64.whl", hash = "sha256:3236754d423955ea08e9bb5f6c04a7895f9e22c290b66aa7653fcb922d839eb0", size = 491037, upload-time = "2026-06-26T19:25:40.672Z" }, { url = "https://files.pythonhosted.org/packages/86/7d/eaf70de20aadca3a5884aec58362861c64ce45e7b277f47ed026926a3b89/greenlet-3.5.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:55cf4d777485d43110e47133cbba6d74a8885a87ec1227ef0267f9ee80c5aa21", size = 1617822, upload-time = "2026-06-26T19:09:06.893Z" }, { url = "https://files.pythonhosted.org/packages/8a/f9/414d38fc400ae4350d4185eaad1827676f7cf5287b9136e0ed1cbbe20a7f/greenlet-3.5.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:12a248ba75f6a9a236375f52296c498c89ff1d8badf32deb9eca7abd5853f7da", size = 1677983, upload-time = "2026-06-26T18:31:49.396Z" }, { url = "https://files.pythonhosted.org/packages/e4/15/7edb977e08f9bff702fe42d6c902702786ff6b9694058b4e6a2a6ac90e57/greenlet-3.5.3-cp314-cp314t-win_amd64.whl", hash = "sha256:efc6bd60ea02e085862c74a3ef64b147ffc6f1a5ea7d9f26e7a939943f68c1e3", size = 243626, upload-time = "2026-06-26T18:24:41.485Z" }, { url = "https://files.pythonhosted.org/packages/2c/8a/93928dce91e6b3598b5e779e8d1fd6576a504640c58e78627077f6a7a91a/greenlet-3.5.3-cp315-cp315-macosx_11_0_universal2.whl", hash = "sha256:ea03f2f04367845d6b58eeed276e1e56e51f0b97d8ad5a88a7d20a91dc9056cc", size = 288860, upload-time = "2026-06-26T18:22:48.07Z" }, { url = "https://files.pythonhosted.org/packages/4f/ca/69db42d447a1378043e2c8f19c09cbbd1263371505053c496b49066d3d16/greenlet-3.5.3-cp315-cp315-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:78dbef602fda6d97d957eb7937f70c9ce9e9527330347f8f6b6f9e554a9e7a47", size = 659747, upload-time = "2026-06-26T19:07:15.565Z" }, { url = "https://files.pythonhosted.org/packages/a8/0b/af7ac2ef8dd41e3da1a40dda6305c23b9a03e13ba975ec916357b50f8575/greenlet-3.5.3-cp315-cp315-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:6f73857adb8fee13fa56c172bd11262f888c0c648f9fea113e777bb2c7904a81", size = 670419, upload-time = "2026-06-26T19:10:12.293Z" }, - { url = "https://files.pythonhosted.org/packages/25/aa/952cf28c2ff949a8c971134fb43854dd7eaa737218723aaef758f8c9aead/greenlet-3.5.3-cp315-cp315-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:cefa9cef4b371f9844c6053db71f1138bc6807bab1578b0dae5149c1f1141357", size = 674261, upload-time = "2026-06-26T19:24:22.79Z" }, { url = "https://files.pythonhosted.org/packages/51/1e/1d51640cacbfc455dbe9f9a9f594c49e4e244f63b9971a2f4764e46cc53d/greenlet-3.5.3-cp315-cp315-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:232fec92e823addaf02d9472cf7381e24a1d046a6ced1103c5caa4c21b9dfc1d", size = 668787, upload-time = "2026-06-26T18:32:24.298Z" }, - { url = "https://files.pythonhosted.org/packages/dc/f2/b00d6f5e63e531a93562b2ec1a4c320fbee91f580fc42e6417af69d706e5/greenlet-3.5.3-cp315-cp315-manylinux_2_39_riscv64.whl", hash = "sha256:6219b6d04dbf6ba6084d77dc609e8473060dc55f759cbf626d512122781fa128", size = 480322, upload-time = "2026-06-26T19:25:41.852Z" }, { url = "https://files.pythonhosted.org/packages/21/66/4030d5b0b5894500023f003bb054d9bb354dfbd1e186c3a296759172f5f5/greenlet-3.5.3-cp315-cp315-musllinux_1_2_aarch64.whl", hash = "sha256:2421c3564da9429d5586d46ca31ebb26516b5498a802cf65c041a8e8a8980d34", size = 1626305, upload-time = "2026-06-26T19:09:08.281Z" }, { url = "https://files.pythonhosted.org/packages/0e/50/5221371c7550108dfa3c378debc41d032aa9c78e89abb01d8011cfc93289/greenlet-3.5.3-cp315-cp315-musllinux_1_2_x86_64.whl", hash = "sha256:e0f0d160f0b2e558e6c75f7930967183255dc9735e5f5b8cae58ee09c9576d8b", size = 1688631, upload-time = "2026-06-26T18:31:51.278Z" }, { url = "https://files.pythonhosted.org/packages/68/5d/00d469daae3c65d2bf620b10eee82eb022127d483c6bc8c69fae6f3fbf17/greenlet-3.5.3-cp315-cp315-win_amd64.whl", hash = "sha256:dd99329bbc15ca78dcc583dba05d0b1b0bae01ab6c2174989f5aaee3e41ac930", size = 241027, upload-time = "2026-06-26T18:22:38.203Z" }, @@ -1052,9 +1042,7 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/1c/da/4f4a8450962fad137c1c8981a3f1b8919d06c829993d4d476f9c525d5173/greenlet-3.5.3-cp315-cp315t-macosx_11_0_universal2.whl", hash = "sha256:176bc16a721fa5fc294d70b87b4dfa5fbdd251b3da5d5372735ecef9bd7d6d0c", size = 297221, upload-time = "2026-06-26T18:23:27.176Z" }, { url = "https://files.pythonhosted.org/packages/57/66/b3bfae3e220a9b63ea539a0eea681800c69ab1aada757eae8789f183e7ce/greenlet-3.5.3-cp315-cp315t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:629b614d2b786e89c50440e246f33eea78f58a962d0bdbbcc809e6d13605903f", size = 657221, upload-time = "2026-06-26T19:07:16.973Z" }, { url = "https://files.pythonhosted.org/packages/7b/81/b6d4d73a709684fc77e7fa034d7c2fe82cffa9fc920fadcaa659c2626213/greenlet-3.5.3-cp315-cp315t-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:2b2e857ae16f5f72142edf75f9f176fe7526ba19a2841df1420516f83831c9f2", size = 663226, upload-time = "2026-06-26T19:10:13.723Z" }, - { url = "https://files.pythonhosted.org/packages/e9/39/0e0938a75115b939d42733a2a12e1d349653c9531fe6fe563e8a681f04e6/greenlet-3.5.3-cp315-cp315t-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:16d192579ed281051396dddd7f7754dac6259e6b1fb26378c87b66622f8e3f91", size = 663706, upload-time = "2026-06-26T19:24:24.312Z" }, { url = "https://files.pythonhosted.org/packages/f5/07/e210b02b589f16e74ff48b730690e4a34ffe984219fce4f3c1a0e7ec8545/greenlet-3.5.3-cp315-cp315t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e515757e2e36bcbf1fad09a46e1557e8b1ae1797d4b44d09da7deed88ad28608", size = 660802, upload-time = "2026-06-26T18:32:26.081Z" }, - { url = "https://files.pythonhosted.org/packages/5b/41/35d1c678cdb3c3b9e6bee691728e563cfb294202b23c7a4c3c2ccc343589/greenlet-3.5.3-cp315-cp315t-manylinux_2_39_riscv64.whl", hash = "sha256:4399eb8d041f20b68d943918bc55502a93d6fdc0a37c14da7881c04139acee9d", size = 498803, upload-time = "2026-06-26T19:25:43.063Z" }, { url = "https://files.pythonhosted.org/packages/eb/2e/5303eb3fa06bca089060f479707182a93e360683bc252acf846c3090d34e/greenlet-3.5.3-cp315-cp315t-musllinux_1_2_aarch64.whl", hash = "sha256:b363d46ed1ea431825fdb01471bb024fc08399bad1572a616e853c7684415adb", size = 1622157, upload-time = "2026-06-26T19:09:09.527Z" }, { url = "https://files.pythonhosted.org/packages/54/70/50de47a488f14df260b50ae34fb5d56016e308b098eab02c878b5223c26a/greenlet-3.5.3-cp315-cp315t-musllinux_1_2_x86_64.whl", hash = "sha256:e44da2f5bbdaabaf7d80b73dbb430c7035771e9f244e3c8b769715c9d8fa0a16", size = 1681159, upload-time = "2026-06-26T18:31:52.986Z" }, { url = "https://files.pythonhosted.org/packages/a7/13/1055e1dda7882073eda533e2b96c62e55bbd2db7fda6d5ece992febc7071/greenlet-3.5.3-cp315-cp315t-win_amd64.whl", hash = "sha256:8ff8bed3e3baa20a3ea261ce00526f1898ad4801d4886fd2220580ee0ad8fadf", size = 244007, upload-time = "2026-06-26T18:22:04.353Z" }, @@ -2219,6 +2207,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl", hash = "sha256:d47fbf302e7d9cbbb9e2555a0d267983d2aa476bac30e90dfbe5669bd57f3762", size = 2068504, upload-time = "2025-12-08T17:02:38.159Z" }, ] +[[package]] +name = "nltk" +version = "3.9.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "click" }, + { name = "joblib" }, + { name = "regex" }, + { name = "tqdm" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/74/a1/b3b4adf15585a5bc4c357adde150c01ebeeb642173ded4d871e89468767c/nltk-3.9.4.tar.gz", hash = "sha256:ed03bc098a40481310320808b2db712d95d13ca65b27372f8a403949c8b523d0", size = 2946864, upload-time = "2026-03-24T06:13:40.641Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9d/91/04e965f8e717ba0ab4bdca5c112deeab11c9e750d94c4d4602f050295d39/nltk-3.9.4-py3-none-any.whl", hash = "sha256:f2fa301c3a12718ce4a0e9305c5675299da5ad9e26068218b69d692fda84828f", size = 1552087, upload-time = "2026-03-24T06:13:38.47Z" }, +] + [[package]] name = "nodeenv" version = "1.10.0" @@ -3821,6 +3824,11 @@ sparse = [ { name = "galois" }, { name = "sparse-transform" }, ] +text = [ + { name = "nltk" }, + { name = "torch" }, + { name = "transformers" }, +] [package.dev-dependencies] all-ml = [ @@ -3902,6 +3910,7 @@ requires-dist = [ { name = "linear-operator", marker = "extra == 'shapleig'", specifier = ">=0.6" }, { name = "matplotlib" }, { name = "networkx" }, + { name = "nltk", marker = "extra == 'text'", specifier = ">=3.9.4" }, { name = "numpy" }, { name = "optuna", marker = "extra == 'benchmark'" }, { name = "pandas" }, @@ -3913,11 +3922,13 @@ requires-dist = [ { name = "sparse-transform", marker = "extra == 'sparse'" }, { name = "tabpfn", marker = "extra == 'benchmark'" }, { name = "torch", marker = "extra == 'shapleig'", specifier = ">=2.9.1" }, + { name = "torch", marker = "extra == 'text'" }, { name = "tqdm" }, + { name = "transformers", marker = "extra == 'text'" }, { name = "xgboost", marker = "extra == 'benchmark'" }, { name = "xgboost", marker = "extra == 'proxy'" }, ] -provides-extras = ["sparse", "proxy", "shapleig", "benchmark"] +provides-extras = ["sparse", "proxy", "shapleig", "text", "benchmark"] [package.metadata.requires-dev] all-ml = [