The MTFraEng pipeline already produces this shifted pair. Same self-supervised setup as a language model — only now the encoder output is concatenated as extra context.
Alternative — scheduled sampling: occasionally feed the prediction back. More realistic at inference but harder to optimize.
Encoder: embed + GRU
Embedding layer for input tokens, then a multilayer GRU. Output: per-step hidden states (top layer) and final state (all layers):
class Seq2SeqEncoder(d2l.Encoder):"""The RNN encoder for sequence-to-sequence learning.""" vocab_size: int embed_size: int num_hiddens: int num_layers: int dropout: float=0def setup(self):self.embedding = nn.Embed(self.vocab_size, self.embed_size)self.rnn = d2l.GRU(self.num_hiddens, self.num_layers, self.dropout)def__call__(self, X, *args, training=False):# X shape: (batch_size, num_steps) embs =self.embedding(d2l.astype(d2l.transpose(X), d2l.int64))# embs shape: (num_steps, batch_size, embed_size) outputs, state =self.rnn(embs, training=training)# outputs shape: (num_steps, batch_size, num_hiddens)# state shape: (num_layers, batch_size, num_hiddens)return outputs, state
Embed the previous target token, concatenate the encoder’s final hidden state at every decoder time step (context broadcast across the sequence), run a GRU, and project to vocab logits:
class Seq2SeqDecoder(d2l.Decoder):"""The RNN decoder for sequence to sequence learning.""" vocab_size: int embed_size: int num_hiddens: int num_layers: int dropout: float=0def setup(self):self.embedding = nn.Embed(self.vocab_size, self.embed_size)self.rnn = d2l.GRU(self.num_hiddens, self.num_layers, self.dropout)self.dense = nn.Dense(self.vocab_size)def init_state(self, enc_all_outputs, *args):return enc_all_outputsdef__call__(self, X, state, training=False):# X shape: (batch_size, num_steps)# embs shape: (num_steps, batch_size, embed_size) embs =self.embedding(d2l.astype(d2l.transpose(X), d2l.int64)) enc_output, hidden_state = state# context shape: (batch_size, num_hiddens) context = enc_output[-1]# Broadcast context to (num_steps, batch_size, num_hiddens) context = jnp.tile(context, (embs.shape[0], 1, 1))# Concat at the feature dimension embs_and_context = d2l.concat((embs, context), -1) outputs, hidden_state =self.rnn(embs_and_context, hidden_state, training=training) outputs = d2l.swapaxes(self.dense(outputs), 0, 1)# outputs shape: (batch_size, num_steps, vocab_size)# hidden_state shape: (num_layers, batch_size, num_hiddens)return outputs, [enc_output, hidden_state]
Decoder shape check
End-to-end forward pass produces (batch, num_steps, vocab) logits and a state of shape (num_layers, batch, num_hiddens):
Run the encoder once, then loop: feed the previous prediction back, take argmax over the vocab. Stop after num_steps (or when <eos> appears — handled by the caller).
Predicting token by token: feed the previous prediction back, stop on <eos>.
@d2l.add_to_class(d2l.EncoderDecoder)def predict_step(self, params, batch, num_steps, save_attention_weights=False): src, tgt, src_valid_len, _ = batch enc_all_outputs, inter_enc_vars =self.encoder.apply( {'params': params['encoder']}, src, src_valid_len, training=False, mutable='intermediates')# Save encoder attention weights if inter_enc_vars containing encoder# attention weights is not empty. (to be covered later) enc_attention_weights = []ifbool(inter_enc_vars) and save_attention_weights:# Encoder Attention Weights saved in the intermediates collection enc_attention_weights = inter_enc_vars['intermediates']['enc_attention_weights'][0] dec_state =self.decoder.init_state(enc_all_outputs, src_valid_len) outputs, attention_weights = [d2l.expand_dims(tgt[:,0], 1), ], []for _ inrange(num_steps): (Y, dec_state), inter_dec_vars =self.decoder.apply( {'params': params['decoder']}, outputs[-1], dec_state, training=False, mutable='intermediates') outputs.append(d2l.argmax(Y, 2))# Save attention weights (to be covered later)if save_attention_weights:# Decoder Attention Weights saved in the intermediates collection dec_attention_weights = inter_dec_vars['intermediates']['dec_attention_weights'][0] attention_weights.append(dec_attention_weights)return d2l.concat(outputs[1:], 1), (attention_weights, enc_attention_weights)
BLEU score
Compare prediction n-grams against reference. Geometric mean of n-gram precisions, with a brevity penalty so the model can’t game it by emitting “the the”.
def bleu(pred_seq, label_seq, k):"""Compute the BLEU.""" pred_tokens, label_tokens = pred_seq.split(' '), label_seq.split(' ') len_pred, len_label =len(pred_tokens), len(label_tokens) score = math.exp(min(0, 1- len_label / len_pred))for n inrange(1, min(k, len_pred) +1): num_matches, label_subs =0, collections.defaultdict(int)for i inrange(len_label - n +1): label_subs[' '.join(label_tokens[i: i + n])] +=1for i inrange(len_pred - n +1):if label_subs[' '.join(pred_tokens[i: i + n])] >0: num_matches +=1 label_subs[' '.join(pred_tokens[i: i + n])] -=1 score *= math.pow(num_matches / (len_pred - n +1), math.pow(0.5, n))return score
Translate four sentences
Run the model on a handful of English sentences and score each. Short, frequent patterns tend to translate cleanly; low BLEU on a sentence usually means one missing or misplaced token is enough to break an n-gram match:
engs = ['i lost .', 'i\'m calm .', 'i\'m home .']fras = ['j\'ai perdu .', 'je suis calme .', 'je suis chez moi .']preds, _ = model.predict_step(trainer.state.params, data.build(engs, fras), data.num_steps)for en, fr, p inzip(engs, fras, preds): translation = []for token in data.tgt_vocab.to_tokens(p):if token =='<eos>':break translation.append(token) print(f'{en} => {translation}, bleu,'f'{bleu(" ".join(translation), fr, k=2):.3f}')
