"""Run the SKYCONTROL command model on one transmission.
    pip install torch safetensors
    python example.py
Prompt: <ctx> the aircraft on frequency <said> what the controller said <cmd>   ->   CALLSIGN command ... <eos>
"""
import torch
from safetensors.torch import load_file
from model import Config, GPT
from bpe import BPE

cfg = Config(vocab=1349, d=192, layers=4, heads=6, max_len=512)
model = GPT(cfg); model.load_state_dict(load_file("model.safetensors")); model.eval()
tok = BPE.load("tokenizer.json")

ctx = ("@APP 07\n"
       "QFA76 qantas 76 ARR APP approach 11000ft h077 07\n"
       "UAE865 emirates 865 ARR APP approach 9000ft h122 07\n"
       "UAE880 emirates 880 ARR APP approach 10000ft h347 07")
said = "emirates eight six five, come on down to seven thousand"

@torch.no_grad()
def command(ctx: str, said: str, max_new: int = 48) -> str:
    ids = tok.encode(f"<ctx>{ctx}<said> {said}<cmd>")
    eos, out = tok.special["<eos>"], []
    caches = model.new_cache()
    logits = model(torch.tensor([ids]), caches, start=0)[0, -1]
    pos = len(ids)
    for _ in range(max_new):
        nxt = int(logits.argmax())
        if nxt == eos: break
        out.append(nxt)
        logits = model(torch.tensor([[nxt]]), caches, start=pos)[0, -1]; pos += 1
    return tok.decode(out).strip()

print(said, "->", command(ctx, said))
for s in ["emirates 880 turn left heading 300", "qantas 76 reduce speed 180 knots", "qantas 76 descend and maintain 5000", "emirates 880 turn left heading three zero zero", "how is your day going"]:
    print(s, "->", command(ctx, s))