Add ChatML model card and local runtime

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+Aurora-3.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,74 @@
+---
+language:
+- en
+library_name: aurora
+pipeline_tag: text-generation
+tags:
+- aurora-proelia
+- north-ml
+- chatml
+- 207m
+license: other
+widget:
+- text: Who are you?
+- text: What is Python?
+- text: Explain photosynthesis in one sentence.
+---
+
+![Project Banner](Aurora-3.png)
+
+# Aurora Proelia ChatML
+
+Aurora Proelia ChatML is a 207M-parameter experimental variant of [Aurora Proelia](https://huggingface.co/North-ML1/Aurora-Proelia). It was SFT-trained on a conventional role-based ChatML surface so applications can send system, user, and assistant turns in a familiar format.
+
+This is a separate candidate. The original `Aurora-Proelia` repository remains the native `Question:` / `Answer:` release.
+
+## ChatML format
+
+Use this format for inference:
+
+```text
+<|im_start|>system
+You are Ember Proelia. Answer directly and concisely.<|im_end|>
+<|im_start|>user
+What is Python?<|im_end|>
+<|im_start|>assistant
+```
+
+The model is a custom Aurora checkpoint, not a Transformers-compatible architecture. Use the included native Aurora runtime, YAML config, and tokenizer.
+
+## What changed
+
+The checkpoint started from the released Aurora candidate and received 2,048 effective ChatML SFT updates over the existing answer-masked ChatML corpus. The pass was intended to teach the input/output surface, not to create a new general-knowledge model.
+
+## Evaluation
+
+On a matched public benchmark mini-slice, the ChatML candidate changed as follows:
+
+| Benchmark | Released Aurora | ChatML candidate |
+|---|---:|---:|
+| MMLU · 57 questions | 14/57 · 24.6% | **16/57 · 28.1%** |
+| ARC-Challenge · 50 questions | 13/50 · 26.0% | **15/50 · 30.0%** |
+| HellaSwag · 50 questions | 19/50 · 38.0% | 19/50 · 38.0% |
+| GSM8K · 50 questions | 1/50 · 2.0% | 0/50 · 0.0% |
+
+The exact runs are in [`benchmarks.json`](./benchmarks.json), [`regression_comparison.json`](./regression_comparison.json), and [`chatml_smoke.json`](./chatml_smoke.json). These are transparent slices of public Hugging Face datasets, not official leaderboard evaluations.
+
+The practical result is clearer than the small score changes: the ChatML candidate answers ordinary identity and Python prompts through the role-based format, while the released checkpoint often echoes the ChatML prompt. Arithmetic and uncertainty handling remain weak.
+
+## Limitations
+
+This remains a small research model. It is unreliable for multi-step arithmetic, deep reasoning, current facts, specialized questions without context, and complex instruction following. Verify important answers and provide retrieval context when freshness or factual accuracy matters.
+
+## Local inference
+
+```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"
+```
+
+## Distribution
+
+This is a public North ML research release. No open-source license is granted; licensing is reserved by the repository owner.
+
+`text-generation` · `aurora-proelia` · `chatml` · `north-ml` · `207m`

aurora/__init__.py

@@ -0,0 +1,4 @@
+from .config import AuroraConfig, load_model_config
+from .model import AuroraForCausalLM
+
+__all__ = ["AuroraConfig", "AuroraForCausalLM", "load_model_config"]

aurora/config.py

@@ -0,0 +1,63 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, fields
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+
+@dataclass
+class AuroraConfig:
+    model_name: str = "Ember Proelia"
+    vocab_size: int = 16000
+    hidden_size: int = 896
+    num_layers: int = 23
+    num_attention_heads: int = 14
+    num_key_value_heads: int = 2
+    intermediate_size: int = 2432
+    context_length: int = 2048
+    rope_theta: float = 500000.0
+    rms_norm_eps: float = 1.0e-5
+    qk_norm: bool = True
+    tie_word_embeddings: bool = True
+    attention_bias: bool = False
+    mlp_bias: bool = False
+    dropout: float = 0.0
+    num_experts: int = 1
+    router_aux_loss_coef: float = 0.0
+    router_z_loss_coef: float = 0.0
+    router_noise_scale: float = 0.0
+    moe_capacity_factor: float = 0.0
+    router_use_gate_weight: bool = False
+
+    @property
+    def head_dim(self) -> int:
+        return self.hidden_size // self.num_attention_heads
+
+    def validate(self) -> None:
+        if self.vocab_size <= 0 or self.hidden_size <= 0 or self.num_layers <= 0:
+            raise ValueError("vocab_size, hidden_size, and num_layers must be positive")
+        if self.hidden_size % self.num_attention_heads != 0:
+            raise ValueError("hidden_size must divide evenly by num_attention_heads")
+        if self.num_attention_heads % self.num_key_value_heads != 0:
+            raise ValueError("num_attention_heads must divide evenly by num_key_value_heads")
+        if self.head_dim % 2 != 0:
+            raise ValueError("head_dim must be even for RoPE")
+        if self.context_length <= 0:
+            raise ValueError("context_length must be positive")
+        if self.num_experts <= 0:
+            raise ValueError("num_experts must be positive")
+
+
+def load_model_config(path: str | Path) -> AuroraConfig:
+    path = Path(path)
+    with path.open("r", encoding="utf-8") as handle:
+        raw: dict[str, Any] = yaml.safe_load(handle) or {}
+    allowed = {field.name for field in fields(AuroraConfig)}
+    unknown = sorted(set(raw) - allowed)
+    if unknown:
+        raise ValueError(f"Unknown model config keys: {unknown}")
+    cfg = AuroraConfig(**raw)
+    cfg.validate()
+    return cfg

aurora/model.py

@@ -0,0 +1,376 @@
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from aurora.config import AuroraConfig
+
+try:
+    from cut_cross_entropy import linear_cross_entropy
+except ImportError:
+    linear_cross_entropy = None
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        scale = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
+        return self.weight * x * scale
+
+
+def precompute_rope_frequencies(
+    seq_len: int, head_dim: int, theta: float, device: torch.device, dtype: torch.dtype
+) -> tuple[torch.Tensor, torch.Tensor]:
+    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
+    positions = torch.arange(seq_len, device=device).float()
+    freqs = torch.outer(positions, inv_freq)
+    return freqs.cos().to(dtype=dtype), freqs.sin().to(dtype=dtype)
+
+
+def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    cos = cos[None, :, None, :]
+    sin = sin[None, :, None, :]
+    x_even = x[..., 0::2]
+    x_odd = x[..., 1::2]
+    out = torch.empty_like(x)
+    out[..., 0::2] = x_even * cos - x_odd * sin
+    out[..., 1::2] = x_even * sin + x_odd * cos
+    return out
+
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, cfg: AuroraConfig) -> None:
+        super().__init__()
+        self.cfg = cfg
+        self.num_heads = cfg.num_attention_heads
+        self.num_kv_heads = cfg.num_key_value_heads
+        self.head_dim = cfg.head_dim
+        self.kv_repeat = self.num_heads // self.num_kv_heads
+
+        self.q_proj = nn.Linear(cfg.hidden_size, cfg.num_attention_heads * self.head_dim, bias=cfg.attention_bias)
+        self.k_proj = nn.Linear(cfg.hidden_size, cfg.num_key_value_heads * self.head_dim, bias=cfg.attention_bias)
+        self.v_proj = nn.Linear(cfg.hidden_size, cfg.num_key_value_heads * self.head_dim, bias=cfg.attention_bias)
+        self.o_proj = nn.Linear(cfg.hidden_size, cfg.hidden_size, bias=cfg.attention_bias)
+        self.q_norm = RMSNorm(self.head_dim, cfg.rms_norm_eps) if cfg.qk_norm else nn.Identity()
+        self.k_norm = RMSNorm(self.head_dim, cfg.rms_norm_eps) if cfg.qk_norm else nn.Identity()
+        self.dropout_p = cfg.dropout
+
+    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+        batch, seq_len, _ = x.shape
+        q = self.q_proj(x).view(batch, seq_len, self.num_heads, self.head_dim)
+        k = self.k_proj(x).view(batch, seq_len, self.num_kv_heads, self.head_dim)
+        v = self.v_proj(x).view(batch, seq_len, self.num_kv_heads, self.head_dim)
+
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+        q = apply_rope(q, cos, sin).transpose(1, 2)
+        k = apply_rope(k, cos, sin).transpose(1, 2)
+        v = v.transpose(1, 2)
+
+        # Explicit K/V expansion is mathematically equivalent to GQA and works
+        # across CUDA, Apple MPS, and CPU PyTorch backends.
+        if self.kv_repeat > 1:
+            k = k.repeat_interleave(self.kv_repeat, dim=1)
+            v = v.repeat_interleave(self.kv_repeat, dim=1)
+        y = F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            attn_mask=None,
+            dropout_p=self.dropout_p if self.training else 0.0,
+            is_causal=True,
+        )
+        y = y.transpose(1, 2).contiguous().view(batch, seq_len, self.cfg.hidden_size)
+        return self.o_proj(y)
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, cfg: AuroraConfig, intermediate_size: int | None = None) -> None:
+        super().__init__()
+        intermediate_size = intermediate_size or cfg.intermediate_size
+        self.gate_proj = nn.Linear(cfg.hidden_size, intermediate_size, bias=cfg.mlp_bias)
+        self.up_proj = nn.Linear(cfg.hidden_size, intermediate_size, bias=cfg.mlp_bias)
+        self.down_proj = nn.Linear(intermediate_size, cfg.hidden_size, bias=cfg.mlp_bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class Top1MoE(nn.Module):
+    """Top-1 routed SwiGLU experts with Switch-style router regularization."""
+
+    def __init__(self, cfg: AuroraConfig) -> None:
+        super().__init__()
+        self.num_experts = cfg.num_experts
+        self.router_aux_loss_coef = cfg.router_aux_loss_coef
+        self.router_z_loss_coef = cfg.router_z_loss_coef
+        self.router_noise_scale = cfg.router_noise_scale
+        self.capacity_factor = cfg.moe_capacity_factor
+        self.use_gate_weight = cfg.router_use_gate_weight
+        # Keep routing in BF16/FP32 rather than quantizing its logits to FP8.
+        self.router = nn.Linear(cfg.hidden_size, cfg.num_experts, bias=False)
+        self.experts = nn.ModuleList([SwiGLU(cfg) for _ in range(cfg.num_experts)])
+        # Detached summaries from the latest batch, for collapse detection in
+        # the trainer.  They are intentionally not persistent model state.
+        self.last_expert_fraction: torch.Tensor | None = None
+        self.last_preferred_expert_fraction: torch.Tensor | None = None
+        self.last_forced_fraction: torch.Tensor | None = None
+        self.last_selected_gate_probability: torch.Tensor | None = None
+
+    def _capacity_constrained_route(
+        self, scores: torch.Tensor, preferred_index: torch.Tensor
+    ) -> torch.Tensor:
+        """Assign exactly one expert/token while bounding every expert load.
+
+        Experts keep their highest-scoring first-choice tokens.  Overflow is
+        deterministically retried against each token's next preference.  With
+        five experts this small eager-only matching pass is far cheaper than
+        an expert MLP and prevents a collapsed router from starving experts.
+        """
+        token_count = scores.size(0)
+        capacity = max(
+            math.ceil(token_count / self.num_experts),
+            math.ceil(token_count * self.capacity_factor / self.num_experts),
+        )
+        rankings = torch.argsort(scores, dim=-1, descending=True)
+        assigned = torch.full_like(preferred_index, -1)
+        remaining = [capacity for _ in range(self.num_experts)]
+
+        for rank in range(self.num_experts):
+            for expert_index in range(self.num_experts):
+                slots = remaining[expert_index]
+                if slots <= 0:
+                    continue
+                candidates = torch.nonzero(
+                    (assigned < 0) & (rankings[:, rank] == expert_index), as_tuple=False
+                ).flatten()
+                candidate_count = candidates.numel()
+                if candidate_count == 0:
+                    continue
+                if candidate_count > slots:
+                    candidate_scores = scores.index_select(0, candidates)[:, expert_index]
+                    best_positions = torch.topk(candidate_scores, k=slots, sorted=False).indices
+                    candidates = candidates.index_select(0, best_positions)
+                    candidate_count = slots
+                assigned.index_fill_(0, candidates, expert_index)
+                remaining[expert_index] -= candidate_count
+
+        # The combined capacity is at least the token count and every token
+        # ranks every expert, so this is a logic invariant rather than an
+        # expected fallback.
+        if bool(torch.any(assigned < 0)):
+            raise RuntimeError("capacity-constrained MoE routing left tokens unassigned")
+        return assigned
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        original_shape = x.shape
+        flat_x = x.reshape(-1, original_shape[-1])
+        router_logits = self.router(flat_x).float()
+        router_probs = torch.softmax(router_logits, dim=-1)
+        # Routing indices are discrete.  Keep this bookkeeping out of the
+        # autograd graph; language gradients still reach the selected gate
+        # probability when gate weighting is enabled, while router losses use
+        # the clean differentiable probabilities below.
+        with torch.no_grad():
+            routing_logits = router_logits
+            if self.training and self.router_noise_scale > 0:
+                # Noisy top-1 routing keeps early training exploratory.  Only
+                # the discrete expert choice is noisy; probability weights and
+                # regularization remain based on clean BF16/FP32 logits.
+                gumbel_noise = -torch.empty_like(router_logits).exponential_().log()
+                routing_logits = router_logits + self.router_noise_scale * gumbel_noise
+            preferred_index = torch.argmax(routing_logits, dim=-1)
+            route_index = (
+                self._capacity_constrained_route(routing_logits, preferred_index)
+                if self.capacity_factor
+                else preferred_index
+            )
+        selected_router_probability = router_probs.gather(1, route_index.unsqueeze(1)).squeeze(1)
+        route_weight = (
+            selected_router_probability
+            if self.use_gate_weight
+            else torch.ones_like(selected_router_probability)
+        )
+
+        output = torch.zeros_like(flat_x)
+        for expert_index, expert in enumerate(self.experts):
+            token_indices = torch.nonzero(route_index == expert_index, as_tuple=False).flatten()
+            if token_indices.numel() == 0:
+                continue
+            expert_input = flat_x.index_select(0, token_indices)
+            real_token_count = expert_input.size(0)
+            # TorchAO's FP8 GEMMs require their M dimension to be divisible
+            # by 16.  Sparse routing gives every expert a variable number of
+            # tokens, so pad only this temporary dispatch buffer and discard
+            # the corresponding outputs.  This changes no real-token math.
+            fp8_padding = (-real_token_count) % 16
+            if fp8_padding:
+                expert_input = torch.cat(
+                    (expert_input, expert_input.new_zeros((fp8_padding, expert_input.size(-1)))), dim=0
+                )
+            expert_output = expert(expert_input)[:real_token_count]
+            routed_output = expert_output * route_weight.index_select(0, token_indices).to(expert_output.dtype).unsqueeze(-1)
+            # RMSNorm can promote the residual stream to FP32, while the FP8
+            # expert projections return BF16 under autocast.  Restore the
+            # residual dtype before scattering selected expert outputs.
+            output = output.index_copy(0, token_indices, routed_output.to(output.dtype))
+
+        expert_fraction = F.one_hot(route_index, num_classes=self.num_experts).to(router_probs.dtype).mean(dim=0)
+        preferred_fraction = F.one_hot(preferred_index, num_classes=self.num_experts).to(router_probs.dtype).mean(dim=0)
+        mean_router_prob = router_probs.mean(dim=0)
+        self.last_expert_fraction = expert_fraction.detach()
+        self.last_preferred_expert_fraction = preferred_fraction.detach()
+        self.last_forced_fraction = (route_index != preferred_index).float().mean().detach()
+        self.last_selected_gate_probability = selected_router_probability.mean().detach()
+        # Balance the router's *clean preference* rather than the capacity-
+        # constrained dispatch, which is intentionally already near-uniform.
+        aux_loss = self.router_aux_loss_coef * self.num_experts * torch.sum(
+            preferred_fraction * mean_router_prob
+        )
+        z_loss = self.router_z_loss_coef * torch.logsumexp(router_logits, dim=-1).square().mean()
+        return output.reshape(original_shape), aux_loss + z_loss
+
+
+class DecoderBlock(nn.Module):
+    def __init__(self, cfg: AuroraConfig) -> None:
+        super().__init__()
+        self.input_layernorm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
+        self.self_attn = CausalSelfAttention(cfg)
+        self.post_attention_layernorm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
+        self.mlp: nn.Module = Top1MoE(cfg) if cfg.num_experts > 1 else SwiGLU(cfg)
+
+    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        x = x + self.self_attn(self.input_layernorm(x), cos, sin)
+        mlp_input = self.post_attention_layernorm(x)
+        if isinstance(self.mlp, Top1MoE):
+            mlp_output, router_loss = self.mlp(mlp_input)
+        else:
+            mlp_output = self.mlp(mlp_input)
+            router_loss = x.new_zeros((), dtype=torch.float32)
+        return x + mlp_output, router_loss
+
+
+class AuroraForCausalLM(nn.Module):
+    def __init__(self, cfg: AuroraConfig) -> None:
+        super().__init__()
+        cfg.validate()
+        self.cfg = cfg
+        self.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.hidden_size)
+        self.layers = nn.ModuleList([DecoderBlock(cfg) for _ in range(cfg.num_layers)])
+        self.norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
+        self.lm_head = nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False)
+        self.last_router_loss: torch.Tensor | None = None
+        self.register_buffer("rope_cos_cached", torch.empty(0), persistent=False)
+        self.register_buffer("rope_sin_cached", torch.empty(0), persistent=False)
+        if cfg.tie_word_embeddings:
+            self.lm_head.weight = self.embed_tokens.weight
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module: nn.Module) -> None:
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+    def _rope_cache(
+        self, seq_len: int, device: torch.device, dtype: torch.dtype
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        cache_miss = (
+            self.rope_cos_cached.numel() == 0
+            or self.rope_cos_cached.size(0) < seq_len
+            or self.rope_cos_cached.device != device
+            or self.rope_cos_cached.dtype != dtype
+        )
+        if cache_miss:
+            cos, sin = precompute_rope_frequencies(
+                self.cfg.context_length,
+                self.cfg.head_dim,
+                self.cfg.rope_theta,
+                device,
+                dtype,
+            )
+            self.rope_cos_cached = cos
+            self.rope_sin_cached = sin
+        return self.rope_cos_cached[:seq_len], self.rope_sin_cached[:seq_len]
+
+    def _rope_dtype(self, x: torch.Tensor) -> torch.dtype:
+        if x.device.type == "cuda" and torch.is_autocast_enabled("cuda"):
+            return torch.get_autocast_dtype("cuda")
+        if x.device.type == "cpu" and torch.is_autocast_enabled("cpu"):
+            return torch.get_autocast_dtype("cpu")
+        return x.dtype
+
+    def forward(
+        self, input_ids: torch.Tensor, labels: torch.Tensor | None = None
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        x = self.embed_tokens(input_ids)
+        cos, sin = self._rope_cache(x.size(1), x.device, self._rope_dtype(x))
+        router_loss = torch.zeros((), device=x.device, dtype=torch.float32)
+        for layer in self.layers:
+            x, layer_router_loss = layer(x, cos, sin)
+            router_loss = router_loss + layer_router_loss
+        # Each layer produces a regularizer of the same scale.  Average them
+        # so the configured coefficient has the same meaning regardless of
+        # depth (instead of becoming 16x stronger in this MoE model).
+        router_loss = router_loss / max(1, len(self.layers))
+        self.last_router_loss = router_loss.detach()
+        x = self.norm(x)
+        if labels is not None and linear_cross_entropy is not None:
+            # RMSNorm may promote activations to FP32, but Cut Cross Entropy's
+            # backward kernel requires BF16/FP16 hidden states.
+            loss = linear_cross_entropy(x.to(self.lm_head.weight.dtype), self.lm_head.weight, labels, shift=True)
+            logits = x.new_empty(0)
+        else:
+            logits = self.lm_head(x)
+            loss = None
+            if labels is not None:
+                loss = F.cross_entropy(
+                    logits[:, :-1].contiguous().view(-1, logits.size(-1)),
+                    labels[:, 1:].contiguous().view(-1),
+                )
+        # Keep evaluation perplexity comparable to dense models: router regularization
+        # shapes gradients only during training and is not language-model loss.
+        if loss is not None and self.training:
+            loss = loss + router_loss
+        return logits, loss
+
+
+def count_parameters(model: nn.Module) -> int:
+    seen: set[int] = set()
+    total = 0
+    for param in model.parameters():
+        ident = id(param)
+        if ident not in seen:
+            seen.add(ident)
+            total += param.numel()
+    return total
+
+
+def count_active_parameters(model: nn.Module) -> int:
+    """Count parameters used by one top-1 path, without double-counting ties."""
+    seen: set[int] = set()
+    total = 0
+    for name, param in model.named_parameters():
+        if ".mlp.experts." in name:
+            expert_index = name.split(".mlp.experts.", 1)[1].split(".", 1)[0]
+            if expert_index != "0":
+                continue
+        ident = id(param)
+        if ident not in seen:
+            seen.add(ident)
+            total += param.numel()
+    return total
+
+
+def estimate_parameter_count(cfg: AuroraConfig) -> int:
+    model = AuroraForCausalLM(cfg)
+    return count_parameters(model)

infer.py

@@ -0,0 +1,197 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import re
+import sys
+import time
+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, count_parameters
+
+ROOT = Path(__file__).resolve().parent
+ROLE_RESTART = re.compile(
+    r"(?:^|\n)\s*(?:(?:question|answer|problem|solution)\s*:|#\s*(?:question|answer|problem|solution|what\s+is)\b)",
+    re.IGNORECASE,
+)
+
+
+def choose_device(requested: str) -> torch.device:
+    if requested != "auto":
+        device = torch.device(requested)
+        if device.type == "mps" and not torch.backends.mps.is_available():
+            raise RuntimeError("MPS was requested, but this PyTorch build cannot access Apple Silicon GPU acceleration.")
+        return device
+    if torch.backends.mps.is_available():
+        return torch.device("mps")
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    return torch.device("cpu")
+
+
+def choose_dtype(device: torch.device, requested: str) -> torch.dtype:
+    if requested == "float32":
+        return torch.float32
+    if requested == "float16":
+        return torch.float16
+    if requested == "bfloat16":
+        return torch.bfloat16
+    if device.type == "mps":
+        return torch.float16
+    if device.type == "cuda":
+        return torch.bfloat16
+    return torch.float32
+
+
+def find_checkpoint(explicit: Path | None) -> Path:
+    if explicit:
+        if not explicit.is_file():
+            raise FileNotFoundError(f"Checkpoint not found: {explicit}")
+        return explicit
+    candidates = sorted(ROOT.glob("*.safetensors"))
+    if not candidates:
+        raise FileNotFoundError(
+            "No .safetensors file found. Drag your model .safetensors file into this folder and run again."
+        )
+    if len(candidates) > 1:
+        names = "\n".join(f"  - {p.name}" for p in candidates)
+        raise RuntimeError(f"More than one .safetensors file was found. Use --checkpoint to choose one:\n{names}")
+    return candidates[0]
+
+
+def repeats_ngram(generated: list[int], candidate: int, n: int) -> bool:
+    if n <= 0 or len(generated) + 1 < n:
+        return False
+    trial = generated + [candidate]
+    target = tuple(trial[-n:])
+    return any(tuple(trial[i:i+n]) == target for i in range(len(trial) - n))
+
+
+def load_model(checkpoint_path: Path, device: torch.device, dtype: torch.dtype):
+    config = load_model_config(ROOT / "model_ember_proelia_207m_16k.yaml")
+    tokenizer = Tokenizer.from_file(str(ROOT / "tokenizer.json"))
+    expected = {"<pad>": 0, "<bos>": 1, "<eos>": 2}
+    actual = {token: tokenizer.token_to_id(token) for token in expected}
+    if actual != expected:
+        raise RuntimeError(f"Unexpected tokenizer special IDs: {actual}")
+
+    print(f"Loading {checkpoint_path.name}…")
+    state_dict = load_file(str(checkpoint_path), device="cpu")
+    model = AuroraForCausalLM(config)
+    missing, unexpected = model.load_state_dict(state_dict, strict=False)
+    missing = [key for key in missing if not key.endswith("._extra_state")]
+    if missing or unexpected:
+        raise RuntimeError(f"Checkpoint mismatch:\nmissing={missing}\nunexpected={unexpected}")
+    del state_dict
+    model = model.to(device=device, dtype=dtype).eval()
+    print(f"Loaded {count_parameters(model):,} parameters on {device} as {str(dtype).replace('torch.', '')}.\n")
+    return model, tokenizer, config, expected
+
+
+def sample_next(logits: torch.Tensor, temperature: float, top_k: int) -> int:
+    if temperature <= 0:
+        return int(torch.argmax(logits).item())
+    logits = logits.float() / temperature
+    if top_k > 0:
+        top_k = min(top_k, logits.numel())
+        values, indices = torch.topk(logits, top_k)
+        pick = torch.multinomial(torch.softmax(values, dim=-1), 1)
+        return int(indices[pick].item())
+    return int(torch.multinomial(torch.softmax(logits, dim=-1), 1).item())
+
+
+def generate(model, tokenizer, config, special, prompt: str, max_new_tokens: int,
+             temperature: float, top_k: int, no_repeat_ngram_size: int) -> tuple[str, str, int, float]:
+    prompt_ids = tokenizer.encode(prompt, add_special_tokens=False).ids
+    ids = [special["<bos>"], *prompt_ids]
+    generated: list[int] = []
+    stop_reason = "max_new_tokens"
+    answer_mode = prompt.rstrip().casefold().endswith("answer:")
+    started = time.perf_counter()
+
+    with torch.inference_mode():
+        for _ in range(max_new_tokens):
+            inputs = torch.tensor([ids[-config.context_length:]], device=next(model.parameters()).device, dtype=torch.long)
+            logits, _ = model(inputs)
+            next_id = sample_next(logits[0, -1], temperature, top_k)
+            if next_id == special["<eos>"]:
+                stop_reason = "eos"
+                break
+            if repeats_ngram(generated, next_id, no_repeat_ngram_size):
+                stop_reason = f"repeat_{no_repeat_ngram_size}gram"
+                break
+            generated.append(next_id)
+            ids.append(next_id)
+            if answer_mode:
+                text = tokenizer.decode(generated, skip_special_tokens=True)
+                match = ROLE_RESTART.search(text)
+                if match:
+                    text = text[:match.start()].rstrip()
+                    return text, "role_restart", len(generated), time.perf_counter() - started
+
+    text = tokenizer.decode(generated, skip_special_tokens=True).strip()
+    return text, stop_reason, len(generated), time.perf_counter() - started
+
+
+def build_prompt(user_text: str, raw_prompt: bool) -> str:
+    return user_text if raw_prompt else f"Question: {user_text.strip()}\nAnswer:"
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Run Ember Proelia SafeTensors inference on macOS, CUDA, or CPU.")
+    parser.add_argument("--checkpoint", type=Path)
+    parser.add_argument("--prompt")
+    parser.add_argument("--raw-prompt", action="store_true")
+    parser.add_argument("--max-new-tokens", type=int, default=128)
+    parser.add_argument("--temperature", type=float, default=0.0, help="0 = greedy decoding")
+    parser.add_argument("--top-k", type=int, default=40)
+    parser.add_argument("--no-repeat-ngram-size", type=int, default=4)
+    parser.add_argument("--device", choices=["auto", "mps", "cpu", "cuda"], default="auto")
+    parser.add_argument("--dtype", choices=["auto", "float16", "float32", "bfloat16"], default="auto")
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    checkpoint = find_checkpoint(args.checkpoint)
+    device = choose_device(args.device)
+    dtype = choose_dtype(device, args.dtype)
+    model, tokenizer, config, special = load_model(checkpoint, device, dtype)
+
+    def run(text: str) -> None:
+        prompt = build_prompt(text, args.raw_prompt)
+        completion, reason, count, elapsed = generate(
+            model, tokenizer, config, special, prompt,
+            args.max_new_tokens, args.temperature, args.top_k, args.no_repeat_ngram_size,
+        )
+        rate = count / elapsed if elapsed > 0 else 0.0
+        print(f"\nEmber: {completion}\n\n[{reason}; {count} tokens; {rate:.1f} tok/s]\n")
+
+    if args.prompt:
+        run(args.prompt)
+        return
+
+    print("Interactive Ember Proelia inference. Type /quit to exit.")
+    while True:
+        try:
+            text = input("\nYou: ").strip()
+        except (EOFError, KeyboardInterrupt):
+            print()
+            return
+        if text.casefold() in {"/quit", "/exit", "quit", "exit"}:
+            return
+        if text:
+            run(text)
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except Exception as exc:
+        print(f"\nERROR: {exc}", file=sys.stderr)
+        raise SystemExit(1)

requirements.txt

@@ -0,0 +1,4 @@
+torch>=2.5
+tokenizers>=0.20
+safetensors>=0.4
+PyYAML>=6.0