jumplander_python_100m.py

"""
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()