// Native sequence execution: choose an available complete backend explicitly.
// Fresh artifact paths are required on each run.
import aner.token
import aner.gen
import aner.nn
import aner.dev

nn.configure("cpu", "Float32", "preserve")
if gen.device_supported("mps", "Float32") {
    nn.configure("mps", "Float32", "preserve")
}
let codec = token.byte_tokenizer()
let training = token.corpus(codec, List(
    "temperature pressure temperature pressure",
    "pressure temperature pressure temperature"), role: "training")
let model = gen.decoder(codec, context: 32, width: 16,
    heads: 4, layers: 2, seed: 63, dropout: 0.1)
print(model.device())
print(model.dtype())
let estimate = model.memory_estimate(batch: 2, length: 32, training: true)
print(estimate.get("total_bytes"))
print(estimate.get("fits"))
print(model.memory_summary(2, 32, true))
let partial = model.train(training, steps: 3, batch: 2, rate: 0.005, seed: 99,
    clip_norm: 1.0, warmup_steps: 2, weight_decay: 0.01)
partial.save("accelerator-partial.anergen")
let restored = gen.load("accelerator-partial.anergen", device: model.device())
let trained = restored.resume(training, steps: 3)
print(trained.summary())
dev.capture("Native device training", trained)
let prompt = codec.encode("temperature", bos: true, eos: false)
let continuation = trained.generate(prompt, max_new_tokens: 8,
    temperature: 0.8, top_k: 20, top_p: 0.95, seed: 7)
print(codec.decode_bytes(continuation, skip_special: true).hex())