Add ready-to-run ChatML inference CLI

README.md

@@ -64,7 +64,13 @@ This remains a small research model. It is unreliable for multi-step arithmetic,
 
 ```bash
 pip install -r requirements.txt
-python infer.py --checkpoint model.safetensors --prompt "<|im_start|>user\nWhat is Python?<|im_end|>\n<|im_start|>assistant\n"
+python inference.py --prompt "What is Python?"
+```
+
+Omit `--prompt` to start an interactive chat:
+
+```bash
+python inference.py
 ```
 
 ## Distribution

inference.py

@@ -0,0 +1,121 @@
+#!/usr/bin/env python3
+"""Run Aurora Proelia ChatML locally or start an interactive chat."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import torch
+from safetensors.torch import load_file
+from tokenizers import Tokenizer
+
+from aurora.config import load_model_config
+from aurora.model import AuroraForCausalLM
+
+
+DEFAULT_SYSTEM = (
+    "You are Ember Proelia, a proprietary language model created by North ML. "
+    "Answer directly and concisely. Do not claim web access or certainty you do not have."
+)
+
+
+def render_chat(messages: list[dict[str, str]], add_generation_prompt: bool = True) -> str:
+    text = "".join(
+        f"<|im_start|>{item['role']}\n{item['content'].strip()}<|im_end|>\n"
+        for item in messages
+    )
+    if add_generation_prompt:
+        text += "<|im_start|>assistant\n"
+    return text
+
+
+def choose_device(value: str) -> torch.device:
+    if value != "auto":
+        return torch.device(value)
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    if torch.backends.mps.is_available():
+        return torch.device("mps")
+    return torch.device("cpu")
+
+
+def load_model(root: Path, device: torch.device):
+    config = load_model_config(root / "model_ember_proelia_207m_16k.yaml")
+    tokenizer = Tokenizer.from_file(str(root / "tokenizer.json"))
+    dtype = torch.float16 if device.type in {"cuda", "mps"} else torch.float32
+    model = AuroraForCausalLM(config).to(device=device, dtype=dtype).eval()
+    state = load_file(str(root / "model.safetensors"), device=str(device))
+    missing, unexpected = model.load_state_dict(state, strict=False)
+    missing = [name for name in missing if not name.endswith("._extra_state")]
+    if missing or unexpected:
+        raise RuntimeError(f"checkpoint mismatch: missing={missing}, unexpected={unexpected}")
+    return model, tokenizer, config
+
+
+def generate(model, tokenizer: Tokenizer, config, prompt: str, device: torch.device, max_new_tokens: int) -> str:
+    bos = tokenizer.token_to_id("<bos>")
+    eos = tokenizer.token_to_id("<eos>")
+    ids = [bos, *tokenizer.encode(prompt, add_special_tokens=False).ids]
+    generated: list[int] = []
+    with torch.inference_mode():
+        for _ in range(max_new_tokens):
+            inputs = torch.tensor([ids[-int(config.context_length):]], dtype=torch.long, device=device)
+            logits, _ = model(inputs)
+            token = int(torch.argmax(logits[0, -1]).item())
+            if token == eos:
+                break
+            generated.append(token)
+            ids.append(token)
+            text = tokenizer.decode(generated, skip_special_tokens=True)
+            if "<|im_end|>" in text or "<|im_start|>" in text:
+                break
+    text = tokenizer.decode(generated, skip_special_tokens=True).strip()
+    for marker in ("<|im_end|>", "<|im_start|>"):
+        if marker in text:
+            text = text.split(marker, 1)[0].strip()
+    return text
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Aurora Proelia ChatML inference")
+    parser.add_argument("--prompt", help="one prompt; omit for interactive chat")
+    parser.add_argument("--system", default=DEFAULT_SYSTEM)
+    parser.add_argument("--checkpoint-dir", type=Path, default=Path(__file__).resolve().parent)
+    parser.add_argument("--device", default="auto", choices=("auto", "cpu", "cuda", "mps"))
+    parser.add_argument("--max-new-tokens", type=int, default=96)
+    args = parser.parse_args()
+    device = choose_device(args.device)
+    model, tokenizer, config = load_model(args.checkpoint_dir, device)
+    history: list[dict[str, str]] = [{"role": "system", "content": args.system}]
+
+    def answer(user_text: str) -> str:
+        history.append({"role": "user", "content": user_text})
+        prompt = render_chat(history)
+        response = generate(model, tokenizer, config, prompt, device, args.max_new_tokens)
+        history.append({"role": "assistant", "content": response})
+        return response
+
+    if args.prompt:
+        print(answer(args.prompt))
+        return
+    print(f"Aurora Proelia ChatML · device={device}")
+    print("Type /quit to exit, /clear to reset the conversation.")
+    while True:
+        try:
+            user_text = input("You: ").strip()
+        except (EOFError, KeyboardInterrupt):
+            print()
+            break
+        if user_text == "/quit":
+            break
+        if user_text == "/clear":
+            history[:] = [{"role": "system", "content": args.system}]
+            print("Conversation cleared.")
+            continue
+        if user_text:
+            print(f"Aurora: {answer(user_text)}")
+
+
+if __name__ == "__main__":
+    main()