# Source: content/notes/nlp/machine-translation.md
# Independent CPU example; use the curriculum environment.
# See /notes/ml/#example-environment or /notes/deep-learning/#example-environment.

import os
os.environ["USE_TF"] = "0"  # select the PyTorch backend before importing Transformers
import torch
from transformers import T5Config, T5ForConditionalGeneration
torch.manual_seed(3)
torch.set_num_threads(1)
config = T5Config(vocab_size=16, d_model=16, d_ff=24, num_layers=1,
    num_decoder_layers=1, num_heads=2, d_kv=8, dropout_rate=0.,
    decoder_start_token_id=0, pad_token_id=0, eos_token_id=1)
model = T5ForConditionalGeneration(config)
source = torch.tensor([[3, 4, 1, 0], [5, 6, 7, 1]])
labels = torch.tensor([[8, 9, 1, -100], [10, 11, 12, 1]])
decoder_input = model.prepare_decoder_input_ids_from_labels(labels)
assert decoder_input.tolist() == [[0, 8, 9, 1], [0, 10, 11, 12]]
optimizer = torch.optim.AdamW(model.parameters(), lr=.001)
output = model(input_ids=source, attention_mask=source.ne(0), labels=labels)
assert output.logits.shape == (2, 4, 16) and torch.isfinite(output.loss)
output.loss.backward()
assert model.shared.weight.grad is not None
optimizer.step()
expected = torch.softmax(torch.tensor([0., torch.log(torch.tensor(2.)).item()]), 0)
assert torch.allclose(expected, torch.tensor([1/3, 2/3]))
print("shifted targets, padding mask, seq2seq update and attention calculation passed")