| """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)) |
|
|
