Inseq allows you to define custom target functions for feature attribution. While the default target is the next token's probability, you can implement custom logic (e.g., probability differences for contrastive explanations) by following the standard StepFunction template.
To use a custom function:
- Define a function that accepts
StepFunctionArgs (and any additional keyword arguments you need). - Register the function using
inseq.register_step_function. - Pass the registered
identifier to the attributed_fn parameter in model.attribute(). - Provide any extra arguments required by your function as keyword arguments in
model.attribute().
Note: If your function returns probabilities or values that should be aggregated over contiguous tokens, specify an aggregate_map during registration to tell Inseq how to handle them (e.g., using prod for products).
import inseq
from inseq.attr.step_functions import probability_fn, StepFunctionArgs
# 1. Define the custom function
def example_prob_diff_fn(args: StepFunctionArgs, contrast_ids, contrast_attention_mask):
# ... implementation logic ...
return model_probs - contrast_probs
# 2. Register the function
inseq.register_step_function(
fn=example_prob_diff_fn,
identifier="example_prob_diff",
aggregate_map={"span_aggregate": lambda x: x.prod(dim=1, keepdim=True)},
)
# 3. Use it in attribution
attribution_model = inseq.load_model("Helsinki-NLP/opus-mt-en-it", "saliency")
contrast = attribution_model.encode("Ho salutato la manager", as_targets=True)
out = attribution_model.attribute(
"I said hi to the manager",
"Ho salutato il manager",
attributed_fn="example_prob_diff",
contrast_ids=contrast.input_ids,
contrast_attention_mask=contrast.attention_mask,
attribute_target=True,
step_scores=["example_prob_diff"]
)
