summarization

Sequence summarization implementation.

Functions

evaluate_summarized(input, output)

Evaluate generated summaries.

summarize_tokenized(input, output, ...)

Summarize a given tokenized dataset.

tokenize_summaries_gpt2(input, output)

Tokenize a summarized dataset for GPT2.

Classes

InstructionTraceTransformerEncoderForSequenceSummarizationGPT2(...)

A transformer encoder for multi-modal sequence summarization.

MLPConnector(hidden_dimensions, ...)

A simple multi-layer perceptron language connector.

SummarizationCollator(summary_length)

Collation function for sequence summarization.

TransformerConnector(hidden_dimensions, ...)

A transformer language connector.
undertale.models.summarization.tokenize_summaries_gpt2(input: str, output: str) → str

Tokenize a summarized dataset for GPT2.

Parameters:

  • input – Path to the summarized dataset.

  • output – Path where the tokenized summarized dataset should be written.

Returns:

The path to the tokenized dataset.

class undertale.models.summarization.SummarizationCollator(summary_length: int)

Bases: object

Collation function for sequence summarization.

Parameters:

summary_length – Optional maximum summary token length. Sequences are truncated after padding. Should be set to the language model’s context window minus the number of prefix tokens produced by the connector.

class undertale.models.summarization.MLPConnector(hidden_dimensions: int, connector_dimensions: int, language_dimensions: int, language_tokens: int)

Bases: Module

A simple multi-layer perceptron language connector.

Parameters:

  • hidden_dimensions – The size of the hidden state space.

  • connector_dimensions – Scaling factor; intermediate MLP size is language_dimensions × connector_dimensions.

  • language_dimensions – The size of the language state space.

  • language_tokens – The number of language tokens to produce.

forward(state: Tensor, mask: Tensor | None = None) → Tensor

Project the input tensor into language token space.

Parameters:

  • state – Encoder hidden states of shape (batch, seq_len, hidden_dimensions).

  • mask – Optional attention mask used for pooling.

Returns:

A batched projection into a tensor of shape (batch, language_tokens, language_dimensions).

class undertale.models.summarization.TransformerConnector(hidden_dimensions: int, connector_dimensions: int, language_dimensions: int, language_tokens: int)

Bases: Module

A transformer language connector.

Design inspiration taken from the paper “ClipCap: CLIP Prefix for Image Captioning.”

Uses learnable prefix queries that attend over the full encoder sequence via transformer self-attention to produce richer prefix representations than the MLP connector.

Parameters:

  • hidden_dimensions – The size of the hidden state space.

  • connector_dimensions – Number of attention heads and transformer layers. Must evenly divide language_dimensions.

  • language_dimensions – The size of the language state space.

  • language_tokens – The number of language tokens to produce.

forward(state: Tensor, mask: Tensor | None = None) → Tensor

Project the input tensor into language token space.

Parameters:

  • state – Encoder hidden states of shape (batch, seq_len, hidden_dimensions).

  • mask – Optional attention mask over the encoder sequence.

Returns:

A batched projection into a tensor of shape (batch, language_tokens, language_dimensions).

class undertale.models.summarization.InstructionTraceTransformerEncoderForSequenceSummarizationGPT2(depth: int, hidden_dimensions: int, vocab_size: int, sequence_length: int, heads: int, intermediate_dimensions: int, next_token_id: int, dropout: float, eps: float, lr: float = 0.0001, warmup: float = 0.025, connector_dimensions: int = 8, language_tokens: int = 40)

Bases: LightningModule, Module

A transformer encoder for multi-modal sequence summarization.

Parameters:

  • depth – The number of stacked transformer layers.

  • hidden_dimensions – The size of the hidden state space.

  • vocab_size – The size of the vocabulary.

  • sequence_length – The fixed size of the input vector.

  • heads – The number of attention heads.

  • intermediate_dimensions – The size of the intermediate state space.

  • next_token_id – The ID of the special NEXT token.

  • connector_dimensions – The size of the intermediate state space.

  • language_dimensions – The size of the language state space.

  • language_tokens – The number of language tokens to produce.

  • dropout – Dropout probability.

  • eps – Layer normalization stabalization parameter.

  • lr – Peak learning rate reached after warmup.

  • warmup – Fraction of total steps used for linear warmup.

encode(state: Tensor, mask: Tensor | None = None) → Tensor

Encode and compute language tokens.

Parameters:

  • state – The tokenized input state tensor.

  • mask – Optional attention mask.

Returns:

A tensor in language token space encoding the given input.

forward(state: Tensor, mask: Tensor | None = None, labels: Tensor | None = None) → Tensor | Tuple[Tensor, Tensor]

Encode and summarize the input sequence.

Parameters:

  • state – The tokenized input state tensor.

  • mask – Optional attention mask.

  • labels – Optional summary ground truth tokens.

Returns:

Language model summary logits. If labels are provided a tuple of summary logits and computed loss are returned.

generate(state: Tensor, mask: Tensor | None = None, **kwargs) → Tensor

Generate a summary from a given input.

Parameters:

  • state – The tokenized input state tensor.

  • mask – Optional attention mask.

  • **kwargs – Forwarded to GPT2LMHeadModel.generate (e.g. max_new_tokens, do_sample, temperature).

Returns:

A tensor of generated token IDs.

undertale.models.summarization.summarize_tokenized(input: str, output: str, tokenizer: str, checkpoint: str) → str

Summarize a given tokenized dataset.

Parameters:

  • input – Path to the tokenized dataset.

  • output – Path where the summarized dataset should be written.

  • tokenizer – Path to a trained tokenizer file.

  • checkpoint – Path to a trained model checkpoint.

Returns:

The path to the summarized dataset - adds a generated field containing the generated summary.

undertale.models.summarization.evaluate_summarized(input: str, output: str) → str

Evaluate generated summaries.

Parameters:

  • input – Path to the summarized dataset.

  • output – Path where the evaluated dataset should be written.

Returns:

The path where the evaluation results are written (JSON).