import argparse import contextlib import dataclasses import sys import time import torch from model import GPT, GPTConfig from tokenizer import TokenizerWrapper # ───────────────────────────────────────────────────────────────────────────── # Device selection # ───────────────────────────────────────────────────────────────────────────── def pick_device(requested: str | None) -> torch.device: if requested is not None: return torch.device(requested) if torch.cuda.is_available(): return torch.device("cuda") try: import torch_xla.core.xla_model as xm # noqa: F401 return torch.device("xla") except ImportError: pass if torch.backends.mps.is_available(): return torch.device("mps") return torch.device("cpu") # ───────────────────────────────────────────────────────────────────────────── # Checkpoint loading # ───────────────────────────────────────────────────────────────────────────── # Static attention params that older checkpoints (trained with those features off) # don't contain; they keep their fresh init on load. _STATIC_ATTN_KEY_MARKERS = (".attn.q_norm.", ".attn.k_norm.", ".attn.xsa_alpha", ".attn.attn_gate_proj.") def _torch_load(path: str): """Load on CPU, memory-mapped when possible: checkpoints carry optimizer state (~2x the weights) that inference never touches, and mmap means it's never read.""" try: return torch.load(path, map_location="cpu", weights_only=True, mmap=True) except (TypeError, RuntimeError, ValueError): # old torch / legacy non-zip file return torch.load(path, map_location="cpu", weights_only=True) def load_model( checkpoint_path: str, device: torch.device, tokenizer: TokenizerWrapper, ) -> tuple[GPT, GPTConfig]: checkpoint = _torch_load(checkpoint_path) if "model_config" not in checkpoint: raise RuntimeError( f"'{checkpoint_path}' has no saved model_config -- it was written " "before save_checkpoint() started saving the architecture config " "alongside the weights. Either retrain (new checkpoints save it " "automatically), or construct GPTConfig(...) by hand here to " "match whatever config.py looked like when this checkpoint was " "produced, and call GPT(that_config) instead of this function." ) saved_cfg = checkpoint["model_config"] known = {f.name for f in dataclasses.fields(GPTConfig)} dropped = sorted(set(saved_cfg) - known) if dropped: print(f"[Inference] Ignoring {len(dropped)} legacy config keys (features are now static): {dropped}", file=sys.stderr) model_config = GPTConfig(**{k: v for k, v in saved_cfg.items() if k in known}) engram_lookup = None if bool(getattr(model_config, "use_engram", False)): from ngram import CompressedTokenizer compressed_tokenizer = CompressedTokenizer(tokenizer) engram_lookup = torch.as_tensor(compressed_tokenizer.lookup_table) model = GPT(model_config, engram_lookup=engram_lookup).to(device) state_dict = checkpoint["model_state_dict"] if any(k.startswith("_orig_mod.") for k in state_dict): state_dict = { (k[len("_orig_mod."):] if k.startswith("_orig_mod.") else k): v for k, v in state_dict.items() } print( "[Inference] Checkpoint was saved from a torch.compile()'d model " "(keys prefixed with '_orig_mod.') -- stripped the prefix to " "load onto this plain GPT instance.", file=sys.stderr, ) # Load fp32 weights into the plain model BEFORE quantizing. # Quantizing first replaces nn.Linear submodules with tensor-subclass # equivalents; load_state_dict then tries to copy float weights into them # which either raises or silently de-quantizes. Correct order: weights → quant. missing, unexpected = model.load_state_dict(state_dict, strict=False) bad_missing = [k for k in missing if not any(m in k for m in _STATIC_ATTN_KEY_MARKERS)] if bad_missing or unexpected: raise RuntimeError(f"state_dict mismatch: missing={bad_missing[:6]} unexpected={list(unexpected)[:6]}") if missing: print(f"[Inference] {len(missing)} static-attention keys absent from checkpoint (legacy run with those " f"features off) -- using fresh init (xsa_alpha/attn gate = identity; q_norm/k_norm are NOT, so " f"outputs differ from what that checkpoint produced when trained).", file=sys.stderr) quant = checkpoint.get("quantization") if quant: import model as _m fn = {"int8": _m.apply_int8_to_model, "nvfp4": _m.apply_nvfp4_to_model}.get(quant) if fn is None: raise RuntimeError(f"checkpoint quantized with {quant!r} but this build can't re-apply it") fn(model) model._quant_mode = quant model.eval() n_params = model.get_num_params() / 1e6 step = checkpoint.get("step", "?") del state_dict, checkpoint # release the mmap print(f"[Inference] Loaded checkpoint (step {step}, {n_params:.2f}M params) onto {device}", file=sys.stderr) return model, model_config # ───────────────────────────────────────────────────────────────────────────── # Sampling # ───────────────────────────────────────────────────────────────────────────── def _top_k_filter(logits: torch.Tensor, top_k: int) -> torch.Tensor: k = min(top_k, logits.size(-1)) values, _ = torch.topk(logits, k) logits[logits < values[:, [-1]]] = float("-inf") return logits def _top_p_filter(logits: torch.Tensor, top_p: float) -> torch.Tensor: """Nucleus sampling: keep the smallest prefix of sorted tokens whose cumulative probability crosses top_p, drop the rest.""" sorted_logits, sorted_idx = torch.sort(logits, descending=True) probs = torch.softmax(sorted_logits, dim=-1) cum_probs = torch.cumsum(probs, dim=-1) # remove token i if the cumulative prob BEFORE it already exceeds top_p # (i.e. it wasn't needed to cross the threshold) -- matches the standard # HF TopPLogitsWarper convention. remove = (cum_probs - probs) > top_p sorted_logits = sorted_logits.masked_fill(remove, float("-inf")) return torch.full_like(logits, float("-inf")).scatter(1, sorted_idx, sorted_logits) def _clip_max_new_tokens(max_new_tokens: int, prompt_ids: list[int], block_size: int) -> int: """Prevent the generation loop from silently overwriting the prompt by clipping the available completion budget to the remaining context window. """ if block_size <= 0: return 0 available = block_size - len(prompt_ids) if available <= 0: raise ValueError( f"prompt is {len(prompt_ids)} tokens >= block_size {block_size}; nothing " f"left to generate. Shorten the prompt (this build has no sliding-window " f"generation)." ) if max_new_tokens > available: print(f"[generate] NOTE: clipped max_new_tokens {max_new_tokens} -> {available} " f"(block_size={block_size}, prompt={len(prompt_ids)}).", file=sys.stderr) return min(max_new_tokens, available) def _ngram_map_init(ids: list[int], n: int) -> dict: """{(n-1)-token prefix -> set of tokens that followed it} over every n-gram in ids.""" m: dict = {} for j in range(n - 1, len(ids)): m.setdefault(tuple(ids[j - n + 1:j]), set()).add(ids[j]) return m class _IncrementalDetok: """Streams text without re-decoding the whole sequence every step: decode a short window, emit what's new, and hold back anything ending in an incomplete multi-byte character.""" def __init__(self, decode, ids: list[int], lookback: int = 6): self.decode = decode self.ids = list(ids) self.lookback = lookback self.read = len(ids) def push(self, tid: int) -> str: self.ids.append(tid) prefix = max(0, len(self.ids) - self.lookback) prev = self.decode(self.ids[prefix:self.read]) cur = self.decode(self.ids[prefix:]) if len(cur) > len(prev) and not cur.endswith("\ufffd"): self.read = len(self.ids) return cur[len(prev):] return "" def _special_ids(tokenizer) -> set[int]: out = set() for name in ("bos_id", "eos_id", "pad_id", "bos_token_id", "eos_token_id", "pad_token_id"): v = getattr(tokenizer, name, None) if isinstance(v, int): out.add(v) return out def _cut_at_eos(seq: list[int], start: int, eos_id: int | None) -> list[int]: """Drop the first EOS at or after `start` and everything after it.""" if eos_id is None: return seq for i in range(start, len(seq)): if seq[i] == eos_id: return seq[:i] return seq @torch.inference_mode() def generate( model: GPT, tokenizer: TokenizerWrapper, prompt: str, device: torch.device, max_new_tokens: int = 200, temperature: float = 1.0, top_k: int | None = None, top_p: float | None = None, repetition_penalty: float = 1.0, no_repeat_ngram_size: int = 0, stream: bool = True, seed: int | None = None, num_samples: int = 1, use_cache: bool = True, stats: bool = False, on_text=None, ) -> list[str]: """KV-cached sampler: one prefill over the prompt, then one-token forwards. Sampling order is unchanged: repetition penalty -> n-gram block -> temperature -> top-k -> top-p -> multinomial. EOS ends a sample and is never part of the output. """ if seed is not None: torch.manual_seed(seed) block_size = model.config.block_size bos_id = getattr(tokenizer, "bos_id", None) eos_id = getattr(tokenizer, "eos_id", None) if eos_id is None: eos_id = getattr(tokenizer, "eos_token_id", None) if eos_id is None: print("[generate] WARNING: tokenizer exposes no eos_id/eos_token_id -- " "early stopping disabled, generations run to budget.", file=sys.stderr) def decode(ids: list[int]) -> str: # Only the leading BOS is dropped. Other ids are never filtered: a special id can coincide # with an ordinary token in some tokenizers, and filtering by id deletes real text. if ids and bos_id is not None and ids[0] == bos_id: ids = ids[1:] return tokenizer.decode(ids) # Keep room to generate: a prompt that fills the whole context leaves nothing to sample. prompt_ids = tokenizer.encode(prompt, add_bos=True, add_eos=False) max_prompt = max(1, block_size - max(1, min(max_new_tokens, block_size // 2))) if len(prompt_ids) > max_prompt: head = prompt_ids[:1] if prompt_ids and prompt_ids[0] == getattr(tokenizer, "bos_id", None) else [] keep = max_prompt - len(head) print(f"[generate] NOTE: prompt is {len(prompt_ids)} tokens; keeping the last {keep} " f"(+BOS) to leave room for generation (block_size={block_size}).", file=sys.stderr) prompt_ids = head + prompt_ids[-keep:] max_new_tokens = _clip_max_new_tokens(max_new_tokens, prompt_ids, block_size) P = len(prompt_ids) B = num_samples V = model.config.vocab_size t0 = time.perf_counter() cache = None supports_kv_cache = bool(getattr(model, "supports_kv_cache", False)) if use_cache and supports_kv_cache: cache = model.new_kv_cache(max_len=P + max_new_tokens) elif use_cache: print("[generate] NOTE: model does not support KV cache (Engram or non-causal attention) " "-- re-encoding the context each step. This is slower but correct.", file=sys.stderr) use_pen = repetition_penalty != 1.0 n_ng = no_repeat_ngram_size if no_repeat_ngram_size > 1 else 0 # n==1 was already a no-op seqs = [list(prompt_ids) for _ in range(B)] # generated tokens are appended, EOS excluded finished = [False] * B fin_t = torch.zeros(B, dtype=torch.bool, device=device) rows = torch.arange(B, device=device) counts = torch.zeros(B, V, dtype=torch.float32, device=device) if use_pen else None ngram_maps = [_ngram_map_init(prompt_ids, n_ng) for _ in range(B)] if n_ng else None # Streaming: print to stdout (stream=True) and/or call on_text(piece) for each new piece of text # (used by the web app). Both need B == 1. print_out = stream and B == 1 detok = None streamed = "" if B == 1 and (print_out or on_text is not None): detok = _IncrementalDetok(decode, prompt_ids) streamed = decode(prompt_ids) if print_out: print(streamed, end="", flush=True) def emit(piece: str) -> None: if print_out: print(piece, end="", flush=True) if on_text is not None: on_text(piece) # Host values are needed per token only for streaming and n-gram blocking. Otherwise # sampling stays on the device and EOS is checked every 16 steps. need_host = detok is not None or n_ng > 0 pending: list[torch.Tensor] = [] idx = torch.tensor([prompt_ids] * B, dtype=torch.long, device=device) last_tok = idx[:, -1] n_done = 0 for step in range(max_new_tokens): if need_host: if all(finished): break elif eos_id is not None and step % 16 == 0 and bool(fin_t.all()): break if cache is not None: inp = idx if step == 0 else last_tok.unsqueeze(1) logits, _ = model(inp, kv_cache=cache) else: # no-cache fallback: re-encode the growing context idx_cond = idx if idx.size(1) <= block_size else idx[:, -block_size:] logits, _ = model(idx_cond) logits = logits[:, -1, :].float() # fp32 sampling math regardless of model dtype if use_pen: # Count-scaled penalty over generated tokens only (prompt is unpenalized); capped at 6 occurrences. factor = repetition_penalty ** counts.clamp(max=6.0) logits = torch.where(logits > 0, logits / factor, logits * factor) if n_ng: br, bc = [], [] for b in range(B): if finished[b] or len(seqs[b]) < n_ng - 1: continue banned = ngram_maps[b].get(tuple(seqs[b][-(n_ng - 1):])) if banned: br.extend([b] * len(banned)); bc.extend(banned) if br: logits[torch.tensor(br, device=device), torch.tensor(bc, device=device)] = float("-inf") logits = logits / max(temperature, 1e-5) if top_k is not None and top_k > 0: logits = _top_k_filter(logits, top_k) if top_p is not None: logits = _top_p_filter(logits, top_p) probs = torch.softmax(logits, dim=-1) dead = ~torch.isfinite(probs).any(dim=-1) | (probs.sum(dim=-1) <= 0) probs = torch.where(dead.unsqueeze(1), torch.full_like(probs, 1.0 / V), probs) sampled = torch.multinomial(probs, num_samples=1).squeeze(1) # Finished rows repeat their last token (batch stays rectangular) and stop counting. active = ~fin_t nxt = torch.where(fin_t, last_tok, sampled) if use_pen: counts[rows, nxt] += active.float() if eos_id is not None: fin_t = fin_t | (active & (nxt == eos_id)) last_tok = nxt if cache is None: idx = torch.cat([idx, nxt.unsqueeze(1)], dim=1) n_done += 1 if need_host: toks = nxt.tolist() # the one host sync this step for b in range(B): if finished[b]: continue t = toks[b] if eos_id is not None and t == eos_id: finished[b] = True continue seqs[b].append(t) if n_ng and len(seqs[b]) >= n_ng: ngram_maps[b].setdefault(tuple(seqs[b][-n_ng:-1]), set()).add(t) if detok is not None and b == 0: piece = detok.push(t) if piece: streamed += piece emit(piece) else: pending.append(nxt) if pending: for b, row in enumerate(torch.stack(pending, dim=1).tolist()): seqs[b].extend(row) seqs = [_cut_at_eos(sq, P, eos_id) for sq in seqs] texts = [decode(sq) for sq in seqs] if detok is not None: # Safety net for tokenizers whose decode isn't prefix-stable: flush what the incremental # path held back, or re-send the full decode if it diverged. if texts[0].startswith(streamed): tail = texts[0][len(streamed):] if tail: emit(tail) else: print(f"\n[generate] WARNING: incremental decode diverged from full decode at " f"len={len(streamed)}. Printing the full decode instead.", file=sys.stderr) if on_text is not None: on_text(f"\n{texts[0]}") else: print() print(texts[0], end="") if print_out: print() if stats: if device.type == "cuda": torch.cuda.synchronize(device) dt = time.perf_counter() - t0 print(f"[generate] {n_done} steps x {B} sample(s) in {dt:.2f}s " f"({n_done / max(dt, 1e-9):.1f} steps/s, {n_done * B / max(dt, 1e-9):.1f} tok/s, " f"kv_cache={'on' if cache is not None else 'off'})", file=sys.stderr) return texts # ───────────────────────────────────────────────────────────────────────────── # CLI # ───────────────────────────────────────────────────────────────────────────── def build_arg_parser() -> argparse.ArgumentParser: p = argparse.ArgumentParser(description="Generate text from a trained Arya/Veylon checkpoint.") p.add_argument("--checkpoint", required=True, help="Path to a checkpoint saved by train.py") p.add_argument("--tokenizer", default="./tokenizer", help="Path passed to TokenizerWrapper") p.add_argument("--prompt", default=None, help="Prompt text (omit if using --interactive)") p.add_argument("--interactive", action="store_true", help="Drop into a REPL: one prompt per line") p.add_argument("--max-new-tokens", type=int, default=200) p.add_argument("--temperature", type=float, default=0.8) p.add_argument("--top-k", type=int, default=50) p.add_argument("--top-p", type=float, default=None, help="Nucleus sampling threshold, e.g. 0.9. Combine freely with --top-k.") p.add_argument("--repetition-penalty", type=float, default=1.3, help="CTRL-style penalty on already-generated tokens, now count-scaled " "(penalty**occurrences). 1.0 disables it -- small models loop into " "degenerate repetition ('so so so so...') without some penalty here, " "so this defaults ON.") p.add_argument("--no-repeat-ngram-size", type=int, default=3, help="Hard-ban repeating any n-gram of this size (HF-style). 0 disables it. " "This is what actually stops infinite word loops -- repetition-penalty " "alone only makes them less likely, this makes them impossible.") p.add_argument("--num-samples", type=int, default=1, help="Generate N completions per prompt (batched)") p.add_argument("--seed", type=int, default=None) p.add_argument("--no-cache", action="store_true", help="Disable the KV cache (re-encode the context every step; for debugging/A-B).") p.add_argument("--stats", action="store_true", help="Print steps/s and cache status to stderr.") p.add_argument("--device", default=None, help="cuda / cpu / mps / xla -- auto-detected if omitted") p.add_argument("--dtype", default="auto", choices=["auto", "fp32", "fp16", "bf16"], help="Inference precision. 'auto' uses fp16 on CUDA, fp32 elsewhere.") return p def _resolve_dtype(dtype_arg: str, device: torch.device) -> torch.dtype | None: if device.type not in ("cuda", "cpu"): # torch.autocast is only exercised here against cuda/cpu; MPS and XLA # handle precision differently (XLA in particular usually wants # XLA_USE_BF16 or torch_xla's own autocast, not torch.autocast). # Guessing at that without real hardware to test against would be # worse than just running those in fp32 and saying so. if dtype_arg != "auto": print(f"[Inference] --dtype={dtype_arg} requested but autocast isn't wired up for device " f"'{device.type}' -- running in fp32.", file=sys.stderr) return None if dtype_arg == "fp32": return None # no autocast if dtype_arg == "fp16": if device.type == "cpu": print("[Inference] fp16 autocast unsupported on CPU -- using bf16 or fp32.", file=sys.stderr) return torch.bfloat16 return torch.float16 if dtype_arg == "bf16": return torch.bfloat16 # auto if device.type == "cuda": major, _ = torch.cuda.get_device_capability(device) return torch.bfloat16 if major >= 8 else torch.float16 return None def main() -> None: args = build_arg_parser().parse_args() if not args.prompt and not args.interactive: print("Provide --prompt \"...\" or pass --interactive.", file=sys.stderr) sys.exit(1) device = pick_device(args.device) tokenizer = TokenizerWrapper(args.tokenizer) model, model_config = load_model(args.checkpoint, device, tokenizer) autocast_dtype = _resolve_dtype(args.dtype, device) def run(prompt: str, stream: bool) -> list[str]: ctx = torch.autocast(device_type=device.type, dtype=autocast_dtype) if autocast_dtype else contextlib.nullcontext() with ctx: return generate( model, tokenizer, prompt, device, max_new_tokens=args.max_new_tokens, temperature=args.temperature, top_k=args.top_k, top_p=args.top_p, repetition_penalty=args.repetition_penalty, no_repeat_ngram_size=args.no_repeat_ngram_size, stream=stream, seed=args.seed, num_samples=args.num_samples, use_cache=not args.no_cache, stats=args.stats, ) if args.interactive: print("[Inference] Interactive mode. Ctrl-C or empty line to exit.", file=sys.stderr) while True: try: prompt = input("\n>>> ") except (EOFError, KeyboardInterrupt): break if not prompt.strip(): break outputs = run(prompt, stream=(args.num_samples == 1)) if args.num_samples > 1: for i, text in enumerate(outputs): print(f"\n--- sample {i} ---\n{text}") else: outputs = run(args.prompt, stream=(args.num_samples == 1)) if args.num_samples > 1: for i, text in enumerate(outputs): print(f"\n--- sample {i} ---\n{text}") if __name__ == "__main__": main()