Data Classes

class inseq.data.data_utils.TensorWrapper[source]

Wrapper for tensors and lists of tensors to allow for easy access to their attributes.

__getitem__(subscript) → TensorClass[source]: By default, idiomatic slicing is used for the sequence dimension across batches. For batching use slice_batch instead.

Batching

class inseq.data.batch.BatchEncoding(input_ids: Int64[Tensor, 'batch_size seq_len'], attention_mask: Int64[Tensor, 'batch_size seq_len'], input_tokens: Sequence[Sequence[str]] | None = None, baseline_ids: Int64[Tensor, 'batch_size seq_len'] | None = None)[source]

Output produced by the tokenization process using encode().

input_ids

Batch of token ids with shape [batch_size, longest_seq_length]. Extra tokens for each sentence are padded, and truncation to max_seq_length is performed.

Type:: torch.Tensor

input_tokens

List of lists containing tokens for each sentence in the batch.

Type:: list(list(str))

attention_mask

Batch of attention masks with shape [batch_size, longest_seq_length]. 1 for positions that are valid, 0 for padded positions.

Type:: torch.Tensor

baseline_ids

Batch of reference token ids with shape [batch_size, longest_seq_length]. Used for attribution methods requiring a baseline input (e.g. IG).

Type:: torch.Tensor, optional

class inseq.data.batch.BatchEmbedding(input_embeds: Float[Tensor, 'batch_size seq_len embed_size'] | None = None, baseline_embeds: Float[Tensor, 'batch_size seq_len embed_size'] | None = None)[source]

Embeddings produced by the embedding process using embed().

input_embeds

Batch of token embeddings with shape [batch_size, longest_seq_length, embedding_size] for each sentence in the batch.

Type:: torch.Tensor

baseline_embeds

Batch of reference token embeddings with shape [batch_size, longest_seq_length, embedding_size] for each sentence in the batch.

Type:: torch.Tensor, optional

class inseq.data.batch.Batch(encoding: BatchEncoding, embedding: BatchEmbedding)[source]

Batch of input data for the attribution model.

encoding

Output produced by the tokenization process using encode().

Type:: BatchEncoding

embedding

Embeddings produced by the embedding process using embed().

Type:: BatchEmbedding

All attribute fields are accessible as properties (e.g. batch.input_ids corresponds to: batch.encoding.input_ids)

class inseq.data.batch.EncoderDecoderBatch(sources: Batch, targets: Batch)[source]

Batch of input data for the encoder-decoder attribution model, including information for the source text and the target prefix.

sources

Batch of input data for the source text.

Type:: Batch

targets

Batch of input data for the target prefix.

Type:: Batch

class inseq.data.batch.DecoderOnlyBatch(encoding: BatchEncoding, embedding: BatchEmbedding)[source]: Input batch adapted for decoder-only attribution models, including information for the target prefix.

Attributions

class inseq.data.attribution.FeatureAttributionSequenceOutput(source: list[TokenWithId], target: list[TokenWithId], source_attributions: Float32[Tensor, 'attributed_seq_len generated_seq_len embed_size'] | Float32[Tensor, 'attributed_seq_len generated_seq_len'] | None = None, target_attributions: Float32[Tensor, 'attributed_seq_len generated_seq_len embed_size'] | Float32[Tensor, 'attributed_seq_len generated_seq_len'] | None = None, step_scores: dict[str, Float32[Tensor, 'generated_seq_len']] | None = None, sequence_scores: dict[str, Float32[Tensor, 'attributed_seq_len generated_seq_len']] | None = None, attr_pos_start: int = 0, attr_pos_end: int | None = None, _aggregator: str | list[str] | None = None, _dict_aggregate_fn: dict[str, str] | None = None)[source]

Output produced by a standard attribution method.

source

Tokenized source sequence.

Type:: list of TokenWithId

target

Tokenized target sequence.

Type:: list of TokenWithId

source_attributions

Tensor of shape (source_len, target_len) plus an optional third dimension if the attribution is granular (e.g. gradient attribution) containing the attribution scores produced at each generation step of the target for every source token.

Type:: SequenceAttributionTensor

target_attributions

Tensor of shape (target_len, target_len), plus an optional third dimension if the attribution is granular containing the attribution scores produced at each generation step of the target for every token in the target prefix.

Type:: SequenceAttributionTensor, optional

step_scores

Dictionary of step scores produced alongside attributions (one per generation step).

Type:: dict[str, SingleScorePerStepTensor], optional

sequence_scores

Dictionary of sequence scores produced alongside attributions (n per generation step, as for attributions).

Type:: dict[str, MultipleScoresPerStepTensor], optional

classmethod from_step_attributions(attributions: list[FeatureAttributionStepOutput], tokenized_target_sentences: list[list[TokenWithId]], pad_token: Any | None = None, attr_pos_end: int | None = None) → list[FeatureAttributionSequenceOutput][source]

Converts a list of FeatureAttributionStepOutput objects containing multiple examples outputs per step into a list of FeatureAttributionSequenceOutput with every object containing all step outputs for an individual example.

Raises:: ValueError – If the number of sequences in the attributions is not the same for all input sequences.
Returns:: List of FeatureAttributionSequenceOutput objects.
Return type:: List[FeatureAttributionSequenceOutput]

Visualize the attributions.

Parameters:

min_val (int, optional, defaults to None) – Minimum value in the color range of the visualization. If None, the minimum value of the attributions across all visualized examples is used.
max_val (int, optional, defaults to None) – Maximum value in the color range of the visualization. If None, the maximum value of the attributions across all visualized examples is used.
display (bool, optional, defaults to True) – Whether to display the visualization. Can be set to False if the visualization is produced and stored for later use.
return_html (bool, optional, defaults to False) – Whether to return the HTML code of the visualization.
aggregator (AggregatorPipeline, optional, defaults to None) – Aggregates attributions before visualizing them. If not specified, the default aggregator for the class is used.
do_aggregation (bool, optional, defaults to True) – Whether to aggregate the attributions before visualizing them. Allows to skip aggregation if the attributions are already aggregated.

Returns:

The HTML code of the visualization if return_html is set to True, otherwise None.

Return type:

str

weight_attributions(step_fn_id: str)[source]

Weights attribution scores in place by the value of the selected step function for every generation step.

Parameters:: step_fn_id (str) – The id of the step function to use for weighting the attributions (e.g. probability)

class inseq.data.attribution.FeatureAttributionStepOutput(source_attributions: ~jaxtyping.Float[Tensor, 'batch_size seq_len embed_size'] | ~jaxtyping.Float32[Tensor, 'batch_size attributed_seq_len'] | None = None, step_scores: dict[str, ~jaxtyping.Float32[Tensor, 'batch_size']] | None = None, target_attributions: ~jaxtyping.Float[Tensor, 'batch_size seq_len embed_size'] | ~jaxtyping.Float32[Tensor, 'batch_size attributed_seq_len'] | None = None, sequence_scores: dict[str, ~jaxtyping.Float32[Tensor, 'batch_size attributed_seq_len']] | None = None, source: ~collections.abc.Sequence[~collections.abc.Sequence[~inseq.utils.typing.TokenWithId]] | None = None, prefix: ~collections.abc.Sequence[~collections.abc.Sequence[~inseq.utils.typing.TokenWithId]] | None = None, target: ~collections.abc.Sequence[~collections.abc.Sequence[~inseq.utils.typing.TokenWithId]] | None = None, _sequence_cls: type[FeatureAttributionSequenceOutput] = <class 'inseq.data.attribution.FeatureAttributionSequenceOutput'>)[source]

Output of a single step of feature attribution, plus extra information related to what was attributed.

remap_from_filtered(target_attention_mask: Int64[Tensor, 'batch_size'], batch: DecoderOnlyBatch | EncoderDecoderBatch, is_final_step_method: bool = False) → None[source]: Remaps the attributions to the original shape of the input sequence.

class inseq.data.attribution.FeatureAttributionOutput(sequence_attributions: list[~inseq.data.attribution.FeatureAttributionSequenceOutput], step_attributions: list[~inseq.data.attribution.FeatureAttributionStepOutput] | None = None, info: dict[str, ~typing.Any] = <factory>)[source]

Output produced by the AttributionModel.attribute method.

sequence_attributions

List containing all attributions performed on input sentences (one per input sentence, including source and optionally target-side attribution).

Type:: list of FeatureAttributionSequenceOutput

step_attributions

List containing all step attributions (one per generation step performed on the batch), returned if output_step_attributions=True.

Type:: list of FeatureAttributionStepOutput, optional

info

Dictionary including all available parameters used to perform the attribution.

Type:: dict with str keys and any values

aggregate(aggregator: AggregatorPipeline | type[Aggregator] | None = None, **kwargs) → FeatureAttributionOutput[source]

Aggregate the sequence attributions using one or more aggregators.

Parameters:: aggregator (AggregatorPipeline or Type[Aggregator], optional) – Aggregator or pipeline to use. If not provided, the default aggregator for every sequence attribution is used.
Returns:: Aggregated attribution output
Return type:: FeatureAttributionOutput

get_scores_dicts(aggregator: AggregatorPipeline | type[Aggregator] | None = None, do_aggregation: bool = True, **kwargs) → list[dict[str, dict[str, dict[str, float]]]][source]

Get all computed scores (attributions and step scores) for all sequences as a list of dictionaries.

Returns:: List containing one dictionary per sequence. Every dictionary contains the keys “source_attributions”, “target_attributions” and “step_scores”. For each of these keys, the value is a dictionary with generated tokens as keys, and for values a final dictionary. For “step_scores”, the keys of the final dictionary are the step score ids, and the values are the scores. For “source_attributions” and “target_attributions”, the keys of the final dictionary are respectively source and target tokens, and the values are the attribution scores.
Return type:: list(dict)

This output is intended to be easily converted to a pandas DataFrame. The following example produces a list of DataFrames, one for each sequence, matching the source attributions that would be visualized by out.show().

`python dfs = [pd.DataFrame(x["source_attributions"]) for x in out.get_scores_dicts()] `

static load(path: PathLike, decompress: bool = False) → FeatureAttributionOutput[source]

Load saved attribution output into a new FeatureAttributionOutput object.

Parameters:

path (str) – Path to the JSON file containing the saved attribution output. Note that the file must have been saved with the save() method with use_primitives=False in order to be loaded correctly.
decompress (bool, optional, defaults to False) – If True, the input file is decompressed using gzip.

Returns:

Loaded attribution output

Return type:

FeatureAttributionOutput

save(path: PathLike, overwrite: bool = False, compress: bool = False, ndarray_compact: bool = True, use_primitives: bool = False, split_sequences: bool = False) → None[source]

Save class contents to a JSON file.

Parameters:

path (os.PathLike) – Path to the folder where the attribution output will be stored (e.g. ./out.json).
overwrite (bool, optional, defaults to False) – If True, overwrite the file if it exists, raise error otherwise.
compress (bool, optional, defaults to False) – If True, the output file is compressed using gzip. Especially useful for large sequences and granular attributions with umerged hidden dimensions.
ndarray_compact (bool, optional, defaults to True) – If True, the arrays for scores and attributions are stored in a compact b64 format. Otherwise, they are stored as plain lists of floats.
use_primitives (bool, optional, defaults to False) – If True, the output is stored as a list of dictionaries with primitive types (e.g. int, float, str). Note that an attribution saved with this option cannot be loaded with the load method.
split_sequences (bool, optional, defaults to False) – If True, the output is split into multiple files, one per sequence. The file names are generated by appending the sequence index to the given path (e.g. ./out.json with two sequences -> ./out_0.json, ./out_1.json)

Visualize the sequence attributions.

Parameters:

min_val (int, optional) – Minimum value for color scale.
max_val (int, optional) – Maximum value for color scale.
display (bool, optional) – If True, display the attribution visualization.
return_html (bool, optional) – If True, return the attribution visualization as HTML.
aggregator (AggregatorPipeline or Type[Aggregator], optional) – Aggregator or pipeline to use. If not provided, the default aggregator for every sequence attribution is used.
do_aggregation (bool, optional, defaults to True) – Whether to aggregate the attributions before visualizing them. Allows to skip aggregation if the attributions are already aggregated.

Returns:

Attribution visualization as HTML if return_html=True, None otherwise.

Return type:

str

class inseq.data.attribution.GranularFeatureAttributionSequenceOutput(source: list[TokenWithId], target: list[TokenWithId], source_attributions: Float32[Tensor, 'attributed_seq_len generated_seq_len embed_size'] | Float32[Tensor, 'attributed_seq_len generated_seq_len'] | None = None, target_attributions: Float32[Tensor, 'attributed_seq_len generated_seq_len embed_size'] | Float32[Tensor, 'attributed_seq_len generated_seq_len'] | None = None, step_scores: dict[str, Float32[Tensor, 'generated_seq_len']] | None = None, sequence_scores: dict[str, Float32[Tensor, 'attributed_seq_len generated_seq_len']] | None = None, attr_pos_start: int = 0, attr_pos_end: int | None = None, _aggregator: str | list[str] | None = None, _dict_aggregate_fn: dict[str, str] | None = None)[source]

Raw output of a single sequence of granular feature attribution.

An example of granular feature attribution methods are gradient-based attribution methods such as Integrated Gradients, returning one score per hidden dimension of the model for every generated token.

Adds the convergence delta and default L2 + normalization merging of attributions to the base class.

class inseq.data.attribution.GranularFeatureAttributionStepOutput(source_attributions: ~jaxtyping.Float[Tensor, 'batch_size seq_len embed_size'] | ~jaxtyping.Float32[Tensor, 'batch_size attributed_seq_len'] | None = None, step_scores: dict[str, ~jaxtyping.Float32[Tensor, 'batch_size']] | None = None, target_attributions: ~jaxtyping.Float[Tensor, 'batch_size seq_len embed_size'] | ~jaxtyping.Float32[Tensor, 'batch_size attributed_seq_len'] | None = None, sequence_scores: dict[str, ~jaxtyping.Float32[Tensor, 'batch_size attributed_seq_len']] | None = None, source: ~collections.abc.Sequence[~collections.abc.Sequence[~inseq.utils.typing.TokenWithId]] | None = None, prefix: ~collections.abc.Sequence[~collections.abc.Sequence[~inseq.utils.typing.TokenWithId]] | None = None, target: ~collections.abc.Sequence[~collections.abc.Sequence[~inseq.utils.typing.TokenWithId]] | None = None, _sequence_cls: type[FeatureAttributionSequenceOutput] = <class 'inseq.data.attribution.GranularFeatureAttributionSequenceOutput'>)[source]: Raw output of a single step of gradient feature attribution.