class

TransformerForSeq2SeqLM

extendsPretrainedModel

TransformerForSeq2SeqLM(config: TransformerConfig)

source edit

Encoder-decoder Transformer with a tied seq2seq language-modeling head.

Wraps TransformerModel with a linear LM head projecting decoder hidden states to the target vocabulary. The head weight is tied to tgt_tok_emb.weight when config.tie_word_embeddings is True, matching the Vaswani 2017 §3.4 convention that halves the parameter cost of the output softmax. Standard interface for translation and summarisation — pass input_ids (the source) and decoder_input_ids (the right-shifted target); supply labels to also compute the shift-loss for training. Greedy inference is exposed through generate.

Parameters

configTransformerConfig

Hyperparameters. config.tie_word_embeddings (default True) controls the LM-head weight tying.

Attributes

transformerTransformerModel

Underlying encoder-decoder trunk.

lm_headnn.Linear

Output projection of shape (d_model, effective_decoder_vocab_size) mapping decoder hidden states to target-vocabulary logits.

Notes

Reference: Vaswani, Shazeer, Parmar, Uszkoreit, Jones, Gomez, Kaiser, and Polosukhin, "Attention Is All You Need", NeurIPS, 2017 (arXiv:1706.03762).

Loss when labels is supplied is the per-token cross-entropy over the right-shifted target sequence:

\mathcal{L}_{\mathrm{S2S}} = -\frac{1}{BT} \sum_{b, t} \log p_\theta(y_{b, t} \mid x_b, y_{b, < t}),

with positions labelled -100 excluded.

Examples

>>> import lucid
>>> from lucid.models.text.transformer import TransformerConfig, TransformerForSeq2SeqLM
>>> cfg = TransformerConfig(num_hidden_layers=2, num_decoder_layers=2,
...                         hidden_size=64, num_attention_heads=2,
...                         intermediate_size=128, vocab_size=1000)
>>> model = TransformerForSeq2SeqLM(cfg).eval()
>>> src = lucid.tensor([[1, 234, 567, 2]])
>>> tgt = lucid.tensor([[1, 100, 200]])
>>> out = model(src, decoder_input_ids=tgt)
>>> out.logits.shape   # (B=1, T_tgt=3, vocab=1000)
(1, 3, 1000)

Used by 2

Constructors

dunder

init

→None

__init__(config: TransformerConfig)

source edit

Instance methods

forward

→Seq2SeqLMOutput

forward(input_ids: Tensor, decoder_input_ids: Tensor, attention_mask: Tensor | None = None, decoder_attention_mask: Tensor | None = None, labels: Tensor | None = None)

source edit

Full encode + decode pass.

Parameters

input_ids

(B, S) source token ids.

decoder_input_ids

(B, T) target token ids (right-shifted).

attention_mask

Optional (B, S) HF-style padding mask.

decoder_attention_mask

Optional (B, T) HF-style padding mask.

Returns

Seq2SeqLMOutput

class:Seq2SeqLMOutput carrying decoder hidden states as

generate

→Tensor

generate(input_ids: Tensor, max_length: int = 32, bos_token_id: int | None = None, eos_token_id: int | None = None, pad_token_id: int | None = None, attention_mask: Tensor | None = None, use_cache: bool = True, do_sample: bool = False, temperature: float = 1.0, top_k: int | None = None, top_p: float | None = None, repetition_penalty: float = 1.0)

source edit

Seq2seq decoding — greedy (default) or stochastic sampling.

Encodes input_ids once, then unrolls the decoder one token at a time (BOS-seeded) until max_length (or every batch row hits eos_token_id). Token selection routes through the shared sampling primitives, so do_sample / temperature / top_k / top_p / repetition_penalty behave exactly as for decoder-only models.

Parameters

input_idsTensor

(B, S) source token ids.

max_lengthint= 32

Cap on generated target length (incl. BOS).

bos_token_idint | None= None

Start-of-sequence id; defaults to config.bos_token_id then 0.

eos_token_idint | None= None

End-of-sequence id; defaults to config.eos_token_id.

pad_token_idint | None= None

Right-padding for finished rows; defaults to config.pad_token_id then 0.

attention_maskTensor | None= None

Optional source padding mask.

use_cachebool= True

Use an EncoderDecoderCache so each step only decodes the new token (default True).

do_samplebool= False

Sample from the (filtered) distribution instead of greedy argmax (default False).

temperaturefloat= 1.0

Logit temperature applied before sampling.

top_kint | None= None

Keep only the top_k highest-probability tokens.

top_pfloat | None= None

Nucleus threshold — keep the smallest set whose cumulative probability reaches top_p.

repetition_penaltyfloat= 1.0

Down-weight logits of already-generated tokens.

Returns

Tensor

(B, T_out) int Tensor of generated target tokens.

>>> import lucid >>> from lucid.models.text.transformer import TransformerConfig, TransformerForSeq2SeqLM >>> cfg = TransformerConfig(num_hidden_layers=2, num_decoder_layers=2, ... hidden_size=64, num_attention_heads=2, ... intermediate_size=128, vocab_size=1000) >>> model = TransformerForSeq2SeqLM(cfg).eval() >>> src = lucid.tensor([[1, 234, 567, 2]]) >>> tgt = lucid.tensor([[1, 100, 200]]) >>> out = model(src, decoder_input_ids=tgt) >>> out.logits.shape # (B=1, T_tgt=3, vocab=1000) (1, 3, 1000)

generate(input_ids: Tensor, max_length: int = 32, bos_token_id: int | None = None, eos_token_id: int | None = None, pad_token_id: int | None = None, attention_mask: Tensor | None = None, use_cache: bool = True, do_sample: bool = False, temperature: float = 1.0, top_k: int | None = None, top_p: float | None = None, repetition_penalty: float = 1.0)