// One explicit differentiable text-classification update on synthetic documents.
import aner.nlp
import aner.token
import aner.nn
import aner.tensor

let documents = nlp.corpus(List("red blue", "red", "blue green"),
    List("document-a", "document-b", "document-c"), role: "training")
let codec = token.byte_tokenizer()
let batch = nlp.batch(documents, codec, max_length: 16,
    truncation: "error", bos: true, eos: true, pad_id: -1)
var table = nn.parameter(nn.xavier(codec.vocabulary(), 4, seed: 17))
var gain = nn.parameter([[1.0, 1.0, 1.0, 1.0]])
var bias = nn.parameter(nn.zeros(1, 4))
var w_in = nn.parameter(nn.xavier(4, 8, seed: 18))
var b_in = nn.parameter(nn.zeros(1, 8))
var w_out = nn.parameter(nn.xavier(8, 4, seed: 19))
var b_out = nn.parameter(nn.zeros(1, 4))
var head = nn.parameter(nn.xavier(4, 2, seed: 20))
var head_bias = nn.parameter(nn.zeros(1, 2))
let optimizer = nn.optimizer("adam", rate: 0.01)
let embedded = nn.embedding(table, batch.token_ids(),
    shape: List(batch.batch_size(), batch.sequence_length()))
let normalized = nn.layer_norm(embedded, gain, bias, epsilon: 0.00001)
// Padding-key masks and pooling masks have different, explicit shapes.
let key_mask = tensor.reshape(batch.mask(),
    List(batch.batch_size(), 1, batch.sequence_length()))
let attended = nn.attention(normalized, normalized, normalized, key_mask,
    causal: false, offset: 0)
let transformed = nn.feed_forward(attended, w_in, b_in, w_out, b_out)
// A custom trainer owns the seed schedule; this is training update zero.
let hidden = nn.add(embedded, nn.dropout(transformed,
    probability: 0.1, seed: 7000, training: true))
let pooled = nn.masked_mean(hidden, batch.mask())
let logits = nn.add_bias(nn.matmul(pooled, head), head_bias)
let loss = nn.cross_entropy_ids(logits, List(0, 0, 1), ignore_index: -1)
nn.backward(loss)
let gradient_norm = nn.clip_grad_norm(List(table, gain, bias, w_in, b_in, w_out, b_out, head, head_bias),
    max_norm: 1.0)
// Each distinct leaf receives exactly one stable optimizer key.
table = nn.update(optimizer, table, "encoder.embedding")
gain = nn.update(optimizer, gain, "encoder.norm.gain")
bias = nn.update(optimizer, bias, "encoder.norm.bias")
w_in = nn.update(optimizer, w_in, "encoder.ffn.in.weight")
b_in = nn.update(optimizer, b_in, "encoder.ffn.in.bias")
w_out = nn.update(optimizer, w_out, "encoder.ffn.out.weight")
b_out = nn.update(optimizer, b_out, "encoder.ffn.out.bias")
head = nn.update(optimizer, head, "classifier.weight")
head_bias = nn.update(optimizer, head_bias, "classifier.bias")
print("Loss before the update")
print(nn.item(loss))
print("Gradient norm before clipping")
print(gradient_norm)
print("Source IDs, retained token lengths and truncation")
print(batch.source_ids())
print(batch.lengths())
print(batch.truncated())
print("Pooled shape")
print(tensor.dim(pooled, 0))
print(tensor.dim(pooled, 1))
print("Exact batch fingerprint")
print(batch.fingerprint())