Upload jumplander_python_100m.py

jumplander_python_100m.py

@@ -0,0 +1,1446 @@
+"""
+JumpLander Python Decoder 100M
+================================
+A single-file, from-scratch Python code language model project.
+
+- Randomly initialized decoder-only Transformer (~97.5M parameters)
+- Custom Byte-Level BPE tokenizer trained from local JSONL files
+- Local train/validation/test datasets beside this script
+- Optional Hugging Face streaming preparation from Python-Edu + MBPP
+- Smoke training, full training, resume, generation, and local web UI
+- No pretrained model weights are loaded
+
+This is an experimental research model. Training 100M parameters from scratch
+requires substantial data and compute even though it fits on an RTX 3060 12GB.
+"""
+
+from __future__ import annotations
+
+import argparse
+import ast
+import contextlib
+import dataclasses
+from dataclasses import dataclass
+import hashlib
+import html
+import io
+import json
+import math
+import os
+from pathlib import Path
+import random
+import re
+import secrets
+import sys
+import threading
+import time
+import traceback
+from typing import Any, Iterable, Iterator, Optional
+import webbrowser
+
+try:
+    import numpy as np
+except ImportError as exc:
+    raise SystemExit("Missing dependency: numpy. Run: pip install -r requirements.txt") from exc
+
+try:
+    import torch
+    import torch.nn as nn
+    import torch.nn.functional as F
+    from torch.utils.data import DataLoader, Dataset
+    from torch.utils.checkpoint import checkpoint as activation_checkpoint
+except ImportError as exc:
+    raise SystemExit("Missing dependency: torch. Run: pip install -r requirements.txt") from exc
+
+
+# -----------------------------------------------------------------------------
+# Paths: all generated files remain beside this single Python source file.
+# -----------------------------------------------------------------------------
+ROOT = Path(__file__).resolve().parent
+TRAIN_JSONL = ROOT / "train.jsonl"
+VALIDATION_JSONL = ROOT / "validation.jsonl"
+TEST_JSONL = ROOT / "test.jsonl"
+TOKENIZER_JSON = ROOT / "tokenizer.json"
+TRAIN_BIN = ROOT / "train_tokens.bin"
+VALIDATION_BIN = ROOT / "validation_tokens.bin"
+SMOKE_CHECKPOINT = ROOT / "smoke_checkpoint.pt"
+MODEL_CHECKPOINT = ROOT / "jumplander_python_100m.pt"
+
+MODEL_NAME = "JumpLander Python Decoder 100M"
+MODEL_ID = "jumplander-python-decoder-100m"
+
+SPECIAL_TOKENS = [
+    "<pad>",
+    "<unk>",
+    "<bos>",
+    "<eos>",
+    "<file_start>",
+    "<file_end>",
+    "<fim_prefix>",
+    "<fim_suffix>",
+    "<fim_middle>",
+    "<instruction>",
+    "<response>",
+]
+
+SECRET_PATTERNS = [
+    re.compile(r"AKIA[0-9A-Z]{16}"),
+    re.compile(r"sk-[A-Za-z0-9_-]{20,}"),
+    re.compile(r"gh[pousr]_[A-Za-z0-9]{20,}"),
+    re.compile(r"-----BEGIN (?:RSA |EC |OPENSSH )?PRIVATE KEY-----"),
+    re.compile(r"(?i)(?:api[_-]?key|secret|password)\s*=\s*['\"][^'\"]{8,}['\"]"),
+]
+
+GENERATED_MARKERS = (
+    "generated file",
+    "auto-generated",
+    "autogenerated",
+    "do not edit",
+    "generated by",
+    "this file was generated",
+)
+
+
+# -----------------------------------------------------------------------------
+# Configuration
+# -----------------------------------------------------------------------------
+@dataclass
+class ModelConfig:
+    name: str = MODEL_NAME
+    vocab_size: int = 16_384
+    max_seq_len: int = 1_024
+    n_layers: int = 12
+    d_model: int = 768
+    n_heads: int = 12
+    d_ff: int = 2_048
+    rope_theta: float = 10_000.0
+    norm_eps: float = 1e-5
+    dropout: float = 0.0
+    tie_embeddings: bool = True
+    gradient_checkpointing: bool = True
+
+    def validate(self) -> None:
+        if self.d_model % self.n_heads != 0:
+            raise ValueError("d_model must be divisible by n_heads")
+        if (self.d_model // self.n_heads) % 2 != 0:
+            raise ValueError("attention head dimension must be even for RoPE")
+        if self.vocab_size > np.iinfo(np.uint16).max:
+            raise ValueError("This project stores token IDs as uint16; vocab is too large")
+
+
+@dataclass
+class TrainConfig:
+    batch_size: int = 1
+    gradient_accumulation: int = 32
+    learning_rate: float = 3e-4
+    min_learning_rate: float = 3e-5
+    weight_decay: float = 0.1
+    warmup_steps: int = 200
+    total_steps: int = 10_000
+    eval_interval: int = 250
+    eval_batches: int = 20
+    save_interval: int = 500
+    log_interval: int = 10
+    grad_clip: float = 1.0
+    seed: int = 1337
+    num_workers: int = 0
+
+
+FULL_MODEL_CONFIG = ModelConfig()
+SMOKE_MODEL_CONFIG = ModelConfig(
+    name="JumpLander Python Decoder Smoke",
+    max_seq_len=256,
+    n_layers=4,
+    d_model=256,
+    n_heads=4,
+    d_ff=768,
+    gradient_checkpointing=False,
+)
+
+
+# -----------------------------------------------------------------------------
+# Utility functions
+# -----------------------------------------------------------------------------
+def require_optional(package: str, install_name: Optional[str] = None) -> Any:
+    try:
+        return __import__(package)
+    except ImportError as exc:
+        target = install_name or package
+        raise SystemExit(
+            f"Missing optional dependency: {target}. Run: pip install -r requirements.txt"
+        ) from exc
+
+
+def seed_everything(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def atomic_write_jsonl(path: Path, rows: Iterable[dict[str, Any]]) -> int:
+    temp = path.with_suffix(path.suffix + ".tmp")
+    count = 0
+    with temp.open("w", encoding="utf-8", newline="\n") as handle:
+        for row in rows:
+            handle.write(json.dumps(row, ensure_ascii=False) + "\n")
+            count += 1
+    temp.replace(path)
+    return count
+
+
+def append_jsonl(path: Path, row: dict[str, Any]) -> None:
+    with path.open("a", encoding="utf-8", newline="\n") as handle:
+        handle.write(json.dumps(row, ensure_ascii=False) + "\n")
+
+
+def iter_jsonl(path: Path) -> Iterator[dict[str, Any]]:
+    if not path.exists():
+        raise FileNotFoundError(f"Dataset file not found: {path}")
+    with path.open("r", encoding="utf-8") as handle:
+        for line_number, line in enumerate(handle, start=1):
+            line = line.strip()
+            if not line:
+                continue
+            try:
+                value = json.loads(line)
+            except json.JSONDecodeError as exc:
+                raise ValueError(f"Invalid JSON at {path.name}:{line_number}: {exc}") from exc
+            if not isinstance(value, dict):
+                raise ValueError(f"Expected a JSON object at {path.name}:{line_number}")
+            yield value
+
+
+def sha256_text(text: str) -> str:
+    return hashlib.sha256(text.encode("utf-8", errors="ignore")).hexdigest()
+
+
+def normalize_code(code: str) -> str:
+    code = code.replace("\r\n", "\n").replace("\r", "\n").replace("\x00", "")
+    lines = [line.rstrip() for line in code.splitlines()]
+    return "\n".join(lines).strip() + "\n"
+
+
+def contains_secret(code: str) -> bool:
+    return any(pattern.search(code) for pattern in SECRET_PATTERNS)
+
+
+def looks_generated(code: str) -> bool:
+    head = code[:2_000].lower()
+    return any(marker in head for marker in GENERATED_MARKERS)
+
+
+def looks_minified(code: str) -> bool:
+    lines = code.splitlines()
+    if not lines:
+        return True
+    longest = max(len(line) for line in lines)
+    average = sum(len(line) for line in lines) / len(lines)
+    return longest > 1_000 or average > 240
+
+
+def is_valid_python(code: str) -> tuple[bool, str]:
+    try:
+        ast.parse(code)
+        return True, "ok"
+    except SyntaxError:
+        return False, "syntax"
+
+
+def basic_code_filter(code: str, min_chars: int = 120, max_chars: int = 50_000) -> tuple[bool, str]:
+    if not isinstance(code, str):
+        return False, "not_string"
+    code = normalize_code(code)
+    if len(code) < min_chars:
+        return False, "too_short"
+    if len(code) > max_chars:
+        return False, "too_large"
+    if contains_secret(code):
+        return False, "secret"
+    if looks_generated(code):
+        return False, "generated"
+    if looks_minified(code):
+        return False, "minified"
+    valid, reason = is_valid_python(code)
+    if not valid:
+        return False, reason
+    return True, "ok"
+
+
+def likely_english(text: str) -> bool:
+    if not text or len(text.split()) < 4:
+        return False
+    ascii_letters = sum(ch.isascii() and ch.isalpha() for ch in text)
+    letters = sum(ch.isalpha() for ch in text)
+    if letters == 0:
+        return False
+    return ascii_letters / letters >= 0.90
+
+
+def extract_instruction_rows(code: str, source: str) -> list[dict[str, Any]]:
+    """Extract function/class docstrings and their source as English->Python rows."""
+    rows: list[dict[str, Any]] = []
+    try:
+        tree = ast.parse(code)
+        source_lines = code.splitlines()
+    except SyntaxError:
+        return rows
+
+    for node in ast.walk(tree):
+        if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+            continue
+        doc = ast.get_docstring(node, clean=True)
+        if not doc or not likely_english(doc):
+            continue
+        first_line = doc.strip().splitlines()[0].strip()
+        if len(first_line) < 20 or len(first_line) > 320:
+            continue
+        if first_line.lower().startswith(("todo", "fixme", "deprecated")):
+            continue
+        if not hasattr(node, "end_lineno") or node.end_lineno is None:
+            continue
+        start = max(0, node.lineno - 1)
+        end = min(len(source_lines), node.end_lineno)
+        function_source = "\n".join(source_lines[start:end]).strip()
+        if len(function_source) < 80 or len(function_source) > 8_000:
+            continue
+        valid, _ = is_valid_python(function_source)
+        if not valid:
+            continue
+        rows.append(
+            {
+                "type": "instruct",
+                "instruction": first_line,
+                "response": function_source + "\n",
+                "source": source,
+                "tests": [],
+            }
+        )
+        if len(rows) >= 8:
+            break
+    return rows
+
+
+def split_python_code(code: str, max_chars: int = 12_000) -> list[str]:
+    """Keep small files whole; split large valid files at top-level AST boundaries."""
+    code = normalize_code(code)
+    if len(code) <= max_chars:
+        return [code]
+    try:
+        tree = ast.parse(code)
+    except SyntaxError:
+        return []
+    lines = code.splitlines()
+    imports: list[str] = []
+    chunks: list[str] = []
+    for node in tree.body:
+        if not hasattr(node, "end_lineno") or node.end_lineno is None:
+            continue
+        segment = "\n".join(lines[node.lineno - 1 : node.end_lineno]).strip()
+        if isinstance(node, (ast.Import, ast.ImportFrom)):
+            imports.append(segment)
+            continue
+        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
+            prefix = "\n".join(imports[-20:])
+            candidate = (prefix + "\n\n" + segment).strip() + "\n"
+            if 120 <= len(candidate) <= max_chars:
+                valid, _ = is_valid_python(candidate)
+                if valid:
+                    chunks.append(candidate)
+    return chunks
+
+
+def serialize_training_row(row: dict[str, Any]) -> str:
+    row_type = row.get("type", "base")
+    if row_type == "instruct":
+        instruction = str(row.get("instruction", "")).strip()
+        response = str(row.get("response", "")).strip()
+        return f"<bos><instruction>\n{instruction}\n<response>\n{response}\n<eos>"
+    text = str(row.get("text", "")).strip()
+    return f"<bos><file_start>\n{text}\n<file_end><eos>"
+
+
+def make_fim_variant(text: str) -> Optional[str]:
+    if len(text) < 300:
+        return None
+    digest = int(hashlib.sha256(text.encode("utf-8")).hexdigest()[:8], 16)
+    rng = random.Random(digest)
+    low = max(50, len(text) // 5)
+    high = min(len(text) - 50, (len(text) * 4) // 5)
+    if high <= low:
+        return None
+    start = rng.randint(low // 2, low)
+    end = rng.randint(max(start + 20, high - low), high)
+    if end <= start:
+        return None
+    prefix, middle, suffix = text[:start], text[start:end], text[end:]
+    return (
+        "<bos><fim_prefix>" + prefix + "<fim_suffix>" + suffix + "<fim_middle>" + middle + "<eos>"
+    )
+
+
+# -----------------------------------------------------------------------------
+# Dataset preparation and checks
+# -----------------------------------------------------------------------------
+def dataset_preview(limit: int) -> None:
+    datasets = require_optional("datasets")
+    print("Connecting to Hugging Face: codeparrot/codeparrot-clean [Python source code]")
+    stream = datasets.load_dataset(
+        "codeparrot/codeparrot-clean",
+        split="train",
+        streaming=True,
+    )
+    for index, sample in enumerate(stream):
+        code = str(sample.get("content") or sample.get("code") or sample.get("text") or "")
+        print("\n" + "=" * 88)
+        print(f"SAMPLE {index + 1}")
+        print("source:", sample.get("repo_name") or sample.get("repo") or sample.get("repository_name") or "unknown")
+        print("path:", sample.get("path") or "unknown")
+        print("score:", sample.get("int_score") or sample.get("score") or "unknown")
+        print("-" * 88)
+        print(code[:2_500])
+        if index + 1 >= limit:
+            break
+
+
+def prepare_remote_dataset(max_base: int, test_mode: bool, min_score: int) -> None:
+    datasets = require_optional("datasets")
+    target_base = min(max_base, 2_000) if test_mode else max_base
+    print(f"Preparing {'test' if test_mode else 'full'} local dataset from Hugging Face")
+    print(f"Target accepted Python base samples: {target_base:,}")
+
+    counters: dict[str, int] = {}
+    train_rows: list[dict[str, Any]] = []
+    validation_rows: list[dict[str, Any]] = []
+    test_rows: list[dict[str, Any]] = []
+    seen: set[str] = set()
+
+    def count(reason: str) -> None:
+        counters[reason] = counters.get(reason, 0) + 1
+
+    stream = datasets.load_dataset(
+        "codeparrot/codeparrot-clean",
+        split="train",
+        streaming=True,
+    )
+
+    for raw in stream:
+        count("read")
+        score_value = raw.get("int_score", raw.get("score"))
+        if score_value is not None:
+            with contextlib.suppress(TypeError, ValueError):
+                if int(float(score_value)) < min_score:
+                    count("low_score")
+                    continue
+        code = normalize_code(str(raw.get("content") or raw.get("code") or raw.get("text") or ""))
+        accepted, reason = basic_code_filter(code)
+        if not accepted:
+            count(reason)
+            continue
+        source = str(raw.get("repo_name") or raw.get("repo") or raw.get("repository_name") or "codeparrot-clean")
+        path = str(raw.get("path") or "unknown.py")
+        for chunk in split_python_code(code):
+            digest = sha256_text(chunk)
+            if digest in seen:
+                count("duplicate")
+                continue
+            seen.add(digest)
+            split_bucket = int(digest[:8], 16) % 1000
+            base_row = {
+                "type": "base",
+                "text": chunk,
+                "source": f"codeparrot-clean:{source}:{path}",
+                "sha256": digest,
+            }
+            destination = train_rows
+            if split_bucket < 15:
+                destination = test_rows
+            elif split_bucket < 30:
+                destination = validation_rows
+            destination.append(base_row)
+            if destination is train_rows:
+                for instruction_row in extract_instruction_rows(chunk, base_row["source"]):
+                    train_rows.append(instruction_row)
+            count("accepted")
+            if counters["accepted"] >= target_base:
+                break
+        if counters.get("accepted", 0) >= target_base:
+            break
+        if counters.get("read", 0) % 2_000 == 0:
+            print(
+                f"read={counters['read']:,} accepted={counters.get('accepted', 0):,} "
+                f"syntax={counters.get('syntax', 0):,} duplicate={counters.get('duplicate', 0):,}"
+            )
+
+    print("Adding MBPP English-to-Python examples and held-out tests...")
+    mbpp_loaded = False
+    mbpp_error = ""
+    for repo, config in [
+        ("google-research-datasets/mbpp", "full"),
+        ("RLAIF/mbpp", None),
+        ("Muennighoff/mbpp", None),
+    ]:
+        try:
+            kwargs: dict[str, Any] = {}
+            if config:
+                kwargs["name"] = config
+            mbpp = datasets.load_dataset(repo, **kwargs)
+            mbpp_loaded = True
+            for split_name in mbpp.keys():
+                split = mbpp[split_name]
+                for sample in split:
+                    instruction = str(sample.get("text", "")).strip()
+                    response = normalize_code(str(sample.get("code", "")))
+                    tests = list(sample.get("test_list") or [])
+                    valid, _ = is_valid_python(response)
+                    if not instruction or not valid:
+                        continue
+                    row = {
+                        "type": "instruct",
+                        "instruction": instruction,
+                        "response": response,
+                        "tests": tests,
+                        "source": f"{repo}:{split_name}:{sample.get('task_id', '')}",
+                    }
+                    split_lower = split_name.lower()
+                    if split_lower in {"test", "prompt"}:
+                        test_rows.append(row)
+                    elif split_lower in {"validation", "valid"}:
+                        validation_rows.append(row)
+                    else:
+                        train_rows.append(row)
+            break
+        except Exception as exc:  # network/schema fallback
+            mbpp_error = f"{type(exc).__name__}: {exc}"
+    if not mbpp_loaded:
+        print("Warning: MBPP could not be loaded; continuing with CodeParrot Python data only.")
+        print("Last MBPP error:", mbpp_error)
+
+    random.Random(1337).shuffle(train_rows)
+    random.Random(7331).shuffle(validation_rows)
+    random.Random(31337).shuffle(test_rows)
+
+    atomic_write_jsonl(TRAIN_JSONL, train_rows)
+    atomic_write_jsonl(VALIDATION_JSONL, validation_rows)
+    atomic_write_jsonl(TEST_JSONL, test_rows)
+
+    print("\nDataset preparation complete")
+    print(f"  train.jsonl:      {len(train_rows):,} rows")
+    print(f"  validation.jsonl: {len(validation_rows):,} rows")
+    print(f"  test.jsonl:       {len(test_rows):,} rows")
+    for key in sorted(counters):
+        print(f"  {key:18s}: {counters[key]:,}")
+    print("Run `check-data` before tokenizer training.")
+
+
+def inspect_dataset_file(path: Path, show_samples: int = 3) -> dict[str, Any]:
+    stats: dict[str, Any] = {
+        "rows": 0,
+        "base": 0,
+        "instruct": 0,
+        "invalid": 0,
+        "duplicates": 0,
+        "chars": 0,
+        "with_tests": 0,
+    }
+    hashes: set[str] = set()
+    previews: list[dict[str, Any]] = []
+    for row in iter_jsonl(path):
+        stats["rows"] += 1
+        row_type = str(row.get("type", "base"))
+        if row_type == "instruct":
+            stats["instruct"] += 1
+            code = str(row.get("response", ""))
+            stats["with_tests"] += int(bool(row.get("tests")))
+        else:
+            stats["base"] += 1
+            code = str(row.get("text", ""))
+        stats["chars"] += len(code)
+        valid, _ = is_valid_python(code)
+        if not valid:
+            stats["invalid"] += 1
+        digest = sha256_text(serialize_training_row(row))
+        if digest in hashes:
+            stats["duplicates"] += 1
+        hashes.add(digest)
+        if len(previews) < show_samples:
+            previews.append(row)
+    stats["average_chars"] = round(stats["chars"] / max(1, stats["rows"]), 1)
+    stats["previews"] = previews
+    return stats
+
+
+def check_data(show_samples: int = 2) -> None:
+    all_good = True
+    for path in (TRAIN_JSONL, VALIDATION_JSONL, TEST_JSONL):
+        print("\n" + "=" * 88)
+        print(path.name)
+        stats = inspect_dataset_file(path, show_samples=show_samples)
+        for key, value in stats.items():
+            if key != "previews":
+                print(f"  {key:16s}: {value}")
+        if stats["invalid"]:
+            all_good = False
+        for index, row in enumerate(stats["previews"], start=1):
+            print(f"\n  Preview {index}: type={row.get('type', 'base')} source={row.get('source', '')}")
+            print("  " + serialize_training_row(row)[:700].replace("\n", "\n  "))
+    print("\nDATASET STATUS:", "PASS" if all_good else "FAIL - invalid Python found")
+
+
+# -----------------------------------------------------------------------------
+# Tokenizer and binary token corpus
+# -----------------------------------------------------------------------------
+def load_tokenizer() -> Any:
+    tokenizers = require_optional("tokenizers")
+    if not TOKENIZER_JSON.exists():
+        raise SystemExit("tokenizer.json not found. Run: python jumplander_python_100m.py tokenizer")
+    return tokenizers.Tokenizer.from_file(str(TOKENIZER_JSON))
+
+
+def train_tokenizer(vocab_size: int = 16_384) -> None:
+    tokenizers = require_optional("tokenizers")
+    from tokenizers import Tokenizer, decoders, models, normalizers, pre_tokenizers, processors, trainers
+
+    if not TRAIN_JSONL.exists():
+        raise SystemExit("train.jsonl not found")
+
+    tokenizer = Tokenizer(models.BPE(unk_token="<unk>", byte_fallback=True))
+    tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=True)
+    tokenizer.decoder = decoders.ByteLevel()
+    trainer = trainers.BpeTrainer(
+        vocab_size=vocab_size,
+        min_frequency=2,
+        show_progress=True,
+        special_tokens=SPECIAL_TOKENS,
+        initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
+    )
+
+    def corpus() -> Iterator[str]:
+        for row in iter_jsonl(TRAIN_JSONL):
+            yield serialize_training_row(row)
+
+    print(f"Training JumpLander Byte-Level BPE tokenizer (target vocab={vocab_size:,})")
+    tokenizer.train_from_iterator(corpus(), trainer=trainer, length=None)
+    tokenizer.save(str(TOKENIZER_JSON))
+    print(f"Saved: {TOKENIZER_JSON}")
+    print("Actual tokenizer vocabulary:", tokenizer.get_vocab_size())
+
+
+def build_token_file(jsonl_path: Path, output_path: Path, add_fim: bool) -> int:
+    tokenizer = load_tokenizer()
+    actual_vocab = tokenizer.get_vocab_size()
+    if actual_vocab > FULL_MODEL_CONFIG.vocab_size:
+        raise SystemExit(
+            f"Tokenizer vocab {actual_vocab} exceeds model vocab {FULL_MODEL_CONFIG.vocab_size}"
+        )
+    temp = output_path.with_suffix(output_path.suffix + ".tmp")
+    total = 0
+    with temp.open("wb") as handle:
+        for row in iter_jsonl(jsonl_path):
+            text = serialize_training_row(row)
+            ids = tokenizer.encode(text, add_special_tokens=False).ids
+            if ids:
+                array = np.asarray(ids, dtype=np.uint16)
+                array.tofile(handle)
+                total += len(ids)
+            if add_fim and row.get("type", "base") == "base":
+                code = str(row.get("text", ""))
+                variant = make_fim_variant(code)
+                if variant:
+                    fim_ids = tokenizer.encode(variant, add_special_tokens=False).ids
+                    np.asarray(fim_ids, dtype=np.uint16).tofile(handle)
+                    total += len(fim_ids)
+    temp.replace(output_path)
+    print(f"Built {output_path.name}: {total:,} tokens")
+    return total
+
+
+def build_tokens() -> None:
+    if not TOKENIZER_JSON.exists():
+        train_tokenizer(FULL_MODEL_CONFIG.vocab_size)
+    train_count = build_token_file(TRAIN_JSONL, TRAIN_BIN, add_fim=True)
+    validation_count = build_token_file(VALIDATION_JSONL, VALIDATION_BIN, add_fim=False)
+    if train_count < 10_000:
+        print("Warning: training corpus is tiny. It is suitable only for a pipeline smoke test.")
+    if validation_count < 1_000:
+        print("Warning: validation corpus is very small.")
+
+
+# -----------------------------------------------------------------------------
+# Model architecture: random initialization, no pretrained checkpoint.
+# -----------------------------------------------------------------------------
+class RMSNorm(nn.Module):
+    def __init__(self, size: int, eps: float) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(size))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        input_dtype = x.dtype
+        x_float = x.float()
+        variance = x_float.pow(2).mean(dim=-1, keepdim=True)
+        normalized = x_float * torch.rsqrt(variance + self.eps)
+        return (normalized.to(input_dtype) * self.weight)
+
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    x1 = x[..., ::2]
+    x2 = x[..., 1::2]
+    return torch.stack((-x2, x1), dim=-1).flatten(-2)
+
+
+def rope_cos_sin(
+    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, dtype=torch.float32) / head_dim)
+    )
+    positions = torch.arange(seq_len, device=device, dtype=torch.float32)
+    frequencies = torch.outer(positions, inv_freq)
+    emb = torch.repeat_interleave(frequencies, 2, dim=-1)
+    cos = emb.cos().to(dtype=dtype)[None, None, :, :]
+    sin = emb.sin().to(dtype=dtype)[None, None, :, :]
+    return cos, sin
+
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, config: ModelConfig) -> None:
+        super().__init__()
+        self.n_heads = config.n_heads
+        self.head_dim = config.d_model // config.n_heads
+        self.rope_theta = config.rope_theta
+        self.dropout = config.dropout
+        self.qkv = nn.Linear(config.d_model, 3 * config.d_model, bias=False)
+        self.out = nn.Linear(config.d_model, config.d_model, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        batch, seq_len, width = x.shape
+        qkv = self.qkv(x)
+        q, k, v = qkv.chunk(3, dim=-1)
+        q = q.view(batch, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
+        k = k.view(batch, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
+        v = v.view(batch, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
+        cos, sin = rope_cos_sin(seq_len, self.head_dim, self.rope_theta, x.device, q.dtype)
+        q = (q * cos) + (rotate_half(q) * sin)
+        k = (k * cos) + (rotate_half(k) * sin)
+        attended = F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            attn_mask=None,
+            dropout_p=self.dropout if self.training else 0.0,
+            is_causal=True,
+        )
+        attended = attended.transpose(1, 2).contiguous().view(batch, seq_len, width)
+        return self.out(attended)
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, config: ModelConfig) -> None:
+        super().__init__()
+        self.gate = nn.Linear(config.d_model, config.d_ff, bias=False)
+        self.up = nn.Linear(config.d_model, config.d_ff, bias=False)
+        self.down = nn.Linear(config.d_ff, config.d_model, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.down(F.silu(self.gate(x)) * self.up(x))
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, config: ModelConfig) -> None:
+        super().__init__()
+        self.attn_norm = RMSNorm(config.d_model, config.norm_eps)
+        self.attn = CausalSelfAttention(config)
+        self.ffn_norm = RMSNorm(config.d_model, config.norm_eps)
+        self.ffn = SwiGLU(config)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x + self.attn(self.attn_norm(x))
+        x = x + self.ffn(self.ffn_norm(x))
+        return x
+
+
+class JumpLanderPythonModel(nn.Module):
+    def __init__(self, config: ModelConfig) -> None:
+        super().__init__()
+        config.validate()
+        self.config = config
+        self.token_embedding = nn.Embedding(config.vocab_size, config.d_model)
+        self.blocks = nn.ModuleList([TransformerBlock(config) for _ in range(config.n_layers)])
+        self.final_norm = RMSNorm(config.d_model, config.norm_eps)
+        if not config.tie_embeddings:
+            self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
+        else:
+            self.lm_head = None
+        self.apply(self._init_weights)
+
+    @staticmethod
+    def _init_weights(module: nn.Module) -> None:
+        if isinstance(module, (nn.Linear, nn.Embedding)):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        targets: Optional[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        if input_ids.ndim != 2:
+            raise ValueError("input_ids must have shape [batch, sequence]")
+        if input_ids.size(1) > self.config.max_seq_len:
+            raise ValueError("Sequence exceeds configured context length")
+        x = self.token_embedding(input_ids)
+        for block in self.blocks:
+            if self.config.gradient_checkpointing and self.training:
+                x = activation_checkpoint(block, x, use_reentrant=False)
+            else:
+                x = block(x)
+        x = self.final_norm(x)
+        if self.lm_head is None:
+            logits = F.linear(x, self.token_embedding.weight)
+        else:
+            logits = self.lm_head(x)
+        loss: Optional[torch.Tensor] = None
+        if targets is not None:
+            loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1))
+        return logits, loss
+
+
+def parameter_report(config: ModelConfig) -> dict[str, int]:
+    model = JumpLanderPythonModel(config)
+    total = sum(parameter.numel() for parameter in model.parameters())
+    trainable = sum(parameter.numel() for parameter in model.parameters() if parameter.requires_grad)
+    del model
+    return {"total": total, "trainable": trainable}
+
+
+# -----------------------------------------------------------------------------
+# Training
+# -----------------------------------------------------------------------------
+class TokenBlockDataset(Dataset[tuple[torch.Tensor, torch.Tensor]]):
+    def __init__(self, path: Path, block_size: int) -> None:
+        if not path.exists():
+            raise FileNotFoundError(f"Token file not found: {path}. Run build-tokens first.")
+        self.tokens = np.memmap(path, dtype=np.uint16, mode="r")
+        self.block_size = block_size
+        if len(self.tokens) < block_size + 1:
+            raise ValueError(
+                f"{path.name} contains only {len(self.tokens):,} tokens; need at least {block_size + 1:,}"
+            )
+        self.examples = (len(self.tokens) - 1) // block_size
+
+    def __len__(self) -> int:
+        return self.examples
+
+    def __getitem__(self, index: int) -> tuple[torch.Tensor, torch.Tensor]:
+        start = index * self.block_size
+        chunk = np.asarray(self.tokens[start : start + self.block_size + 1], dtype=np.int64)
+        x = torch.from_numpy(chunk[:-1].copy())
+        y = torch.from_numpy(chunk[1:].copy())
+        return x, y
+
+
+def cosine_lr(step: int, config: TrainConfig) -> float:
+    if step < config.warmup_steps:
+        return config.learning_rate * (step + 1) / max(1, config.warmup_steps)
+    progress = (step - config.warmup_steps) / max(1, config.total_steps - config.warmup_steps)
+    progress = min(max(progress, 0.0), 1.0)
+    coefficient = 0.5 * (1.0 + math.cos(math.pi * progress))
+    return config.min_learning_rate + coefficient * (
+        config.learning_rate - config.min_learning_rate
+    )
+
+
+def choose_precision(device: torch.device) -> tuple[torch.dtype, bool]:
+    if device.type != "cuda":
+        return torch.float32, False
+    if torch.cuda.is_bf16_supported():
+        return torch.bfloat16, True
+    return torch.float16, True
+
+
+def make_optimizer(model: nn.Module, train_config: TrainConfig, device: torch.device) -> torch.optim.Optimizer:
+    kwargs: dict[str, Any] = {
+        "lr": train_config.learning_rate,
+        "betas": (0.9, 0.95),
+        "eps": 1e-8,
+        "weight_decay": train_config.weight_decay,
+    }
+    if device.type == "cuda":
+        try:
+            return torch.optim.AdamW(model.parameters(), fused=True, **kwargs)
+        except (TypeError, RuntimeError):
+            pass
+    return torch.optim.AdamW(model.parameters(), **kwargs)
+
+
+@torch.no_grad()
+def evaluate(
+    model: JumpLanderPythonModel,
+    loader: DataLoader[Any],
+    device: torch.device,
+    dtype: torch.dtype,
+    batches: int,
+) -> float:
+    model.eval()
+    losses: list[float] = []
+    iterator = iter(loader)
+    amp_enabled = device.type == "cuda"
+    for _ in range(batches):
+        try:
+            x, y = next(iterator)
+        except StopIteration:
+            break
+        x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)
+        with torch.autocast(device_type=device.type, dtype=dtype, enabled=amp_enabled):
+            _, loss = model(x, y)
+        if loss is not None:
+            losses.append(float(loss.item()))
+    model.train()
+    return float(sum(losses) / max(1, len(losses)))
+
+
+def save_checkpoint(
+    path: Path,
+    model: JumpLanderPythonModel,
+    optimizer: torch.optim.Optimizer,
+    step: int,
+    tokens_seen: int,
+    train_config: TrainConfig,
+) -> None:
+    payload = {
+        "model_name": model.config.name,
+        "model_config": dataclasses.asdict(model.config),
+        "train_config": dataclasses.asdict(train_config),
+        "model_state": model.state_dict(),
+        "optimizer_state": optimizer.state_dict(),
+        "step": step,
+        "tokens_seen": tokens_seen,
+        "saved_at": time.time(),
+        "format_version": 1,
+    }
+    temp = path.with_suffix(path.suffix + ".tmp")
+    torch.save(payload, temp)
+    temp.replace(path)
+
+
+def load_checkpoint_payload(path: Path, device: torch.device) -> dict[str, Any]:
+    if not path.exists():
+        raise FileNotFoundError(f"Checkpoint not found: {path}")
+    try:
+        return torch.load(path, map_location=device, weights_only=False)
+    except TypeError:
+        return torch.load(path, map_location=device)
+
+
+def train_model(
+    smoke: bool,
+    resume: bool,
+    total_steps: Optional[int],
+    batch_size: Optional[int],
+    accumulation: Optional[int],
+) -> None:
+    if not TRAIN_BIN.exists() or not VALIDATION_BIN.exists():
+        print("Token files are missing; building them now.")
+        build_tokens()
+
+    model_config = dataclasses.replace(SMOKE_MODEL_CONFIG if smoke else FULL_MODEL_CONFIG)
+    if smoke:
+        train_config = TrainConfig(
+            batch_size=batch_size or 2,
+            gradient_accumulation=accumulation or 4,
+            learning_rate=5e-4,
+            min_learning_rate=5e-5,
+            warmup_steps=20,
+            total_steps=total_steps or 300,
+            eval_interval=50,
+            eval_batches=10,
+            save_interval=100,
+            log_interval=5,
+        )
+        checkpoint_path = SMOKE_CHECKPOINT
+    else:
+        train_config = TrainConfig(
+            batch_size=batch_size or 1,
+            gradient_accumulation=accumulation or 32,
+            total_steps=total_steps or 10_000,
+        )
+        checkpoint_path = MODEL_CHECKPOINT
+
+    seed_everything(train_config.seed)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    dtype, amp_enabled = choose_precision(device)
+    print(f"Device: {device}; precision: {dtype}; AMP: {amp_enabled}")
+    if device.type != "cuda":
+        print("Warning: CUDA is not available. Full 100M training on CPU is impractical.")
+
+    train_dataset = TokenBlockDataset(TRAIN_BIN, model_config.max_seq_len)
+    validation_dataset = TokenBlockDataset(VALIDATION_BIN, model_config.max_seq_len)
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=train_config.batch_size,
+        shuffle=True,
+        num_workers=train_config.num_workers,
+        pin_memory=device.type == "cuda",
+        drop_last=True,
+    )
+    validation_loader = DataLoader(
+        validation_dataset,
+        batch_size=train_config.batch_size,
+        shuffle=False,
+        num_workers=0,
+        pin_memory=device.type == "cuda",
+        drop_last=False,
+    )
+
+    model = JumpLanderPythonModel(model_config).to(device)
+    report = parameter_report(model_config)
+    print(f"Model parameters: {report['total']:,}")
+    optimizer = make_optimizer(model, train_config, device)
+    scaler = torch.amp.GradScaler("cuda", enabled=(device.type == "cuda" and dtype == torch.float16))
+
+    start_step = 0
+    tokens_seen = 0
+    if resume:
+        payload = load_checkpoint_payload(checkpoint_path, device)
+        model.load_state_dict(payload["model_state"])
+        optimizer.load_state_dict(payload["optimizer_state"])
+        start_step = int(payload.get("step", 0))
+        tokens_seen = int(payload.get("tokens_seen", 0))
+        print(f"Resumed {checkpoint_path.name} at step {start_step:,}")
+
+    model.train()
+    iterator = iter(train_loader)
+    optimizer.zero_grad(set_to_none=True)
+    last_log_time = time.time()
+    running_loss = 0.0
+    micro_steps = 0
+
+    for step in range(start_step, train_config.total_steps):
+        lr = cosine_lr(step, train_config)
+        for group in optimizer.param_groups:
+            group["lr"] = lr
+
+        for _ in range(train_config.gradient_accumulation):
+            try:
+                x, y = next(iterator)
+            except StopIteration:
+                iterator = iter(train_loader)
+                x, y = next(iterator)
+            x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)
+            with torch.autocast(device_type=device.type, dtype=dtype, enabled=amp_enabled):
+                _, loss = model(x, y)
+                if loss is None:
+                    raise RuntimeError("Training loss was not produced")
+                scaled_loss = loss / train_config.gradient_accumulation
+            scaler.scale(scaled_loss).backward()
+            running_loss += float(loss.detach().item())
+            micro_steps += 1
+            tokens_seen += x.numel()
+
+        scaler.unscale_(optimizer)
+        gradient_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), train_config.grad_clip)
+        scaler.step(optimizer)
+        scaler.update()
+        optimizer.zero_grad(set_to_none=True)
+
+        completed_step = step + 1
+        if completed_step % train_config.log_interval == 0:
+            now = time.time()
+            elapsed = max(now - last_log_time, 1e-6)
+            recent_tokens = (
+                train_config.log_interval
+                * train_config.gradient_accumulation
+                * train_config.batch_size
+                * model_config.max_seq_len
+            )
+            tokens_per_second = recent_tokens / elapsed
+            average_loss = running_loss / max(1, micro_steps)
+            memory = ""
+            if device.type == "cuda":
+                allocated = torch.cuda.max_memory_allocated() / (1024**3)
+                memory = f" vram={allocated:.2f}GB"
+                torch.cuda.reset_peak_memory_stats()
+            print(
+                f"step={completed_step:,}/{train_config.total_steps:,} "
+                f"loss={average_loss:.4f} lr={lr:.2e} grad={float(gradient_norm):.3f} "
+                f"tok/s={tokens_per_second:,.0f}{memory}"
+            )
+            running_loss = 0.0
+            micro_steps = 0
+            last_log_time = now
+
+        if completed_step % train_config.eval_interval == 0:
+            validation_loss = evaluate(
+                model, validation_loader, device, dtype, train_config.eval_batches
+            )
+            print(f"validation_loss={validation_loss:.4f}")
+
+        if completed_step % train_config.save_interval == 0:
+            save_checkpoint(
+                checkpoint_path,
+                model,
+                optimizer,
+                completed_step,
+                tokens_seen,
+                train_config,
+            )
+            print(f"Saved checkpoint: {checkpoint_path.name}")
+
+    save_checkpoint(
+        checkpoint_path,
+        model,
+        optimizer,
+        train_config.total_steps,
+        tokens_seen,
+        train_config,
+    )
+    print(f"Training complete. Saved: {checkpoint_path}")
+
+
+# -----------------------------------------------------------------------------
+# Inference and web UI
+# -----------------------------------------------------------------------------
+def load_model_for_inference(checkpoint_path: Path) -> tuple[JumpLanderPythonModel, Any, torch.device]:
+    tokenizer = load_tokenizer()
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    payload = load_checkpoint_payload(checkpoint_path, device)
+    config = ModelConfig(**payload["model_config"])
+    config.gradient_checkpointing = False
+    model = JumpLanderPythonModel(config)
+    model.load_state_dict(payload["model_state"])
+    model.to(device)
+    model.eval()
+    return model, tokenizer, device
+
+
+def top_p_sample(logits: torch.Tensor, temperature: float, top_p: float) -> torch.Tensor:
+    if temperature <= 0:
+        return torch.argmax(logits, dim=-1, keepdim=True)
+    logits = logits / max(temperature, 1e-5)
+    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+    probabilities = F.softmax(sorted_logits, dim=-1)
+    cumulative = torch.cumsum(probabilities, dim=-1)
+    remove = cumulative > top_p
+    remove[..., 1:] = remove[..., :-1].clone()
+    remove[..., 0] = False
+    sorted_logits = sorted_logits.masked_fill(remove, float("-inf"))
+    probabilities = F.softmax(sorted_logits, dim=-1)
+    selected = torch.multinomial(probabilities, num_samples=1)
+    return sorted_indices.gather(-1, selected)
+
+
+@torch.inference_mode()
+def generate_text(
+    model: JumpLanderPythonModel,
+    tokenizer: Any,
+    device: torch.device,
+    prompt: str,
+    mode: str,
+    max_new_tokens: int,
+    temperature: float,
+    top_p: float,
+) -> str:
+    if mode == "instruction":
+        formatted = f"<bos><instruction>\n{prompt.strip()}\n<response>\n"
+    else:
+        formatted = f"<bos><file_start>\n{prompt}"
+    encoded = tokenizer.encode(formatted, add_special_tokens=False).ids
+    if not encoded:
+        raise ValueError("Prompt produced no tokens")
+    max_context = model.config.max_seq_len
+    encoded = encoded[-max_context:]
+    ids = torch.tensor([encoded], dtype=torch.long, device=device)
+    eos_id = tokenizer.token_to_id("<eos>")
+    file_end_id = tokenizer.token_to_id("<file_end>")
+    actual_vocab = tokenizer.get_vocab_size()
+
+    for _ in range(max_new_tokens):
+        context = ids[:, -max_context:]
+        logits, _ = model(context)
+        next_logits = logits[:, -1, :]
+        if actual_vocab < next_logits.size(-1):
+            next_logits[:, actual_vocab:] = float("-inf")
+        next_id = top_p_sample(next_logits, temperature, top_p)
+        ids = torch.cat([ids, next_id], dim=1)
+        token_value = int(next_id.item())
+        if token_value in {value for value in (eos_id, file_end_id) if value is not None}:
+            break
+
+    generated_ids = ids[0, len(encoded) :].tolist()
+    output = tokenizer.decode(generated_ids, skip_special_tokens=True)
+    return output.strip()
+
+
+def resolve_checkpoint(use_smoke: bool) -> Path:
+    selected = SMOKE_CHECKPOINT if use_smoke else MODEL_CHECKPOINT
+    if not selected.exists() and not use_smoke and SMOKE_CHECKPOINT.exists():
+        print("Full model checkpoint not found; using smoke checkpoint.")
+        return SMOKE_CHECKPOINT
+    return selected
+
+
+def terminal_generate(args: argparse.Namespace) -> None:
+    checkpoint_path = resolve_checkpoint(args.smoke)
+    model, tokenizer, device = load_model_for_inference(checkpoint_path)
+    result = generate_text(
+        model,
+        tokenizer,
+        device,
+        args.prompt,
+        args.mode,
+        args.max_new_tokens,
+        args.temperature,
+        args.top_p,
+    )
+    print(result)
+
+
+WEB_PAGE = """<!doctype html>
+<html lang="en">
+<head>
+<meta charset="utf-8">
+<meta name="viewport" content="width=device-width,initial-scale=1">
+<title>JumpLander Python Decoder 100M</title>
+<style>
+:root{color-scheme:dark;--bg:#0c0c0e;--panel:#151719;--line:#28392b;--accent:#819e2e;--text:#f9f9f9;--muted:#a7aaa4}
+*{box-sizing:border-box}body{margin:0;background:var(--bg);color:var(--text);font:15px/1.55 Inter,Segoe UI,Arial,sans-serif}
+main{max-width:980px;margin:34px auto;padding:0 18px}.brand{display:flex;align-items:center;gap:12px;margin-bottom:20px}
+.logo{width:42px;height:42px;border-radius:14px;background:linear-gradient(145deg,#819e2e,#4b5d2a);display:grid;place-items:center;font-weight:800;color:#0c0c0e}
+h1{font-size:22px;margin:0}.sub{color:var(--muted);font-size:13px}.card{background:var(--panel);border:1px solid var(--line);border-radius:22px;padding:18px}
+.controls{display:grid;grid-template-columns:1fr 140px 120px 120px;gap:10px;margin-bottom:12px}select,input,button,textarea{border:1px solid #303430;background:#101210;color:var(--text);border-radius:12px;padding:11px;font:inherit}
+textarea{width:100%;min-height:190px;resize:vertical;font-family:Consolas,monospace}button{background:var(--accent);color:#0c0c0e;border:0;font-weight:700;cursor:pointer}button:disabled{opacity:.55;cursor:wait}
+pre{white-space:pre-wrap;min-height:220px;background:#0d0f0d;border:1px solid #252a25;border-radius:14px;padding:16px;overflow:auto;font-family:Consolas,monospace}
+.status{color:var(--muted);font-size:13px;margin:10px 2px}.note{margin-top:14px;color:var(--muted);font-size:13px}@media(max-width:760px){.controls{grid-template-columns:1fr 1fr}.controls button{grid-column:1/-1}}
+</style>
+</head>
+<body><main>
+<div class="brand"><div class="logo">JL</div><div><h1>JumpLander Python Decoder 100M</h1><div class="sub">Local from-scratch Python model test console</div></div></div>
+<div class="card">
+<div class="controls">
+<select id="mode"><option value="instruction">English instruction → Python</option><option value="completion">Python code completion</option></select>
+<input id="tokens" type="number" min="1" max="512" value="160" title="Max new tokens">
+<input id="temp" type="number" min="0" max="2" step="0.05" value="0.20" title="Temperature">
+<button id="run">Generate</button>
+</div>
+<textarea id="prompt" spellcheck="false">Write a Python function that removes duplicate items while preserving their original order.</textarea>
+<div id="status" class="status">Ready</div>
+<pre id="output"></pre>
+<div class="note">This is a small research model. The smoke checkpoint validates the pipeline; useful quality requires a much larger clean corpus and longer training.</div>
+</div></main>
+<script>
+const run=document.getElementById('run'),status=document.getElementById('status'),output=document.getElementById('output');
+run.onclick=async()=>{run.disabled=true;status.textContent='Generating…';output.textContent='';const started=performance.now();
+try{const response=await fetch('/generate',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify({prompt:document.getElementById('prompt').value,mode:document.getElementById('mode').value,max_new_tokens:Number(document.getElementById('tokens').value),temperature:Number(document.getElementById('temp').value),top_p:.95})});const data=await response.json();if(!response.ok)throw new Error(data.error||'Generation failed');output.textContent=data.output;status.textContent=`Done in ${((performance.now()-started)/1000).toFixed(2)}s`}
+catch(error){status.textContent='Error';output.textContent=String(error)}finally{run.disabled=false}};
+</script></body></html>"""
+
+
+def run_web_chat(args: argparse.Namespace) -> None:
+    from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+
+    checkpoint_path = resolve_checkpoint(args.smoke)
+    model, tokenizer, device = load_model_for_inference(checkpoint_path)
+    generation_lock = threading.Lock()
+
+    class Handler(BaseHTTPRequestHandler):
+        def log_message(self, format_string: str, *values: Any) -> None:
+            print("[web]", format_string % values)
+
+        def send_bytes(self, status: int, content_type: str, body: bytes) -> None:
+            self.send_response(status)
+            self.send_header("Content-Type", content_type)
+            self.send_header("Content-Length", str(len(body)))
+            self.send_header("Cache-Control", "no-store")
+            self.end_headers()
+            self.wfile.write(body)
+
+        def do_GET(self) -> None:  # noqa: N802
+            if self.path == "/":
+                self.send_bytes(200, "text/html; charset=utf-8", WEB_PAGE.encode("utf-8"))
+            else:
+                self.send_bytes(404, "text/plain; charset=utf-8", b"Not found")
+
+        def do_POST(self) -> None:  # noqa: N802
+            if self.path != "/generate":
+                self.send_bytes(404, "application/json", b'{"error":"Not found"}')
+                return
+            try:
+                length = int(self.headers.get("Content-Length", "0"))
+                if length <= 0 or length > 100_000:
+                    raise ValueError("Invalid request body")
+                payload = json.loads(self.rfile.read(length))
+                prompt = str(payload.get("prompt", ""))
+                if not prompt.strip() or len(prompt) > 20_000:
+                    raise ValueError("Prompt is empty or too long")
+                mode = str(payload.get("mode", "instruction"))
+                if mode not in {"instruction", "completion"}:
+                    raise ValueError("Invalid mode")
+                max_new_tokens = min(max(int(payload.get("max_new_tokens", 160)), 1), 512)
+                temperature = min(max(float(payload.get("temperature", 0.2)), 0.0), 2.0)
+                top_p = min(max(float(payload.get("top_p", 0.95)), 0.05), 1.0)
+                with generation_lock:
+                    result = generate_text(
+                        model,
+                        tokenizer,
+                        device,
+                        prompt,
+                        mode,
+                        max_new_tokens,
+                        temperature,
+                        top_p,
+                    )
+                body = json.dumps({"output": result}, ensure_ascii=False).encode("utf-8")
+                self.send_bytes(200, "application/json; charset=utf-8", body)
+            except Exception as exc:
+                body = json.dumps({"error": str(exc)}, ensure_ascii=False).encode("utf-8")
+                self.send_bytes(400, "application/json; charset=utf-8", body)
+
+    address = (args.host, args.port)
+    server = ThreadingHTTPServer(address, Handler)
+    url = f"http://{args.host}:{args.port}"
+    print(f"Loaded checkpoint: {checkpoint_path.name}")
+    print(f"Web UI: {url}")
+    if not args.no_browser:
+        threading.Timer(0.8, lambda: webbrowser.open(url)).start()
+    try:
+        server.serve_forever()
+    except KeyboardInterrupt:
+        print("\nStopping web server")
+    finally:
+        server.server_close()
+
+
+# -----------------------------------------------------------------------------
+# Diagnostics and CLI
+# -----------------------------------------------------------------------------
+def system_check() -> None:
+    print(MODEL_NAME)
+    print("Python:", sys.version.replace("\n", " "))
+    print("PyTorch:", torch.__version__)
+    print("CUDA available:", torch.cuda.is_available())
+    if torch.cuda.is_available():
+        print("GPU:", torch.cuda.get_device_name(0))
+        print("VRAM GB:", round(torch.cuda.get_device_properties(0).total_memory / 1024**3, 2))
+        print("BF16 supported:", torch.cuda.is_bf16_supported())
+    for package in ("datasets", "tokenizers"):
+        try:
+            module = __import__(package)
+            print(f"{package}:", getattr(module, "__version__", "installed"))
+        except ImportError:
+            print(f"{package}: MISSING")
+    for path in (TRAIN_JSONL, VALIDATION_JSONL, TEST_JSONL, TOKENIZER_JSON):
+        print(f"{path.name}:", "present" if path.exists() else "missing")
+
+
+def model_info() -> None:
+    for config in (FULL_MODEL_CONFIG, SMOKE_MODEL_CONFIG):
+        report = parameter_report(config)
+        print("\n" + config.name)
+        print(json.dumps(dataclasses.asdict(config), indent=2))
+        print(f"parameters: {report['total']:,}")
+        print(f"estimated FP16 weights: {report['total'] * 2 / 1024**2:.1f} MiB")
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        description="JumpLander Python Decoder 100M — one-file from-scratch model project"
+    )
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    subparsers.add_parser("check", help="Check Python, PyTorch, CUDA, and local files")
+    subparsers.add_parser("info", help="Show architecture and exact parameter count")
+
+    preview = subparsers.add_parser("dataset-preview", help="Stream and display Python-Edu rows")
+    preview.add_argument("--limit", type=int, default=5)
+
+    prepare_test = subparsers.add_parser("prepare-test", help="Build a small remote test dataset")
+    prepare_test.add_argument("--max-base", type=int, default=2_000)
+    prepare_test.add_argument("--min-score", type=int, default=3)
+
+    prepare = subparsers.add_parser("prepare", help="Build/replace the local dataset from Hugging Face")
+    prepare.add_argument("--max-base", type=int, default=50_000)
+    prepare.add_argument("--min-score", type=int, default=3)
+
+    check_data_parser = subparsers.add_parser("check-data", help="Validate and preview local JSONL files")
+    check_data_parser.add_argument("--show", type=int, default=2)
+
+    tokenizer_parser = subparsers.add_parser("tokenizer", help="Train tokenizer.json from train.jsonl")
+    tokenizer_parser.add_argument("--vocab-size", type=int, default=16_384)
+
+    subparsers.add_parser("build-tokens", help="Create flat uint16 train/validation token files")
+
+    smoke = subparsers.add_parser("smoke", help="Train the small pipeline-validation model")
+    smoke.add_argument("--steps", type=int, default=300)
+    smoke.add_argument("--batch-size", type=int, default=None)
+    smoke.add_argument("--accumulation", type=int, default=None)
+    smoke.add_argument("--resume", action="store_true")
+
+    train = subparsers.add_parser("train", help="Train the ~97.5M parameter model")
+    train.add_argument("--steps", type=int, default=10_000)
+    train.add_argument("--batch-size", type=int, default=None)
+    train.add_argument("--accumulation", type=int, default=None)
+    train.add_argument("--resume", action="store_true")
+
+    generate = subparsers.add_parser("generate", help="Generate code in the terminal")
+    generate.add_argument("prompt")
+    generate.add_argument("--mode", choices=["instruction", "completion"], default="instruction")
+    generate.add_argument("--max-new-tokens", type=int, default=160)
+    generate.add_argument("--temperature", type=float, default=0.2)
+    generate.add_argument("--top-p", type=float, default=0.95)
+    generate.add_argument("--smoke", action="store_true")
+
+    chat = subparsers.add_parser("chat", help="Launch the local browser UI")
+    chat.add_argument("--host", default="127.0.0.1")
+    chat.add_argument("--port", type=int, default=7860)
+    chat.add_argument("--smoke", action="store_true")
+    chat.add_argument("--no-browser", action="store_true")
+
+    return parser
+
+
+def main() -> None:
+    parser = build_parser()
+    args = parser.parse_args()
+    try:
+        if args.command == "check":
+            system_check()
+        elif args.command == "info":
+            model_info()
+        elif args.command == "dataset-preview":
+            dataset_preview(args.limit)
+        elif args.command == "prepare-test":
+            prepare_remote_dataset(args.max_base, test_mode=True, min_score=args.min_score)
+        elif args.command == "prepare":
+            prepare_remote_dataset(args.max_base, test_mode=False, min_score=args.min_score)
+        elif args.command == "check-data":
+            check_data(args.show)
+        elif args.command == "tokenizer":
+            train_tokenizer(args.vocab_size)
+        elif args.command == "build-tokens":
+            build_tokens()
+        elif args.command == "smoke":
+            train_model(True, args.resume, args.steps, args.batch_size, args.accumulation)
+        elif args.command == "train":
+            train_model(False, args.resume, args.steps, args.batch_size, args.accumulation)
+        elif args.command == "generate":
+            terminal_generate(args)
+        elif args.command == "chat":
+            run_web_chat(args)
+        else:
+            parser.error("Unknown command")
+    except KeyboardInterrupt:
+        print("\nCancelled")
+    except Exception as exc:
+        print(f"ERROR: {type(exc).__name__}: {exc}", file=sys.stderr)
+        if os.environ.get("JL_DEBUG") == "1":
+            traceback.print_exc()
+        raise SystemExit(1) from exc
+
+
+if __name__ == "__main__":
+    main()