# GenAI Program g7 - Transformer sequence reversal (PyTorch)
!pip install torch


import torch
import torch.nn as nn
import torch.optim as optim

vocab_size, d_model, max_len, batch_size, epochs = 20, 32, 6, 16, 10

embed = nn.Embedding(vocab_size, d_model)
model = nn.Transformer(d_model=d_model, nhead=2, num_encoder_layers=2,
                       num_decoder_layers=2, batch_first=True)
fc    = nn.Linear(d_model, vocab_size)

params    = list(model.parameters()) + list(embed.parameters()) + list(fc.parameters())
optimizer = optim.Adam(params, lr=0.01)
loss_fn   = nn.CrossEntropyLoss()

def gen():
    x    = torch.randint(1, vocab_size, (batch_size, max_len))
    y    = torch.flip(x, dims=[1])
    y_in = torch.cat([torch.zeros(batch_size, 1, dtype=torch.long), y[:, :-1]], dim=1)
    return x, y_in, y

for e in range(epochs):
    src, tgt_in, tgt_out = gen()
    out  = fc(model(embed(src), embed(tgt_in)))
    loss = loss_fn(out.reshape(-1, vocab_size), tgt_out.reshape(-1))
    optimizer.zero_grad(); loss.backward(); optimizer.step()
    print(f"Epoch {e+1:2d} | Loss: {loss.item():.4f}")

src, tgt_in, tgt_out = gen()
with torch.no_grad():
    pred = fc(model(embed(src), embed(tgt_in))).argmax(-1)

print(f"\nInput  : {src[0].tolist()}")
print(f"Target : {tgt_out[0].tolist()}")
print(f"Pred   : {pred[0].tolist()}")
