A native language model research workflow

// Native from-scratch workflow and executable correctness checks.
// Small synthetic strings expose mechanisms; they are not a language-quality benchmark.
// Run from a fresh output directory using a runtime with aner.token and aner.gen.
import aner.token
import aner.gen
import aner.nn
import aner.dev
import aner.math

fn verify(label: String, condition: Bool) -> Unit {
    print("check_" + label)
    if !condition {
        let failure: Result<Bool, String> = Result.err(label)
        print(failure.unwrap())
    }
    print(condition)
}
fn trace(model: LanguageModel, training: TokenCorpus, validation: TokenCorpus) -> Unit {
    print("trace_step")
    print(model.steps())
    print("trace_training_loss")
    print(model.loss(training))
    print("trace_validation_loss")
    print(model.loss(validation))
}

nn.configure("cpu", "Float64", "preserve")

// Partition real research documents by subject, experiment or time before this step.
// Learn BPE only from training documents, retaining exact UTF-8 byte fallback.
let fitting_documents = List(
    "temperature α pressure β temperature α pressure β",
    "pressure β temperature α pressure β temperature α")
let fitting = token.corpus(token.byte_tokenizer(), fitting_documents, role: "training")
let bpe = token.bpe(fitting, vocabulary: 280, min_frequency: 2)
let symbols = "new scientific symbols α β 🔬"
verify("bpe_roundtrip", bpe.decode(bpe.encode(symbols, bos: true, eos: true), skip_special: true) == symbols)
bpe.save("research-bpe.anertoken")
print("bpe_vocabulary")
print(bpe.vocabulary())

// Fixed byte IDs make every point on the learning trace directly comparable.
let codec = token.byte_tokenizer()
let training = token.corpus(codec, List(
    "abababababababab", "babababababababa"), role: "training")
let validation = token.corpus(codec, List(
    "abababababababa", "bababababababab"), role: "validation")
let testing = token.corpus(codec, List(
    "ababababab", "babababab"), role: "test")
let initial = gen.decoder(codec, context: 24, width: 8,
    heads: 2, layers: 1, seed: 63, dropout: 0.0)
let initial_validation = initial.loss(validation)
print("uniform_loss")
print(math.log(259.0))
trace(initial, training, validation)
let step30 = initial.train(training, steps: 30, batch: 2, rate: 0.01, seed: 99,
    clip_norm: 1.0, warmup_steps: 0, weight_decay: 0.0)
trace(step30, training, validation)
let step60 = step30.resume(training, steps: 30)
trace(step60, training, validation)
step60.save("research-step60.anergen")
let loaded = gen.load("research-step60.anergen")
verify("checkpoint_roundtrip", loaded.fingerprint() == step60.fingerprint())
let step90 = loaded.resume(training, steps: 30)
trace(step90, training, validation)
let model = step90.resume(training, steps: 30)
trace(model, training, validation)
let direct = initial.train(training, steps: 120, batch: 2, rate: 0.01, seed: 99,
    clip_norm: 1.0, warmup_steps: 0, weight_decay: 0.0)
verify("exact_split_resume", model.fingerprint() == direct.fingerprint())
verify("training_improves", model.loss(training) < initial.loss(training))
verify("validation_improves", model.loss(validation) < initial_validation && model.loss(validation) < math.log(259.0))
print("test_loss")
print(model.loss(testing))

// Future-token changes cannot change earlier logits, across the full vocabulary.
let left = model.logits(List(256, 97, 98, 97), batch: 1, length: 4, training: false)
let right = model.logits(List(256, 97, 122, 122), batch: 1, length: 4, training: false)
var prefix_equal = true
var row = 0
while row < 2 {
    var column = 0
    while column < codec.vocabulary() {
        if nn.value(left, row, column) != nn.value(right, row, column) {
            prefix_equal = false
        }
        column = column + 1
    }
    row = row + 1
}
verify("causal_prefix", prefix_equal)
let prompt = codec.encode("ab", bos: true, eos: false)
let generated = model.generate(prompt, max_new_tokens: 8,
    temperature: 0.8, top_k: 20, top_p: 0.95, seed: 7)
let repeated = direct.generate(prompt, max_new_tokens: 8,
    temperature: 0.8, top_k: 20, top_p: 0.95, seed: 7)
var generation_equal = generated.len() == repeated.len()
var token_index = 0
while token_index < generated.len() && generation_equal {
    generation_equal = cast<Int64>(generated.get(token_index)) == cast<Int64>(repeated.get(token_index))
    token_index = token_index + 1
}
verify("seeded_generation", generation_equal)
print("continuation_hex")
print(codec.decode_bytes(generated, skip_special: true).hex())
model.save("research-trained.anergen")
codec.save("research-byte.anertoken")
dev.capture("From-scratch native decoder", model)

// A changed training corpus starts an explicit fine-tuning run.
let adaptation = token.corpus(codec, List(
    "abababab science", "babababa research"), role: "training")
let full = model.fine_tune(adaptation, steps: 4, batch: 2, rate: 0.001, seed: 101)
verify("full_fine_tune_steps", full.steps() == 4)
let selected = model.freeze(List("embedding.token", "embedding.position"))
verify("freezing_reduces_selection", selected.trainable_parameters().len() < model.trainable_parameters().len())
let adapted = model.lora(List("block.0.qkv", "block.0.attention.output"),
    rank: 2, alpha: 4.0, seed: 23)
let lora = adapted.fine_tune(adaptation, steps: 4, batch: 2, rate: 0.001, seed: 102)
verify("lora_training_steps", lora.steps() == 4)
lora.save("research-lora.anergen")
print("workflow_checks_passed")