Spaces:
Sleeping
Sleeping
| 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 | |
| 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() |