# Source: content/notes/deep-learning/attention-and-transformers.md
# Independent CPU example; use the curriculum environment.
# See /notes/ml/#example-environment or /notes/deep-learning/#example-environment.

import torch
from torch import nn

torch.manual_seed(12)
torch.set_num_threads(1)
key_count, value_count, pairs, width = 8, 8, 4, 32
mask_value = value_count

def make_examples(count, generator, excluded=None):
    seen = set() if excluded is None else set(excluded)
    keys, values, answers = [], [], []
    while len(keys) < count:
        context_keys = torch.randperm(key_count, generator=generator)[:pairs]
        context_values = torch.randint(value_count, (pairs,), generator=generator)
        chosen = int(torch.randint(pairs, (), generator=generator))
        key_row = torch.cat((context_keys, context_keys[chosen:chosen+1]))
        value_row = torch.cat((context_values, torch.tensor([mask_value])))
        signature = tuple(key_row.tolist() + value_row.tolist())
        if signature in seen:
            continue
        seen.add(signature)
        keys.append(key_row)
        values.append(value_row)
        answers.append(int(context_values[chosen]))
    return torch.stack(keys), torch.stack(values), torch.tensor(answers), seen

train_keys, train_values, train_y, train_rows = make_examples(
    2048, torch.Generator().manual_seed(120))
test_keys, test_values, test_y, all_rows = make_examples(
    384, torch.Generator().manual_seed(121), train_rows)
assert len(all_rows) == len(train_rows) + len(test_y)

class TinyCausalRetriever(nn.Module):
    def __init__(self):
        super().__init__()
        self.key = nn.Embedding(key_count, width)
        self.value = nn.Embedding(value_count + 1, width)
        self.layer = nn.TransformerEncoderLayer(
            width, nhead=4, dim_feedforward=64, dropout=0.0,
            batch_first=True, norm_first=True)
        self.norm = nn.LayerNorm(width)
        self.head = nn.Linear(width, value_count)

    def forward(self, keys, values):
        x = self.key(keys) + self.value(values)
        blocked = torch.ones(keys.shape[1], keys.shape[1], dtype=torch.bool).triu(1)
        return self.head(self.norm(self.layer(x, src_mask=blocked)))

model = TinyCausalRetriever()
optimizer = torch.optim.AdamW(model.parameters(), lr=0.005, weight_decay=0.01)
loss_fn = nn.CrossEntropyLoss()
sampler = torch.Generator().manual_seed(122)
initial = loss_fn(model(train_keys[:128], train_values[:128])[:, -1],
                  train_y[:128]).item()
for _ in range(320):
    ids = torch.randint(len(train_y), (64,), generator=sampler)
    optimizer.zero_grad(set_to_none=True)
    loss = loss_fn(model(train_keys[ids], train_values[ids])[:, -1], train_y[ids])
    loss.backward()
    nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    optimizer.step()
model.eval()
with torch.no_grad():
    logits = model(test_keys, test_values)
    test_loss = loss_fn(logits[:, -1], test_y).item()
    accuracy = (logits[:, -1].argmax(-1) == test_y).float().mean().item()
    no_context = model(test_keys[:, -1:], test_values[:, -1:])[:, -1]
    no_context_accuracy = (no_context.argmax(-1) == test_y).float().mean().item()
    changed_future = test_values.clone()
    changed_future[:, 2:pairs] = (changed_future[:, 2:pairs] + 1) % value_count
    torch.testing.assert_close(model(test_keys, changed_future)[:, :2],
                               logits[:, :2], atol=1e-6, rtol=1e-5)
    matched = (test_keys[:, :pairs] == test_keys[:, -1:]).long().argmax(-1)
    changed_context = test_values.clone()
    changed_answers = (test_y + 1) % value_count
    changed_context[torch.arange(len(test_y)), matched] = changed_answers
    revised = model(test_keys, changed_context)[:, -1].argmax(-1)
    intervention_accuracy = (revised == changed_answers).float().mean().item()
assert test_loss < initial * 0.25 and accuracy > 0.90
assert no_context_accuracy < 0.25
assert intervention_accuracy > 0.85
print({"initial_training_loss": initial, "test_loss": test_loss,
       "retrieval_accuracy": accuracy, "no_context_accuracy": no_context_accuracy,
       "changed_value_accuracy": intervention_accuracy})