{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Pepaman \u00b7 Universe Lab \u2014 train from scratch\n",
    "Randomly initialized causal transformer with a fixed 256-byte vocabulary. No pretrained weights, model download, external training API, or paid service. This is next-byte language modeling, not a finished chat assistant or truth detector.\n",
    "\n",
    "The built-in corpus is a tiny demonstration. Expect memorization and unreadable output. To build a useful science assistant, collect a much larger, diverse, reviewed corpus; keep source attribution, rights, dates and uncertainty labels. Preserve separate source/topic groups for validation and test. Training does not guarantee truth.\n",
    "\n",
    "**Kaggle:** import this `.ipynb`, select an available GPU and enable internet only if PyTorch installation is needed. **Colab:** File \u2192 Upload notebook, then Runtime \u2192 Change runtime type \u2192 GPU. Free availability and quotas vary. Use your own legitimate account. Pepaman hosts the notebook; computation happens in your chosen runtime. Neither service is connected just by opening this file."
   ],
   "id": "pepaman-01"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "import importlib.util, subprocess, sys\n",
    "if importlib.util.find_spec('torch') is None:\n",
    "    subprocess.check_call([sys.executable, '-m', 'pip', 'install', 'torch>=2.5,<3'])\n",
    "import torch\n",
    "print('PyTorch:', torch.__version__)\n",
    "print('GPU:', torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'none \u2014 GPU training is blocked by default')"
   ],
   "execution_count": null,
   "outputs": [],
   "id": "pepaman-02"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Choose storage and data\n",
    "Set `DATA_DIR` to a directory containing UTF-8 `train.txt`, `validation.txt`, and `test.txt`. For the first demonstration, leave it as `None` to write the built-in summaries into a new directory. They never overwrite your own files.\n",
    "\n",
    "On Colab you may mount Drive yourself and set `OUTPUT_DIR` to a Drive folder. On Kaggle, `/kaggle/working` is only retained when outputs are saved using the platform. Download checkpoints regularly; writing a checkpoint on an ephemeral disk does not preserve it after a reset. To resume elsewhere, restore `latest.pt`, use identical data and architecture, set `RESUME=True`, and increase `STEPS` to the total target step."
   ],
   "id": "pepaman-03"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "from pathlib import Path\n",
    "DATA_DIR = None  # Example: '/kaggle/input/my-reviewed-corpus'\n",
    "OUTPUT_DIR = Path('/kaggle/working/pepaman-run' if Path('/kaggle').exists() else '/content/pepaman-run')\n",
    "RESUME = False\n",
    "ALLOW_CPU = False  # True only for a deliberately small local experiment\n",
    "STEPS = 500         # Total target step, including any resumed steps\n",
    "BATCH_SIZE = 8\n",
    "ACCUMULATION = 2"
   ],
   "execution_count": null,
   "outputs": [],
   "id": "pepaman-04"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "import json\n",
    "DEMO = json.loads(\"{\\\"train\\\": \\\"COSMIC HISTORY\\\\nThe universe developed from a hot, dense early state and has expanded and cooled. The Big Bang model describes its early evolution; it does not establish why anything exists. Observations constrain scientific models, and explanations should be revised when stronger evidence requires it.\\\\n\\\\nEarly cosmic conditions allowed protons and neutrons to form light atomic nuclei. Much later, cooling allowed electrons to become bound in atoms. Light from the early universe remains observable as the cosmic microwave background. It provides evidence about early conditions rather than a photograph of an absolute beginning.\\\\n\\\\nGravity gathered unevenly distributed gas into denser regions. Compression heated some regions until fusion began, forming stars. Stars later collected into galaxies. Scientists reconstruct this history from measurements and models with stated assumptions. The earliest stages and ultimate future still contain unanswered questions.\\\\n\\\\nSource: https://science.nasa.gov/universe/overview/\\\\n\\\\nHOW THE SUN WORKS\\\\nThe Sun is a star. Its core releases energy through nuclear fusion, joining hydrogen nuclei into helium. Some mass becomes energy. This differs from ordinary chemical burning and does not require atmospheric oxygen.\\\\n\\\\nThe Sun contains electrically charged matter called plasma. Its gravity pulls inward while pressure helps support it. Energy travels outward and eventually reaches space as radiation. Visible light is one part of that radiation.\\\\n\\\\nSolar magnetic activity can produce sunspots, flares and eruptions. A scientific explanation should distinguish the mechanism producing energy in the core from activity in the outer atmosphere. A simplified explanation is useful when its limits are stated.\\\\n\\\\nSource: https://science.nasa.gov/sun/facts/\\\\n\\\\nBLACK HOLES\\\\nA black hole has a region from which light cannot escape to a distant observer. The event horizon marks that boundary. It is not a material wall. Hot matter outside the horizon can emit radiation that telescopes detect.\\\\n\\\\nSome black holes form when massive stars collapse. Supermassive black holes are found in the centers of many galaxies. Their formation and growth remain research subjects.\\\\n\\\\nBlack holes affect surrounding matter through gravity. They do not automatically consume everything in the universe. Observations of nearby stars and gas help scientists infer their properties. Claims about what happens deep inside require caution because direct observations are unavailable.\\\\n\\\\nSource: https://science.nasa.gov/universe/black-holes/\\\\n\\\\nANSWERING CAREFULLY\\\\nAn answer should explain a mechanism, describe the evidence, and name relevant uncertainty. A confident sentence alone is not evidence. A model is a useful description that can have limits. An unanswered scientific question should remain unanswered rather than be filled with an invented fact.\\\\n\\\\nPhilosophical and religious traditions offer interpretations of meaning and existence. Describe a tradition respectfully and attribute its claims to that tradition. Do not present a personal belief as an experimentally established universal fact. Ask what the speaker means when a question has several interpretations.\\\\n\\\\nThese answer-style instructions are authored for this project, not quoted from a scientific authority.\\\\n\\\", \\\"validation\\\": \\\"DARK MATTER: A HELD-OUT TOPIC\\\\nDark matter is inferred from gravitational effects. Visible matter alone does not explain several observations of galaxies and clusters under the usual models. Astronomers also study gravitational lensing, where the paths of light reveal the influence of mass.\\\\n\\\\nThe name does not identify a confirmed particle. Evidence for an unseen gravitational component and certainty about its microscopic nature are different claims. The first has observational support; the second remains unresolved.\\\\n\\\\nAn explanation should identify what was observed and what was inferred. A proposed dark-matter candidate requires tests and evidence. Calling a candidate established would overstate current knowledge. There is no justification for claiming a universal answer just because a model can generate fluent text.\\\\n\\\\nSource for the physical claims: https://science.nasa.gov/missions/hubble/dark-matter/\\\\nThe distinction between answer confidence and evidence is project-authored guidance.\\\\n\\\", \\\"test\\\": \\\"EARTH CLIMATE: FINAL HELD-OUT TOPIC\\\\nGreenhouse gases absorb and emit infrared radiation, affecting how energy moves through Earth's atmosphere. The natural greenhouse effect helps keep Earth warmer than it would otherwise be. Increasing greenhouse gas concentrations changes the planet's energy balance.\\\\n\\\\nHuman activities, especially burning fossil fuels, have increased atmospheric carbon dioxide. The current warming trend is primarily driven by human greenhouse gas emissions. Measurements and physical explanations support this conclusion.\\\\n\\\\nWeather varies over short periods. Climate describes patterns over longer periods. A single cold day does not establish the absence of global warming. Specific regional changes can involve several influences, so an explanation should distinguish the broad mechanism from a local prediction.\\\\n\\\\nSource: https://science.nasa.gov/climate-change/causes/\\\\nThis held-out passage evaluates text prediction on a topic excluded from training. It is too short to establish a model's factual reliability or general scientific ability.\\\\n\\\"}\")\n",
    "SOURCES = json.loads(\"{\\\"purpose\\\": \\\"Original, brief educational summaries and project-authored answer-style guidance. Demonstration only.\\\", \\\"reviewed_on\\\": \\\"2026-10-19\\\", \\\"notices\\\": \\\"No endorsement by NASA. Linked sources support physical explanations; project guidance is separately identified. Review claims before expanding the dataset. Follow each source's reuse terms, including exceptions for third-party material.\\\", \\\"splits\\\": {\\\"train\\\": {\\\"bytes\\\": 3280, \\\"purpose\\\": \\\"Optimization only\\\"}, \\\"validation\\\": {\\\"bytes\\\": 1000, \\\"purpose\\\": \\\"Model selection\\\"}, \\\"test\\\": {\\\"bytes\\\": 1066, \\\"purpose\\\": \\\"Final evaluation only\\\"}}, \\\"sources\\\": [{\\\"split\\\": \\\"train\\\", \\\"url\\\": \\\"https://science.nasa.gov/universe/overview/\\\"}, {\\\"split\\\": \\\"train\\\", \\\"url\\\": \\\"https://science.nasa.gov/sun/facts/\\\"}, {\\\"split\\\": \\\"train\\\", \\\"url\\\": \\\"https://science.nasa.gov/universe/black-holes/\\\"}, {\\\"split\\\": \\\"validation\\\", \\\"url\\\": \\\"https://science.nasa.gov/missions/hubble/dark-matter/\\\"}, {\\\"split\\\": \\\"test\\\", \\\"url\\\": \\\"https://science.nasa.gov/climate-change/causes/\\\"}]}\")\n",
    "if DATA_DIR is None:\n",
    "    DATA_DIR = Path('/kaggle/working/pepaman-demo' if Path('/kaggle').exists() else '/content/pepaman-demo')\n",
    "    DATA_DIR.mkdir(parents=True, exist_ok=True)\n",
    "    for name, text in DEMO.items():\n",
    "        destination = DATA_DIR / (name + '.txt')\n",
    "        if not destination.exists():\n",
    "            destination.write_text(text, encoding='utf-8')\n",
    "    print('DEMONSTRATION CORPUS ONLY \u2014 replace it before serious training.')\n",
    "else:\n",
    "    DATA_DIR = Path(DATA_DIR)\n",
    "print('Data:', DATA_DIR, 'Output:', OUTPUT_DIR)"
   ],
   "execution_count": null,
   "outputs": [],
   "id": "pepaman-05"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Model and training engine\n",
    "All trainable weights below start from random numbers. UTF-8 bytes are the tokenizer: there is no pretrained tokenizer. This small model uses causal attention, gradient clipping, AdamW, source-separated data files and resumable checkpoints. Dataset checks catch identical files and copied paragraphs; also review semantic overlap yourself. Changing hardware can change numerical results."
   ],
   "id": "pepaman-06"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "\"\"\"Pepaman byte-level causal transformer. Random weights, no pretrained assets.\"\"\"\n",
    "import hashlib\n",
    "import json\n",
    "import math\n",
    "import random\n",
    "import time\n",
    "from dataclasses import asdict, dataclass\n",
    "from pathlib import Path\n",
    "\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch.nn import functional as F\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class Config:\n",
    "    width: int = 192\n",
    "    layers: int = 4\n",
    "    heads: int = 6\n",
    "    context: int = 128\n",
    "    dropout: float = 0.1\n",
    "\n",
    "\n",
    "class Block(nn.Module):\n",
    "    def __init__(self, c):\n",
    "        super().__init__()\n",
    "        self.heads = c.heads\n",
    "        self.norm1 = nn.LayerNorm(c.width)\n",
    "        self.norm2 = nn.LayerNorm(c.width)\n",
    "        self.qkv = nn.Linear(c.width, 3 * c.width)\n",
    "        self.proj = nn.Linear(c.width, c.width)\n",
    "        self.ff = nn.Sequential(nn.Linear(c.width, 4*c.width), nn.GELU(),\n",
    "                                nn.Linear(4*c.width, c.width), nn.Dropout(c.dropout))\n",
    "        self.dropout = c.dropout\n",
    "\n",
    "    def forward(self, x):\n",
    "        b, t, w = x.shape\n",
    "        q, k, v = self.qkv(self.norm1(x)).chunk(3, dim=-1)\n",
    "        q, k, v = [z.view(b, t, self.heads, w//self.heads).transpose(1, 2)\n",
    "                   for z in (q, k, v)]\n",
    "        a = F.scaled_dot_product_attention(q, k, v, is_causal=True,\n",
    "                  dropout_p=self.dropout if self.training else 0.0)\n",
    "        x = x + self.proj(a.transpose(1, 2).contiguous().view(b, t, w))\n",
    "        return x + self.ff(self.norm2(x))\n",
    "\n",
    "\n",
    "class ScratchGPT(nn.Module):\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        if config.width % config.heads:\n",
    "            raise ValueError('Width must be divisible by number of heads.')\n",
    "        self.config = config\n",
    "        self.tokens = nn.Embedding(256, config.width)\n",
    "        self.positions = nn.Embedding(config.context, config.width)\n",
    "        self.blocks = nn.Sequential(*[Block(config) for _ in range(config.layers)])\n",
    "        self.norm = nn.LayerNorm(config.width)\n",
    "        self.output = nn.Linear(config.width, 256, bias=False)\n",
    "        self.output.weight = self.tokens.weight\n",
    "        self.apply(self._init)\n",
    "\n",
    "    @staticmethod\n",
    "    def _init(module):\n",
    "        if isinstance(module, (nn.Linear, nn.Embedding)):\n",
    "            nn.init.normal_(module.weight, mean=0, std=0.02)\n",
    "            if isinstance(module, nn.Linear) and module.bias is not None:\n",
    "                nn.init.zeros_(module.bias)\n",
    "\n",
    "    def forward(self, tokens, targets=None):\n",
    "        if tokens.size(1) > self.config.context:\n",
    "            raise ValueError('Sequence exceeds context.')\n",
    "        positions = torch.arange(tokens.size(1), device=tokens.device)\n",
    "        x = self.tokens(tokens) + self.positions(positions)\n",
    "        logits = self.output(self.norm(self.blocks(x)))\n",
    "        loss = None if targets is None else F.cross_entropy(\n",
    "            logits.reshape(-1, 256), targets.reshape(-1))\n",
    "        return logits, loss\n",
    "\n",
    "    @torch.no_grad()\n",
    "    def generate(self, prompt, length=300, temperature=0.8, seed=123):\n",
    "        if temperature <= 0:\n",
    "            raise ValueError('Temperature must be positive.')\n",
    "        self.eval()\n",
    "        device = next(self.parameters()).device\n",
    "        ids = list(prompt.encode('utf-8')) or [10]\n",
    "        tokens = torch.tensor([ids], device=device, dtype=torch.long)\n",
    "        rng = torch.Generator(device=device).manual_seed(seed)\n",
    "        for _ in range(length):\n",
    "            logits, _ = self(tokens[:, -self.config.context:])\n",
    "            probabilities = F.softmax(logits[:, -1] / temperature, dim=-1)\n",
    "            next_token = torch.multinomial(probabilities, 1, generator=rng)\n",
    "            tokens = torch.cat([tokens, next_token], dim=1)\n",
    "        return bytes(tokens[0].tolist()).decode('utf-8', errors='replace')\n",
    "\n",
    "\n",
    "def load_splits(folder, context):\n",
    "    result = {}\n",
    "    fingerprints = {}\n",
    "    for name in ('train', 'validation', 'test'):\n",
    "        raw = (Path(folder) / (name + '.txt')).read_bytes()\n",
    "        raw.decode('utf-8')  # Reject non-text inputs.\n",
    "        if len(raw) <= context:\n",
    "            raise ValueError(f'{name}.txt needs more than {context} UTF-8 bytes.')\n",
    "        fingerprints[name] = hashlib.sha256(raw).hexdigest()\n",
    "        result[name] = torch.tensor(list(raw), dtype=torch.long)\n",
    "    if len(set(fingerprints.values())) != 3:\n",
    "        raise ValueError('Training, validation and test files must differ.')\n",
    "    # Catch exact copied paragraphs; human review must also catch paraphrases.\n",
    "    paragraphs = {}\n",
    "    for name in result:\n",
    "        text = (Path(folder) / (name + '.txt')).read_text(encoding='utf-8')\n",
    "        paragraphs[name] = {' '.join(p.lower().split()) for p in text.split('\\n\\n')\n",
    "                            if len(p.strip()) >= 80}\n",
    "    for a, b in [('train', 'validation'), ('train', 'test'), ('validation', 'test')]:\n",
    "        if paragraphs[a] & paragraphs[b]:\n",
    "            raise ValueError(f'Copied paragraphs overlap between {a} and {b}.')\n",
    "    return result, fingerprints\n",
    "\n",
    "\n",
    "def batch(data, context, size, device, generator=None):\n",
    "    starts = torch.randint(len(data)-context, (size,), generator=generator)\n",
    "    x = torch.stack([data[i:i+context] for i in starts])\n",
    "    y = torch.stack([data[i+1:i+context+1] for i in starts])\n",
    "    return x.to(device), y.to(device)\n",
    "\n",
    "\n",
    "@torch.no_grad()\n",
    "def evaluate(model, data, batches=8, size=4):\n",
    "    was_training = model.training\n",
    "    model.eval()\n",
    "    rng = torch.Generator().manual_seed(987)\n",
    "    losses = []\n",
    "    for _ in range(batches):\n",
    "        x, y = batch(data, model.config.context, size, next(model.parameters()).device, rng)\n",
    "        losses.append(float(model(x, y)[1]))\n",
    "    model.train(was_training)\n",
    "    return sum(losses) / len(losses)\n",
    "\n",
    "\n",
    "def train(data_dir, output_dir, config=None, steps=500, batch_size=8,\n",
    "          accumulation=2, learning_rate=3e-4, save_every=50, resume=False,\n",
    "          allow_cpu=False, seed=42):\n",
    "    \"\"\"steps is total target step, also when resuming. No network or paid APIs.\"\"\"\n",
    "    if min(steps, batch_size, accumulation, save_every) < 1:\n",
    "        raise ValueError('Steps, batch size, accumulation and save interval must be positive.')\n",
    "    config = config or Config()\n",
    "    device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "    if device == 'cpu' and not allow_cpu:\n",
    "        raise RuntimeError('Select a GPU runtime, or explicitly set allow_cpu=True for a small test.')\n",
    "    random.seed(seed)\n",
    "    torch.manual_seed(seed)\n",
    "    if device == 'cpu':\n",
    "        torch.set_num_threads(min(4, torch.get_num_threads()))\n",
    "    data, fingerprints = load_splits(data_dir, config.context)\n",
    "    output = Path(output_dir)\n",
    "    output.mkdir(parents=True, exist_ok=True)\n",
    "    checkpoint = output / 'latest.pt'\n",
    "    if checkpoint.exists() and not resume:\n",
    "        raise ValueError('An existing run is present. Resume it or choose a different output folder.')\n",
    "    model = ScratchGPT(config).to(device)\n",
    "    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=0.1)\n",
    "    start = 0\n",
    "    history = []\n",
    "    settings = dict(batch_size=batch_size, accumulation=accumulation, learning_rate=learning_rate)\n",
    "    if resume:\n",
    "        if not checkpoint.exists():\n",
    "            raise FileNotFoundError('Restore latest.pt into the output folder before resuming.')\n",
    "        saved = torch.load(checkpoint, map_location=device, weights_only=True)\n",
    "        if saved['config'] != asdict(config) or saved['data_sha256'] != fingerprints:\n",
    "            raise ValueError('Resume requires the same architecture and exact dataset files.')\n",
    "        if saved['training_settings'] != settings:\n",
    "            raise ValueError('Resume requires the same batch, accumulation and learning rate.')\n",
    "        model.load_state_dict(saved['model'])\n",
    "        optimizer.load_state_dict(saved['optimizer'])\n",
    "        torch.set_rng_state(saved['rng'].cpu())\n",
    "        if device == 'cuda' and saved.get('cuda_rng'):\n",
    "            if len(saved['cuda_rng']) == torch.cuda.device_count():\n",
    "                torch.cuda.set_rng_state_all([r.cpu() for r in saved['cuda_rng']])\n",
    "        start, history = saved['step'], saved['history']\n",
    "    params = sum(p.numel() for p in model.parameters())\n",
    "    print(f'Device: {device}; parameters: {params:,}; completed steps: {start}')\n",
    "    if len(data['train']) < 1_000_000:\n",
    "        print('DEMO CORPUS: too small for a useful language or science model. Expect memorization.')\n",
    "    model.train()\n",
    "    started = time.time()\n",
    "\n",
    "    def save(step):\n",
    "        state = dict(model=model.state_dict(), optimizer=optimizer.state_dict(),\n",
    "                     config=asdict(config), data_sha256=fingerprints,\n",
    "                     step=step, history=history, training_settings=settings,\n",
    "                     rng=torch.get_rng_state(),\n",
    "                     cuda_rng=torch.cuda.get_rng_state_all() if device == 'cuda' else [])\n",
    "        temporary = output / 'latest.tmp'\n",
    "        torch.save(state, temporary)\n",
    "        temporary.replace(checkpoint)\n",
    "\n",
    "    for step in range(start+1, steps+1):\n",
    "        optimizer.zero_grad(set_to_none=True)\n",
    "        total = 0.0\n",
    "        for _ in range(accumulation):\n",
    "            x, y = batch(data['train'], config.context, batch_size, device)\n",
    "            _, loss = model(x, y)\n",
    "            if not torch.isfinite(loss):\n",
    "                raise RuntimeError('Nonfinite loss. Restore the previous checkpoint and inspect data/settings.')\n",
    "            (loss / accumulation).backward()\n",
    "            total += float(loss.detach()) / accumulation\n",
    "        nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",
    "        optimizer.step()\n",
    "        if step % save_every == 0 or step == steps:\n",
    "            val_loss = evaluate(model, data['validation'])\n",
    "            entry = dict(step=step, train_loss=total, validation_loss=val_loss)\n",
    "            history.append(entry)\n",
    "            save(step)\n",
    "            print(json.dumps(entry))\n",
    "\n",
    "    # Test is deliberately separate, run once after selecting the final model.\n",
    "    report = dict(parameters=params, device=device, completed_steps=max(start, steps),\n",
    "                  elapsed_seconds=round(time.time()-started, 2), history=history,\n",
    "                  data_sha256=fingerprints, config=asdict(config),\n",
    "                  warning='Loss is a text-prediction metric, not a factual-accuracy or truth score.')\n",
    "    (output / 'metrics.json').write_text(json.dumps(report, indent=2), encoding='utf-8')\n",
    "    torch.save(dict(model=model.state_dict(), config=asdict(config)), output / 'model.pt')\n",
    "    return model, report\n",
    "\n",
    "\n",
    "def load_model(path, device='cpu'):\n",
    "    saved = torch.load(path, map_location=device, weights_only=True)\n",
    "    model = ScratchGPT(Config(**saved['config'])).to(device)\n",
    "    model.load_state_dict(saved['model'])\n",
    "    model.eval()\n",
    "    return model",
    "\nENGINE_SOURCE = ['\"\"\"Pepaman byte-level causal transformer. Random weights, no pretrained assets.\"\"\"\\n', 'import hashlib\\n', 'import json\\n', 'import math\\n', 'import random\\n', 'import time\\n', 'from dataclasses import asdict, dataclass\\n', 'from pathlib import Path\\n', '\\n', 'import torch\\n', 'from torch import nn\\n', 'from torch.nn import functional as F\\n', '\\n', '\\n', '@dataclass\\n', 'class Config:\\n', '    width: int = 192\\n', '    layers: int = 4\\n', '    heads: int = 6\\n', '    context: int = 128\\n', '    dropout: float = 0.1\\n', '\\n', '\\n', 'class Block(nn.Module):\\n', '    def __init__(self, c):\\n', '        super().__init__()\\n', '        self.heads = c.heads\\n', '        self.norm1 = nn.LayerNorm(c.width)\\n', '        self.norm2 = nn.LayerNorm(c.width)\\n', '        self.qkv = nn.Linear(c.width, 3 * c.width)\\n', '        self.proj = nn.Linear(c.width, c.width)\\n', '        self.ff = nn.Sequential(nn.Linear(c.width, 4*c.width), nn.GELU(),\\n', '                                nn.Linear(4*c.width, c.width), nn.Dropout(c.dropout))\\n', '        self.dropout = c.dropout\\n', '\\n', '    def forward(self, x):\\n', '        b, t, w = x.shape\\n', '        q, k, v = self.qkv(self.norm1(x)).chunk(3, dim=-1)\\n', '        q, k, v = [z.view(b, t, self.heads, w//self.heads).transpose(1, 2)\\n', '                   for z in (q, k, v)]\\n', '        a = F.scaled_dot_product_attention(q, k, v, is_causal=True,\\n', '                  dropout_p=self.dropout if self.training else 0.0)\\n', '        x = x + self.proj(a.transpose(1, 2).contiguous().view(b, t, w))\\n', '        return x + self.ff(self.norm2(x))\\n', '\\n', '\\n', 'class ScratchGPT(nn.Module):\\n', '    def __init__(self, config):\\n', '        super().__init__()\\n', '        if config.width % config.heads:\\n', \"            raise ValueError('Width must be divisible by number of heads.')\\n\", '        self.config = config\\n', '        self.tokens = nn.Embedding(256, config.width)\\n', '        self.positions = nn.Embedding(config.context, config.width)\\n', '        self.blocks = nn.Sequential(*[Block(config) for _ in range(config.layers)])\\n', '        self.norm = nn.LayerNorm(config.width)\\n', '        self.output = nn.Linear(config.width, 256, bias=False)\\n', '        self.output.weight = self.tokens.weight\\n', '        self.apply(self._init)\\n', '\\n', '    @staticmethod\\n', '    def _init(module):\\n', '        if isinstance(module, (nn.Linear, nn.Embedding)):\\n', '            nn.init.normal_(module.weight, mean=0, std=0.02)\\n', '            if isinstance(module, nn.Linear) and module.bias is not None:\\n', '                nn.init.zeros_(module.bias)\\n', '\\n', '    def forward(self, tokens, targets=None):\\n', '        if tokens.size(1) > self.config.context:\\n', \"            raise ValueError('Sequence exceeds context.')\\n\", '        positions = torch.arange(tokens.size(1), device=tokens.device)\\n', '        x = self.tokens(tokens) + self.positions(positions)\\n', '        logits = self.output(self.norm(self.blocks(x)))\\n', '        loss = None if targets is None else F.cross_entropy(\\n', '            logits.reshape(-1, 256), targets.reshape(-1))\\n', '        return logits, loss\\n', '\\n', '    @torch.no_grad()\\n', '    def generate(self, prompt, length=300, temperature=0.8, seed=123):\\n', '        if temperature <= 0:\\n', \"            raise ValueError('Temperature must be positive.')\\n\", '        self.eval()\\n', '        device = next(self.parameters()).device\\n', \"        ids = list(prompt.encode('utf-8')) or [10]\\n\", '        tokens = torch.tensor([ids], device=device, dtype=torch.long)\\n', '        rng = torch.Generator(device=device).manual_seed(seed)\\n', '        for _ in range(length):\\n', '            logits, _ = self(tokens[:, -self.config.context:])\\n', '            probabilities = F.softmax(logits[:, -1] / temperature, dim=-1)\\n', '            next_token = torch.multinomial(probabilities, 1, generator=rng)\\n', '            tokens = torch.cat([tokens, next_token], dim=1)\\n', \"        return bytes(tokens[0].tolist()).decode('utf-8', errors='replace')\\n\", '\\n', '\\n', 'def load_splits(folder, context):\\n', '    result = {}\\n', '    fingerprints = {}\\n', \"    for name in ('train', 'validation', 'test'):\\n\", \"        raw = (Path(folder) / (name + '.txt')).read_bytes()\\n\", \"        raw.decode('utf-8')  # Reject non-text inputs.\\n\", '        if len(raw) <= context:\\n', \"            raise ValueError(f'{name}.txt needs more than {context} UTF-8 bytes.')\\n\", '        fingerprints[name] = hashlib.sha256(raw).hexdigest()\\n', '        result[name] = torch.tensor(list(raw), dtype=torch.long)\\n', '    if len(set(fingerprints.values())) != 3:\\n', \"        raise ValueError('Training, validation and test files must differ.')\\n\", '    # Catch exact copied paragraphs; human review must also catch paraphrases.\\n', '    paragraphs = {}\\n', '    for name in result:\\n', \"        text = (Path(folder) / (name + '.txt')).read_text(encoding='utf-8')\\n\", \"        paragraphs[name] = {' '.join(p.lower().split()) for p in text.split('\\\\n\\\\n')\\n\", '                            if len(p.strip()) >= 80}\\n', \"    for a, b in [('train', 'validation'), ('train', 'test'), ('validation', 'test')]:\\n\", '        if paragraphs[a] & paragraphs[b]:\\n', \"            raise ValueError(f'Copied paragraphs overlap between {a} and {b}.')\\n\", '    return result, fingerprints\\n', '\\n', '\\n', 'def batch(data, context, size, device, generator=None):\\n', '    starts = torch.randint(len(data)-context, (size,), generator=generator)\\n', '    x = torch.stack([data[i:i+context] for i in starts])\\n', '    y = torch.stack([data[i+1:i+context+1] for i in starts])\\n', '    return x.to(device), y.to(device)\\n', '\\n', '\\n', '@torch.no_grad()\\n', 'def evaluate(model, data, batches=8, size=4):\\n', '    was_training = model.training\\n', '    model.eval()\\n', '    rng = torch.Generator().manual_seed(987)\\n', '    losses = []\\n', '    for _ in range(batches):\\n', '        x, y = batch(data, model.config.context, size, next(model.parameters()).device, rng)\\n', '        losses.append(float(model(x, y)[1]))\\n', '    model.train(was_training)\\n', '    return sum(losses) / len(losses)\\n', '\\n', '\\n', 'def train(data_dir, output_dir, config=None, steps=500, batch_size=8,\\n', '          accumulation=2, learning_rate=3e-4, save_every=50, resume=False,\\n', '          allow_cpu=False, seed=42):\\n', '    \"\"\"steps is total target step, also when resuming. No network or paid APIs.\"\"\"\\n', '    if min(steps, batch_size, accumulation, save_every) < 1:\\n', \"        raise ValueError('Steps, batch size, accumulation and save interval must be positive.')\\n\", '    config = config or Config()\\n', \"    device = 'cuda' if torch.cuda.is_available() else 'cpu'\\n\", \"    if device == 'cpu' and not allow_cpu:\\n\", \"        raise RuntimeError('Select a GPU runtime, or explicitly set allow_cpu=True for a small test.')\\n\", '    random.seed(seed)\\n', '    torch.manual_seed(seed)\\n', \"    if device == 'cpu':\\n\", '        torch.set_num_threads(min(4, torch.get_num_threads()))\\n', '    data, fingerprints = load_splits(data_dir, config.context)\\n', '    output = Path(output_dir)\\n', '    output.mkdir(parents=True, exist_ok=True)\\n', \"    checkpoint = output / 'latest.pt'\\n\", '    if checkpoint.exists() and not resume:\\n', \"        raise ValueError('An existing run is present. Resume it or choose a different output folder.')\\n\", '    model = ScratchGPT(config).to(device)\\n', '    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=0.1)\\n', '    start = 0\\n', '    history = []\\n', '    settings = dict(batch_size=batch_size, accumulation=accumulation, learning_rate=learning_rate)\\n', '    if resume:\\n', '        if not checkpoint.exists():\\n', \"            raise FileNotFoundError('Restore latest.pt into the output folder before resuming.')\\n\", '        saved = torch.load(checkpoint, map_location=device, weights_only=True)\\n', \"        if saved['config'] != asdict(config) or saved['data_sha256'] != fingerprints:\\n\", \"            raise ValueError('Resume requires the same architecture and exact dataset files.')\\n\", \"        if saved['training_settings'] != settings:\\n\", \"            raise ValueError('Resume requires the same batch, accumulation and learning rate.')\\n\", \"        model.load_state_dict(saved['model'])\\n\", \"        optimizer.load_state_dict(saved['optimizer'])\\n\", \"        torch.set_rng_state(saved['rng'].cpu())\\n\", \"        if device == 'cuda' and saved.get('cuda_rng'):\\n\", \"            if len(saved['cuda_rng']) == torch.cuda.device_count():\\n\", \"                torch.cuda.set_rng_state_all([r.cpu() for r in saved['cuda_rng']])\\n\", \"        start, history = saved['step'], saved['history']\\n\", '    params = sum(p.numel() for p in model.parameters())\\n', \"    print(f'Device: {device}; parameters: {params:,}; completed steps: {start}')\\n\", \"    if len(data['train']) < 1_000_000:\\n\", \"        print('DEMO CORPUS: too small for a useful language or science model. Expect memorization.')\\n\", '    model.train()\\n', '    started = time.time()\\n', '\\n', '    def save(step):\\n', '        state = dict(model=model.state_dict(), optimizer=optimizer.state_dict(),\\n', '                     config=asdict(config), data_sha256=fingerprints,\\n', '                     step=step, history=history, training_settings=settings,\\n', '                     rng=torch.get_rng_state(),\\n', \"                     cuda_rng=torch.cuda.get_rng_state_all() if device == 'cuda' else [])\\n\", \"        temporary = output / 'latest.tmp'\\n\", '        torch.save(state, temporary)\\n', '        temporary.replace(checkpoint)\\n', '\\n', '    for step in range(start+1, steps+1):\\n', '        optimizer.zero_grad(set_to_none=True)\\n', '        total = 0.0\\n', '        for _ in range(accumulation):\\n', \"            x, y = batch(data['train'], config.context, batch_size, device)\\n\", '            _, loss = model(x, y)\\n', '            if not torch.isfinite(loss):\\n', \"                raise RuntimeError('Nonfinite loss. Restore the previous checkpoint and inspect data/settings.')\\n\", '            (loss / accumulation).backward()\\n', '            total += float(loss.detach()) / accumulation\\n', '        nn.utils.clip_grad_norm_(model.parameters(), 1.0)\\n', '        optimizer.step()\\n', '        if step % save_every == 0 or step == steps:\\n', \"            val_loss = evaluate(model, data['validation'])\\n\", '            entry = dict(step=step, train_loss=total, validation_loss=val_loss)\\n', '            history.append(entry)\\n', '            save(step)\\n', '            print(json.dumps(entry))\\n', '\\n', '    # Test is deliberately separate, run once after selecting the final model.\\n', '    report = dict(parameters=params, device=device, completed_steps=max(start, steps),\\n', '                  elapsed_seconds=round(time.time()-started, 2), history=history,\\n', '                  data_sha256=fingerprints, config=asdict(config),\\n', \"                  warning='Loss is a text-prediction metric, not a factual-accuracy or truth score.')\\n\", \"    (output / 'metrics.json').write_text(json.dumps(report, indent=2), encoding='utf-8')\\n\", \"    torch.save(dict(model=model.state_dict(), config=asdict(config)), output / 'model.pt')\\n\", '    return model, report\\n', '\\n', '\\n', \"def load_model(path, device='cpu'):\\n\", '    saved = torch.load(path, map_location=device, weights_only=True)\\n', \"    model = ScratchGPT(Config(**saved['config'])).to(device)\\n\", \"    model.load_state_dict(saved['model'])\\n\", '    model.eval()\\n', '    return model\\n']\n"
   ],
   "execution_count": null,
   "outputs": [],
   "id": "pepaman-07"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "CONFIG = Config(width=192, layers=4, heads=6, context=128, dropout=0.1)\n",
    "# For a CPU test only: Config(width=64, layers=2, heads=4, context=64)\n",
    "splits, hashes = load_splits(DATA_DIR, CONFIG.context)\n",
    "print({name: len(values) for name, values in splits.items()}, 'bytes')\n",
    "print('SHA256:', hashes)\n",
    "print('Model parameters:', sum(p.numel() for p in ScratchGPT(CONFIG).parameters()))"
   ],
   "execution_count": null,
   "outputs": [],
   "id": "pepaman-08"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Train\n",
    "Run this cell while actively using the notebook. It uses only the selected runtime. Reduce batch size/context if memory runs out. It saves `latest.pt` every 50 steps, with optimizer and random state for resumption. A new run will not overwrite an existing checkpoint unless you explicitly resume it."
   ],
   "id": "pepaman-09"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "model, report = train(DATA_DIR, OUTPUT_DIR, config=CONFIG, steps=STEPS,\n",
    "    batch_size=BATCH_SIZE, accumulation=ACCUMULATION, save_every=50,\n",
    "    resume=RESUME, allow_cpu=ALLOW_CPU)"
   ],
   "execution_count": null,
   "outputs": [],
   "id": "pepaman-10"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Inspect and evaluate\n",
    "Validation loss measures next-byte prediction, not scientific accuracy. Compare factual answers against sources and have a human review uncertainty, fabricated citations, repetition and misleading confidence. A tiny raw language model may not follow questions at all. Test data should be consulted once after model selection, not used to tune repeatedly."
   ],
   "id": "pepaman-11"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "print(model.generate('The universe ', length=250))\n",
    "# Only after choosing the final configuration:\n",
    "RUN_FINAL_TEST = False\n",
    "if RUN_FINAL_TEST:\n",
    "    test_loss = evaluate(model, splits['test'])\n",
    "    print('Held-out next-byte loss:', test_loss, '(not a truth score)')\n",
    "    (OUTPUT_DIR/'test-metrics.json').write_text(json.dumps({'test_loss':test_loss}),encoding='utf-8')"
   ],
   "execution_count": null,
   "outputs": [],
   "id": "pepaman-12"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Export and resume\n",
    "The archive contains weights, optimizer checkpoint, configuration, metrics and the model code. Keep your exact three dataset files separately; their hashes are checked on resume. For a custom corpus, the bundled NASA notes describe only the demonstration, not your custom sources. Free runtimes stop, and GPU access is not guaranteed for life. Checkpoint portability lets you continue work; it does not pool separate providers into one GPU."
   ],
   "id": "pepaman-13"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "import shutil\n",
    "(OUTPUT_DIR/'scratch.py').write_text(''.join(ENGINE_SOURCE),encoding='utf-8')\n",
    "(OUTPUT_DIR/'demo-sources.json').write_text(json.dumps(SOURCES,indent=2),encoding='utf-8')\n",
    "archive = shutil.make_archive(str(OUTPUT_DIR)+'-export', 'zip', OUTPUT_DIR)\n",
    "print('Download:', archive)\n",
    "if Path('/kaggle').exists():\n",
    "    from IPython.display import FileLink, display\n",
    "    display(FileLink(Path(archive).name))\n",
    "    print('Also available in the Kaggle Output panel. Save a version with outputs before stopping the session.')\n",
    "else:\n",
    "    from google.colab import files\n",
    "    files.download(archive)"
   ],
   "execution_count": null,
   "outputs": [],
   "id": "pepaman-14"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  },
  "accelerator": "GPU"
 },
 "nbformat": 4,
 "nbformat_minor": 5
}