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

import numpy as np
import torch
torch.manual_seed(9)
torch.set_num_threads(1)
counts = np.array([[2., 0., 1.], [0., 0., 0.]])
smoothed = (counts+.5)/(counts.sum(1, keepdims=True)+.5*3)
assert np.allclose(smoothed.sum(1), 1)
assert np.allclose(smoothed[1], np.full(3, 1/3))
tokens = torch.tensor([[1, 3, 4, 2, 0], [1, 3, 5, 4, 2]])  # PAD=0 BOS=1 EOS=2
inputs = tokens[:, :-1]
targets = tokens[:, 1:].clone()
targets[targets == 0] = -100
model = torch.nn.Embedding(6, 6)
optimizer = torch.optim.Adam(model.parameters(), lr=.15)
loss_fn = torch.nn.CrossEntropyLoss(ignore_index=-100)
initial = loss_fn(model(inputs).reshape(-1, 6), targets.reshape(-1)).item()
for _ in range(30):
    optimizer.zero_grad()
    loss = loss_fn(model(inputs).reshape(-1, 6), targets.reshape(-1))
    loss.backward()
    optimizer.step()
with torch.no_grad():
    token_losses = torch.nn.functional.cross_entropy(model(inputs).reshape(-1, 6),
        targets.reshape(-1), ignore_index=-100, reduction="none")
    count = (targets != -100).sum()
    nll = token_losses.sum()/count
assert count.item() == 7 and nll.item() < initial
assert np.isfinite(np.exp(nll.item()))
print("valid targets", count.item(), "mean NLL", nll.item(), "PPL", np.exp(nll.item()))