{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cf36d705",
   "metadata": {},
   "source": [
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "ae78ea2d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: datasets in /Users/sameerchandra/.pyenv/versions/3.13.0/lib/python3.13/site-packages (5.0.0)\n",
      "Requirement already satisfied: filelock in /Users/sameerchandra/.pyenv/versions/3.13.0/lib/python3.13/site-packages (from datasets) (3.18.0)\n",
      "Requirement already satisfied: numpy>=1.17 in /Users/sameerchandra/.pyenv/versions/3.13.0/lib/python3.13/site-packages (from datasets) (2.2.6)\n",
      "Requirement already satisfied: pyarrow>=21.0.0 in /Users/sameerchandra/.pyenv/versions/3.13.0/lib/python3.13/site-packages (from datasets) (24.0.0)\n",
      "Requirement already satisfied: dill<0.4.2,>=0.3.0 in /Users/sameerchandra/.pyenv/versions/3.13.0/lib/python3.13/site-packages (from datasets) (0.4.1)\n",
      "Requirement already satisfied: pandas in /Users/sameerchandra/.pyenv/versions/3.13.0/lib/python3.13/site-packages (from datasets) (3.0.3)\n",
      "Requirement already satisfied: requests>=2.32.2 in /Users/sameerchandra/.pyenv/versions/3.13.0/lib/python3.13/site-packages (from datasets) (2.32.4)\n",
      "Requirement already satisfied: httpx<1.0.0 in /Users/sameerchandra/.pyenv/versions/3.13.0/lib/python3.13/site-packages (from datasets) (0.28.1)\n",
      "Requirement already satisfied: tqdm>=4.66.3 in /Users/sameerchandra/.pyenv/versions/3.13.0/lib/python3.13/site-packages (from datasets) (4.67.1)\n",
      "Requirement already satisfied: xxhash in /Users/sameerchandra/.pyenv/versions/3.13.0/lib/python3.13/site-packages (from datasets) (3.7.0)\n",
      "Requirement already satisfied: multiprocess<0.70.20 in /Users/sameerchandra/.pyenv/versions/3.13.0/lib/python3.13/site-packages (from datasets) (0.70.19)\n",
      "Requirement already satisfied: fsspec<=2026.4.0,>=2023.1.0 in /Users/sameerchandra/.pyenv/versions/3.13.0/lib/python3.13/site-packages (from fsspec[http]<=2026.4.0,>=2023.1.0->datasets) (2025.7.0)\n",
      "Requirement already satisfied: huggingface-hub<2.0,>=0.25.0 in /Users/sameerchandra/.pyenv/versions/3.13.0/lib/python3.13/site-packages (from datasets) (1.19.0)\n",
      "Requirement already satisfied: packaging in /Users/sameerchandra/.pyenv/versions/3.13.0/lib/python3.13/site-packages (from datasets) (25.0)\n",
      "Requirement already satisfied: pyyaml>=5.1 in /Users/sameerchandra/.pyenv/versions/3.13.0/lib/python3.13/site-packages (from datasets) (6.0.2)\n",
      "Requirement already satisfied: aiohttp!=4.0.0a0,!=4.0.0a1 in /Users/sameerchandra/.pyenv/versions/3.13.0/lib/python3.13/site-packages (from fsspec[http]<=2026.4.0,>=2023.1.0->datasets) (3.12.15)\n",
      "Requirement already satisfied: anyio in /Users/sameerchandra/.pyenv/versions/3.13.0/lib/python3.13/site-packages (from httpx<1.0.0->datasets) (4.13.0)\n",
      "Requirement already satisfied: certifi in /Users/sameerchandra/.pyenv/versions/3.13.0/lib/python3.13/site-packages (from httpx<1.0.0->datasets) (2025.8.3)\n",
      "Requirement already satisfied: httpcore==1.* in /Users/sameerchandra/.pyenv/versions/3.13.0/lib/python3.13/site-packages (from httpx<1.0.0->datasets) (1.0.9)\n",
      "Requirement already satisfied: idna in /Users/sameerchandra/.pyenv/versions/3.13.0/lib/python3.13/site-packages (from httpx<1.0.0->datasets) (3.10)\n",
      "Requirement already satisfied: h11>=0.16 in /Users/sameerchandra/.pyenv/versions/3.13.0/lib/python3.13/site-packages (from httpcore==1.*->httpx<1.0.0->datasets) (0.16.0)\n",
      "Requirement already satisfied: click>=8.4.0 in /Users/sameerchandra/.pyenv/versions/3.13.0/lib/python3.13/site-packages (from huggingface-hub<2.0,>=0.25.0->datasets) (8.4.1)\n",
      "Requirement already satisfied: hf-xet<2.0.0,>=1.5.1 in /Users/sameerchandra/.pyenv/versions/3.13.0/lib/python3.13/site-packages (from huggingface-hub<2.0,>=0.25.0->datasets) (1.5.1)\n",
      "Requirement already satisfied: typer<0.26.0,>=0.20.0 in /Users/sameerchandra/.pyenv/versions/3.13.0/lib/python3.13/site-packages (from huggingface-hub<2.0,>=0.25.0->datasets) (0.25.1)\n",
      "Requirement already satisfied: typing-extensions>=4.1.0 in /Users/sameerchandra/.pyenv/versions/3.13.0/lib/python3.13/site-packages (from huggingface-hub<2.0,>=0.25.0->datasets) (4.15.0)\n",
      "Requirement already satisfied: shellingham>=1.3.0 in /Users/sameerchandra/.pyenv/versions/3.13.0/lib/python3.13/site-packages (from typer<0.26.0,>=0.20.0->huggingface-hub<2.0,>=0.25.0->datasets) (1.5.4)\n",
      "Requirement already satisfied: rich>=13.8.0 in /Users/sameerchandra/.pyenv/versions/3.13.0/lib/python3.13/site-packages (from typer<0.26.0,>=0.20.0->huggingface-hub<2.0,>=0.25.0->datasets) (15.0.0)\n",
      "Requirement already satisfied: annotated-doc>=0.0.2 in /Users/sameerchandra/.pyenv/versions/3.13.0/lib/python3.13/site-packages (from typer<0.26.0,>=0.20.0->huggingface-hub<2.0,>=0.25.0->datasets) (0.0.4)\n",
      "Requirement already satisfied: aiohappyeyeballs>=2.5.0 in /Users/sameerchandra/.pyenv/versions/3.13.0/lib/python3.13/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<=2026.4.0,>=2023.1.0->datasets) (2.6.1)\n",
      "Requirement already satisfied: aiosignal>=1.4.0 in /Users/sameerchandra/.pyenv/versions/3.13.0/lib/python3.13/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<=2026.4.0,>=2023.1.0->datasets) (1.4.0)\n",
      "Requirement already satisfied: attrs>=17.3.0 in /Users/sameerchandra/.pyenv/versions/3.13.0/lib/python3.13/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<=2026.4.0,>=2023.1.0->datasets) (25.3.0)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in /Users/sameerchandra/.pyenv/versions/3.13.0/lib/python3.13/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<=2026.4.0,>=2023.1.0->datasets) (1.7.0)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /Users/sameerchandra/.pyenv/versions/3.13.0/lib/python3.13/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<=2026.4.0,>=2023.1.0->datasets) (6.6.3)\n",
      "Requirement already satisfied: propcache>=0.2.0 in /Users/sameerchandra/.pyenv/versions/3.13.0/lib/python3.13/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<=2026.4.0,>=2023.1.0->datasets) (0.3.2)\n",
      "Requirement already satisfied: yarl<2.0,>=1.17.0 in /Users/sameerchandra/.pyenv/versions/3.13.0/lib/python3.13/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<=2026.4.0,>=2023.1.0->datasets) (1.20.1)\n",
      "Requirement already satisfied: charset_normalizer<4,>=2 in /Users/sameerchandra/.pyenv/versions/3.13.0/lib/python3.13/site-packages (from requests>=2.32.2->datasets) (3.4.2)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /Users/sameerchandra/.pyenv/versions/3.13.0/lib/python3.13/site-packages (from requests>=2.32.2->datasets) (2.5.0)\n",
      "Requirement already satisfied: markdown-it-py>=2.2.0 in /Users/sameerchandra/.pyenv/versions/3.13.0/lib/python3.13/site-packages (from rich>=13.8.0->typer<0.26.0,>=0.20.0->huggingface-hub<2.0,>=0.25.0->datasets) (4.2.0)\n",
      "Requirement already satisfied: pygments<3.0.0,>=2.13.0 in /Users/sameerchandra/.pyenv/versions/3.13.0/lib/python3.13/site-packages (from rich>=13.8.0->typer<0.26.0,>=0.20.0->huggingface-hub<2.0,>=0.25.0->datasets) (2.20.0)\n",
      "Requirement already satisfied: mdurl~=0.1 in /Users/sameerchandra/.pyenv/versions/3.13.0/lib/python3.13/site-packages (from markdown-it-py>=2.2.0->rich>=13.8.0->typer<0.26.0,>=0.20.0->huggingface-hub<2.0,>=0.25.0->datasets) (0.1.2)\n",
      "Requirement already satisfied: python-dateutil>=2.8.2 in /Users/sameerchandra/.pyenv/versions/3.13.0/lib/python3.13/site-packages (from pandas->datasets) (2.9.0.post0)\n",
      "Requirement already satisfied: six>=1.5 in /Users/sameerchandra/.pyenv/versions/3.13.0/lib/python3.13/site-packages (from python-dateutil>=2.8.2->pandas->datasets) (1.17.0)\n",
      "\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m25.2\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m26.1.2\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "%pip install datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "192bbd24",
   "metadata": {},
   "outputs": [],
   "source": [
    "import mlx.core as mx\n",
    "import mlx.nn as nn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "213ca177",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TokenAndPositionEmbedding(nn.Module):\n",
    "    def __init__(self, maxlen: int, vocab_size:int, embed_dim: int):\n",
    "        super().__init__()\n",
    "        self.token_emb = nn.Embedding(vocab_size, embed_dim)\n",
    "        self.pos_emb = nn.Embedding(maxlen, embed_dim)\n",
    "\n",
    "    def __call__(self, x):\n",
    "        seq_len = x.shape[1]\n",
    "        positions = mx.arange(seq_len)[None, :]\n",
    "        return self.token_emb(x) + self.pos_emb(positions) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "58ab2f35",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TransformerBlock(nn.Module):\n",
    "    def __init__(self, emed_dim: int, num_heads:int, ff_dim: int):\n",
    "        super().__init__()\n",
    "        self.attention = nn.MultiHeadAttention(emed_dim, num_heads)\n",
    "    def __call__(self, x, mask=None):\n",
    "        attn_out = self.attention(x, x, x, mask=mask)\n",
    "        x = x + attn_out\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "1713a3a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class NanoLLM(nn.Module):\n",
    "    def __init__(self, maxlen: int, vocab_size:int, embed_dim:int, num_heads:int, feed_forward_dim:int , num_transformer_blocks: int):\n",
    "        super().__init__()\n",
    "        self.maxlen = maxlen\n",
    "        self.embedding = TokenAndPositionEmbedding(maxlen, vocab_size, embed_dim)\n",
    "        self.transformer_blocks = [\n",
    "            TransformerBlock(embed_dim, num_heads, feed_forward_dim)\n",
    "            for _ in range(num_transformer_blocks)\n",
    "        ]\n",
    "        self.output_layer = nn.Linear(embed_dim, vocab_size, bias=False)\n",
    "\n",
    "    def __call__(self, token_ids):\n",
    "        seq_len = token_ids.shape[1]\n",
    "\n",
    "        mask = nn.MultiHeadAttention.create_additive_causal_mask(seq_len)\n",
    "\n",
    "        x = self.embedding(token_ids)\n",
    "        for block in self.transformer_blocks:\n",
    "            x = block(x, mask=mask)\n",
    "\n",
    "        logits = self.output_layer(x)\n",
    "        return logits\n",
    "\n",
    "\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "cdd65e7c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading TinyStories from Hugging Face...\n",
      "Selecting 100000 stories...\n",
      "Formatting and writing to TinyStories-100k.txt...\n",
      "Done! Dataset saved to TinyStories-100k.txt (85.77 MB)\n"
     ]
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "import os\n",
    "\n",
    "def download_tinystories(output_file=\"TinyStories-100k.txt\", num_stories=100000):\n",
    "    print(f\"Downloading TinyStories from Hugging Face...\")\n",
    "    # Load the training split\n",
    "    dataset = load_dataset(\"roneneldan/TinyStories\", split=\"train\")\n",
    "    \n",
    "    # Slice exactly the number of stories you want\n",
    "    print(f\"Selecting {num_stories} stories...\")\n",
    "    subset = dataset.select(range(num_stories))\n",
    "    \n",
    "    # Write to a local text file\n",
    "    print(f\"Formatting and writing to {output_file}...\")\n",
    "    with open(output_file, \"w\", encoding=\"utf-8\") as f:\n",
    "        for item in subset:\n",
    "            story = item['text'].strip()\n",
    "            if story:\n",
    "                # Append the EOS token between stories\n",
    "                f.write(story + \"\\n<|endoftext|>\\n\")\n",
    "                \n",
    "    file_size_mb = os.path.getsize(output_file) / (1024 * 1024)\n",
    "    print(f\"Done! Dataset saved to {output_file} ({file_size_mb:.2f} MB)\")\n",
    "\n",
    "# Run the download step\n",
    "download_tinystories(output_file=\"TinyStories-100k.txt\", num_stories=100000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "861052da",
   "metadata": {},
   "outputs": [],
   "source": [
    "import mlx.core as mx\n",
    "import mlx.nn as nn\n",
    "import mlx.optimizers as optim\n",
    "import numpy as np\n",
    "import tiktoken\n",
    "\n",
    "def get_batches(file_path, tokenizer, batch_size=32, maxlen=128):\n",
    "    \"\"\"A lightweight generator to replace the Grain dataloader.\"\"\"\n",
    "    with open(file_path, 'r', encoding='utf-8') as f:\n",
    "        data = f.read()\n",
    "    \n",
    "    # Split by the end token just like the JAX setup\n",
    "    stories = data.split('<|endoftext|>')\n",
    "    # Remove empty strings\n",
    "    stories = [s.strip() + '<|endoftext|>' for s in stories if s.strip()]\n",
    "    \n",
    "    # Shuffle the dataset\n",
    "    np.random.shuffle(stories)\n",
    "    \n",
    "    batch_x, batch_y = [], []\n",
    "    \n",
    "    for story in stories:\n",
    "        tokens = tokenizer.encode(story, allowed_special={'<|endoftext|>'})\n",
    "        \n",
    "        # We need maxlen + 1 tokens because we shift the targets by 1\n",
    "        # Example: if maxlen is 128, we need 129 tokens to get 128 inputs and 128 targets\n",
    "        if len(tokens) > maxlen + 1:\n",
    "            tokens = tokens[:maxlen + 1]\n",
    "        \n",
    "        # Pad with 0s if it's too short\n",
    "        if len(tokens) < maxlen + 1:\n",
    "            tokens.extend([0] * (maxlen + 1 - len(tokens)))\n",
    "            \n",
    "        # x is the input sequence, y is the target (shifted by 1)\n",
    "        batch_x.append(tokens[:-1])\n",
    "        batch_y.append(tokens[1:])\n",
    "        \n",
    "        if len(batch_x) == batch_size:\n",
    "            # Yield MLX arrays directly\n",
    "            yield mx.array(batch_x), mx.array(batch_y)\n",
    "            batch_x, batch_y = [], []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "e3cc2b44",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1. Initialize Tokenizer and Model\n",
    "tokenizer = tiktoken.get_encoding(\"gpt2\")\n",
    "model = NanoLLM(\n",
    "    maxlen=128, \n",
    "    vocab_size=tokenizer.n_vocab, \n",
    "    embed_dim=192, \n",
    "    num_heads=6, \n",
    "    feed_forward_dim=512, \n",
    "    num_transformer_blocks=6\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "966d0878",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 1 | Step: 0 | Loss: 10.8266\n",
      "Epoch: 1 | Step: 50 | Loss: 5.9461\n",
      "Epoch: 1 | Step: 100 | Loss: 5.6107\n",
      "Epoch: 1 | Step: 150 | Loss: 5.2950\n",
      "Epoch: 1 | Step: 200 | Loss: 4.9307\n",
      "Epoch: 1 | Step: 250 | Loss: 4.5847\n",
      "Epoch: 1 | Step: 300 | Loss: 4.2845\n",
      "Epoch: 1 | Step: 350 | Loss: 4.3064\n",
      "Epoch: 1 | Step: 400 | Loss: 4.0243\n",
      "Epoch: 1 | Step: 450 | Loss: 3.8484\n",
      "Epoch: 1 | Step: 500 | Loss: 3.7691\n",
      "Epoch: 1 | Step: 550 | Loss: 3.8548\n",
      "Epoch: 1 | Step: 600 | Loss: 3.9880\n",
      "Epoch: 1 | Step: 650 | Loss: 3.8109\n",
      "Epoch: 1 | Step: 700 | Loss: 3.8814\n",
      "Epoch: 1 | Step: 750 | Loss: 3.6469\n",
      "Epoch: 1 | Step: 800 | Loss: 3.6962\n",
      "Epoch: 1 | Step: 850 | Loss: 3.5966\n",
      "Epoch: 1 | Step: 900 | Loss: 3.5697\n",
      "Epoch: 1 | Step: 950 | Loss: 3.6199\n",
      "Epoch: 1 | Step: 1000 | Loss: 3.7395\n",
      "Epoch: 1 | Step: 1050 | Loss: 3.5999\n",
      "Epoch: 1 | Step: 1100 | Loss: 3.6092\n",
      "Epoch: 1 | Step: 1150 | Loss: 3.6909\n",
      "Epoch: 1 | Step: 1200 | Loss: 3.4631\n",
      "Epoch: 1 | Step: 1250 | Loss: 3.6353\n",
      "Epoch: 1 | Step: 1300 | Loss: 3.4745\n",
      "Epoch: 1 | Step: 1350 | Loss: 3.5697\n",
      "Epoch: 1 | Step: 1400 | Loss: 3.6310\n",
      "Epoch: 1 | Step: 1450 | Loss: 3.3946\n",
      "Epoch: 1 | Step: 1500 | Loss: 3.3186\n",
      "Epoch: 1 | Step: 1550 | Loss: 3.4553\n",
      "Epoch: 1 | Step: 1600 | Loss: 3.5321\n",
      "Epoch: 1 | Step: 1650 | Loss: 3.5969\n",
      "Epoch: 1 | Step: 1700 | Loss: 3.3159\n",
      "Epoch: 1 | Step: 1750 | Loss: 3.2411\n",
      "Epoch: 1 | Step: 1800 | Loss: 3.1364\n",
      "Epoch: 1 | Step: 1850 | Loss: 3.4197\n",
      "Epoch: 1 | Step: 1900 | Loss: 3.3262\n",
      "Epoch: 1 | Step: 1950 | Loss: 3.2812\n",
      "Epoch: 1 | Step: 2000 | Loss: 3.4379\n",
      "Epoch: 1 | Step: 2050 | Loss: 3.2536\n",
      "Epoch: 1 | Step: 2100 | Loss: 3.3224\n",
      "Epoch: 1 | Step: 2150 | Loss: 3.2211\n",
      "Epoch: 1 | Step: 2200 | Loss: 3.2087\n",
      "Epoch: 1 | Step: 2250 | Loss: 3.2078\n",
      "Epoch: 1 | Step: 2300 | Loss: 3.1197\n",
      "Epoch: 1 | Step: 2350 | Loss: 3.2666\n",
      "Epoch: 1 | Step: 2400 | Loss: 3.2225\n",
      "Epoch: 1 | Step: 2450 | Loss: 3.2007\n",
      "Epoch: 1 | Step: 2500 | Loss: 3.1944\n",
      "Epoch: 1 | Step: 2550 | Loss: 3.0893\n",
      "Epoch: 1 | Step: 2600 | Loss: 3.2330\n",
      "Epoch: 1 | Step: 2650 | Loss: 3.1280\n",
      "Epoch: 1 | Step: 2700 | Loss: 3.1640\n",
      "Epoch: 1 | Step: 2750 | Loss: 2.9679\n",
      "Epoch: 1 | Step: 2800 | Loss: 2.9453\n",
      "Epoch: 1 | Step: 2850 | Loss: 3.0761\n",
      "Epoch: 1 | Step: 2900 | Loss: 3.2077\n",
      "Epoch: 1 | Step: 2950 | Loss: 3.0634\n",
      "Epoch: 1 | Step: 3000 | Loss: 3.1369\n",
      "Epoch: 1 | Step: 3050 | Loss: 3.0469\n",
      "Epoch: 1 | Step: 3100 | Loss: 2.9706\n",
      "Epoch: 2 | Step: 3150 | Loss: 3.0505\n",
      "Epoch: 2 | Step: 3200 | Loss: 2.9739\n",
      "Epoch: 2 | Step: 3250 | Loss: 2.9699\n",
      "Epoch: 2 | Step: 3300 | Loss: 3.0688\n",
      "Epoch: 2 | Step: 3350 | Loss: 2.9738\n",
      "Epoch: 2 | Step: 3400 | Loss: 2.9960\n",
      "Epoch: 2 | Step: 3450 | Loss: 3.1820\n",
      "Epoch: 2 | Step: 3500 | Loss: 3.0580\n",
      "Epoch: 2 | Step: 3550 | Loss: 3.0920\n",
      "Epoch: 2 | Step: 3600 | Loss: 2.8166\n",
      "Epoch: 2 | Step: 3650 | Loss: 3.0154\n",
      "Epoch: 2 | Step: 3700 | Loss: 2.9677\n",
      "Epoch: 2 | Step: 3750 | Loss: 3.1327\n",
      "Epoch: 2 | Step: 3800 | Loss: 2.9152\n",
      "Epoch: 2 | Step: 3850 | Loss: 3.0307\n",
      "Epoch: 2 | Step: 3900 | Loss: 2.9989\n",
      "Epoch: 2 | Step: 3950 | Loss: 2.9257\n",
      "Epoch: 2 | Step: 4000 | Loss: 2.9208\n",
      "Epoch: 2 | Step: 4050 | Loss: 3.0419\n",
      "Epoch: 2 | Step: 4100 | Loss: 2.9438\n",
      "Epoch: 2 | Step: 4150 | Loss: 2.9466\n",
      "Epoch: 2 | Step: 4200 | Loss: 2.8878\n",
      "Epoch: 2 | Step: 4250 | Loss: 2.8965\n",
      "Epoch: 2 | Step: 4300 | Loss: 2.9472\n",
      "Epoch: 2 | Step: 4350 | Loss: 2.9109\n",
      "Epoch: 2 | Step: 4400 | Loss: 3.0593\n",
      "Epoch: 2 | Step: 4450 | Loss: 2.9651\n",
      "Epoch: 2 | Step: 4500 | Loss: 3.1125\n",
      "Epoch: 2 | Step: 4550 | Loss: 2.9089\n",
      "Epoch: 2 | Step: 4600 | Loss: 2.8422\n",
      "Epoch: 2 | Step: 4650 | Loss: 2.8573\n",
      "Epoch: 2 | Step: 4700 | Loss: 2.8248\n",
      "Epoch: 2 | Step: 4750 | Loss: 3.0166\n",
      "Epoch: 2 | Step: 4800 | Loss: 2.8956\n",
      "Epoch: 2 | Step: 4850 | Loss: 2.6628\n",
      "Epoch: 2 | Step: 4900 | Loss: 2.9485\n",
      "Epoch: 2 | Step: 4950 | Loss: 2.8816\n",
      "Epoch: 2 | Step: 5000 | Loss: 2.9333\n",
      "Epoch: 2 | Step: 5050 | Loss: 2.9712\n",
      "Epoch: 2 | Step: 5100 | Loss: 2.8881\n",
      "Epoch: 2 | Step: 5150 | Loss: 2.7893\n",
      "Epoch: 2 | Step: 5200 | Loss: 2.9218\n",
      "Epoch: 2 | Step: 5250 | Loss: 2.8605\n",
      "Epoch: 2 | Step: 5300 | Loss: 2.8556\n",
      "Epoch: 2 | Step: 5350 | Loss: 2.8889\n",
      "Epoch: 2 | Step: 5400 | Loss: 2.8702\n",
      "Epoch: 2 | Step: 5450 | Loss: 2.8950\n",
      "Epoch: 2 | Step: 5500 | Loss: 2.9168\n",
      "Epoch: 2 | Step: 5550 | Loss: 2.8748\n",
      "Epoch: 2 | Step: 5600 | Loss: 2.6852\n",
      "Epoch: 2 | Step: 5650 | Loss: 2.7915\n",
      "Epoch: 2 | Step: 5700 | Loss: 2.9245\n",
      "Epoch: 2 | Step: 5750 | Loss: 2.8624\n",
      "Epoch: 2 | Step: 5800 | Loss: 2.6552\n",
      "Epoch: 2 | Step: 5850 | Loss: 2.8393\n",
      "Epoch: 2 | Step: 5900 | Loss: 2.7675\n",
      "Epoch: 2 | Step: 5950 | Loss: 2.9464\n",
      "Epoch: 2 | Step: 6000 | Loss: 2.6946\n",
      "Epoch: 2 | Step: 6050 | Loss: 2.6613\n",
      "Epoch: 2 | Step: 6100 | Loss: 2.7430\n",
      "Epoch: 2 | Step: 6150 | Loss: 2.6010\n",
      "Epoch: 2 | Step: 6200 | Loss: 2.8388\n",
      "Epoch: 3 | Step: 6250 | Loss: 2.6970\n",
      "Epoch: 3 | Step: 6300 | Loss: 2.7332\n",
      "Epoch: 3 | Step: 6350 | Loss: 2.7604\n",
      "Epoch: 3 | Step: 6400 | Loss: 2.5661\n",
      "Epoch: 3 | Step: 6450 | Loss: 2.7858\n",
      "Epoch: 3 | Step: 6500 | Loss: 2.6917\n",
      "Epoch: 3 | Step: 6550 | Loss: 2.7398\n",
      "Epoch: 3 | Step: 6600 | Loss: 2.8848\n",
      "Epoch: 3 | Step: 6650 | Loss: 2.7489\n",
      "Epoch: 3 | Step: 6700 | Loss: 2.8082\n",
      "Epoch: 3 | Step: 6750 | Loss: 2.8065\n",
      "Epoch: 3 | Step: 6800 | Loss: 2.6998\n",
      "Epoch: 3 | Step: 6850 | Loss: 2.7331\n",
      "Epoch: 3 | Step: 6900 | Loss: 2.7613\n",
      "Epoch: 3 | Step: 6950 | Loss: 2.6102\n",
      "Epoch: 3 | Step: 7000 | Loss: 2.7100\n",
      "Epoch: 3 | Step: 7050 | Loss: 2.7291\n",
      "Epoch: 3 | Step: 7100 | Loss: 2.5198\n",
      "Epoch: 3 | Step: 7150 | Loss: 2.7852\n",
      "Epoch: 3 | Step: 7200 | Loss: 2.8854\n",
      "Epoch: 3 | Step: 7250 | Loss: 2.6228\n",
      "Epoch: 3 | Step: 7300 | Loss: 2.6152\n",
      "Epoch: 3 | Step: 7350 | Loss: 2.6976\n",
      "Epoch: 3 | Step: 7400 | Loss: 2.7868\n",
      "Epoch: 3 | Step: 7450 | Loss: 2.7638\n",
      "Epoch: 3 | Step: 7500 | Loss: 2.7816\n",
      "Epoch: 3 | Step: 7550 | Loss: 2.7030\n",
      "Epoch: 3 | Step: 7600 | Loss: 2.7045\n",
      "Epoch: 3 | Step: 7650 | Loss: 2.6442\n",
      "Epoch: 3 | Step: 7700 | Loss: 2.7383\n",
      "Epoch: 3 | Step: 7750 | Loss: 2.5767\n",
      "Epoch: 3 | Step: 7800 | Loss: 2.8687\n",
      "Epoch: 3 | Step: 7850 | Loss: 2.6407\n",
      "Epoch: 3 | Step: 7900 | Loss: 2.6357\n",
      "Epoch: 3 | Step: 7950 | Loss: 2.5778\n",
      "Epoch: 3 | Step: 8000 | Loss: 2.5556\n",
      "Epoch: 3 | Step: 8050 | Loss: 2.6383\n",
      "Epoch: 3 | Step: 8100 | Loss: 2.6702\n",
      "Epoch: 3 | Step: 8150 | Loss: 2.4866\n",
      "Epoch: 3 | Step: 8200 | Loss: 2.6706\n",
      "Epoch: 3 | Step: 8250 | Loss: 2.6504\n",
      "Epoch: 3 | Step: 8300 | Loss: 2.6218\n",
      "Epoch: 3 | Step: 8350 | Loss: 2.7245\n",
      "Epoch: 3 | Step: 8400 | Loss: 2.5353\n",
      "Epoch: 3 | Step: 8450 | Loss: 2.6789\n",
      "Epoch: 3 | Step: 8500 | Loss: 2.7370\n",
      "Epoch: 3 | Step: 8550 | Loss: 2.7378\n",
      "Epoch: 3 | Step: 8600 | Loss: 2.4548\n",
      "Epoch: 3 | Step: 8650 | Loss: 2.5813\n",
      "Epoch: 3 | Step: 8700 | Loss: 2.6413\n",
      "Epoch: 3 | Step: 8750 | Loss: 2.5437\n",
      "Epoch: 3 | Step: 8800 | Loss: 2.5465\n",
      "Epoch: 3 | Step: 8850 | Loss: 2.6350\n",
      "Epoch: 3 | Step: 8900 | Loss: 2.5841\n",
      "Epoch: 3 | Step: 8950 | Loss: 2.5196\n",
      "Epoch: 3 | Step: 9000 | Loss: 2.5717\n",
      "Epoch: 3 | Step: 9050 | Loss: 2.4508\n",
      "Epoch: 3 | Step: 9100 | Loss: 2.4996\n",
      "Epoch: 3 | Step: 9150 | Loss: 2.5525\n",
      "Epoch: 3 | Step: 9200 | Loss: 2.6256\n",
      "Epoch: 3 | Step: 9250 | Loss: 2.4567\n",
      "Epoch: 3 | Step: 9300 | Loss: 2.5538\n",
      "Epoch: 3 | Step: 9350 | Loss: 2.6922\n",
      "Epoch: 4 | Step: 9400 | Loss: 2.4322\n",
      "Epoch: 4 | Step: 9450 | Loss: 2.4887\n",
      "Epoch: 4 | Step: 9500 | Loss: 2.4750\n",
      "Epoch: 4 | Step: 9550 | Loss: 2.4780\n",
      "Epoch: 4 | Step: 9600 | Loss: 2.5154\n",
      "Epoch: 4 | Step: 9650 | Loss: 2.3938\n",
      "Epoch: 4 | Step: 9700 | Loss: 2.5977\n",
      "Epoch: 4 | Step: 9750 | Loss: 2.5060\n",
      "Epoch: 4 | Step: 9800 | Loss: 2.4144\n",
      "Epoch: 4 | Step: 9850 | Loss: 2.5329\n",
      "Epoch: 4 | Step: 9900 | Loss: 2.6268\n",
      "Epoch: 4 | Step: 9950 | Loss: 2.6731\n",
      "Epoch: 4 | Step: 10000 | Loss: 2.4936\n",
      "Epoch: 4 | Step: 10050 | Loss: 2.4888\n",
      "Epoch: 4 | Step: 10100 | Loss: 2.5665\n",
      "Epoch: 4 | Step: 10150 | Loss: 2.6172\n",
      "Epoch: 4 | Step: 10200 | Loss: 2.4674\n",
      "Epoch: 4 | Step: 10250 | Loss: 2.5077\n",
      "Epoch: 4 | Step: 10300 | Loss: 2.6004\n",
      "Epoch: 4 | Step: 10350 | Loss: 2.6375\n",
      "Epoch: 4 | Step: 10400 | Loss: 2.4944\n",
      "Epoch: 4 | Step: 10450 | Loss: 2.3994\n",
      "Epoch: 4 | Step: 10500 | Loss: 2.5516\n",
      "Epoch: 4 | Step: 10550 | Loss: 2.5688\n",
      "Epoch: 4 | Step: 10600 | Loss: 2.7096\n",
      "Epoch: 4 | Step: 10650 | Loss: 2.5643\n",
      "Epoch: 4 | Step: 10700 | Loss: 2.4758\n",
      "Epoch: 4 | Step: 10750 | Loss: 2.6422\n",
      "Epoch: 4 | Step: 10800 | Loss: 2.4667\n",
      "Epoch: 4 | Step: 10850 | Loss: 2.4922\n",
      "Epoch: 4 | Step: 10900 | Loss: 2.5602\n",
      "Epoch: 4 | Step: 10950 | Loss: 2.5130\n",
      "Epoch: 4 | Step: 11000 | Loss: 2.2740\n",
      "Epoch: 4 | Step: 11050 | Loss: 2.4862\n",
      "Epoch: 4 | Step: 11100 | Loss: 2.5169\n",
      "Epoch: 4 | Step: 11150 | Loss: 2.4537\n",
      "Epoch: 4 | Step: 11200 | Loss: 2.4462\n",
      "Epoch: 4 | Step: 11250 | Loss: 2.5049\n",
      "Epoch: 4 | Step: 11300 | Loss: 2.5403\n",
      "Epoch: 4 | Step: 11350 | Loss: 2.4348\n",
      "Epoch: 4 | Step: 11400 | Loss: 2.4458\n",
      "Epoch: 4 | Step: 11450 | Loss: 2.4711\n",
      "Epoch: 4 | Step: 11500 | Loss: 2.5001\n",
      "Epoch: 4 | Step: 11550 | Loss: 2.5460\n",
      "Epoch: 4 | Step: 11600 | Loss: 2.4432\n",
      "Epoch: 4 | Step: 11650 | Loss: 2.3907\n",
      "Epoch: 4 | Step: 11700 | Loss: 2.3861\n",
      "Epoch: 4 | Step: 11750 | Loss: 2.5116\n",
      "Epoch: 4 | Step: 11800 | Loss: 2.3755\n",
      "Epoch: 4 | Step: 11850 | Loss: 2.3396\n",
      "Epoch: 4 | Step: 11900 | Loss: 2.5026\n",
      "Epoch: 4 | Step: 11950 | Loss: 2.4337\n",
      "Epoch: 4 | Step: 12000 | Loss: 2.4622\n",
      "Epoch: 4 | Step: 12050 | Loss: 2.3881\n",
      "Epoch: 4 | Step: 12100 | Loss: 2.5008\n",
      "Epoch: 4 | Step: 12150 | Loss: 2.4145\n",
      "Epoch: 4 | Step: 12200 | Loss: 2.4365\n",
      "Epoch: 4 | Step: 12250 | Loss: 2.3167\n",
      "Epoch: 4 | Step: 12300 | Loss: 2.3190\n",
      "Epoch: 4 | Step: 12350 | Loss: 2.3477\n",
      "Epoch: 4 | Step: 12400 | Loss: 2.4438\n",
      "Epoch: 4 | Step: 12450 | Loss: 2.5531\n",
      "Epoch: 5 | Step: 12500 | Loss: 2.2985\n",
      "Epoch: 5 | Step: 12550 | Loss: 2.3663\n",
      "Epoch: 5 | Step: 12600 | Loss: 2.4244\n",
      "Epoch: 5 | Step: 12650 | Loss: 2.4081\n",
      "Epoch: 5 | Step: 12700 | Loss: 2.5068\n",
      "Epoch: 5 | Step: 12750 | Loss: 2.5635\n",
      "Epoch: 5 | Step: 12800 | Loss: 2.2140\n",
      "Epoch: 5 | Step: 12850 | Loss: 2.3762\n",
      "Epoch: 5 | Step: 12900 | Loss: 2.3729\n",
      "Epoch: 5 | Step: 12950 | Loss: 2.3644\n",
      "Epoch: 5 | Step: 13000 | Loss: 2.4486\n",
      "Epoch: 5 | Step: 13050 | Loss: 2.3588\n",
      "Epoch: 5 | Step: 13100 | Loss: 2.3154\n",
      "Epoch: 5 | Step: 13150 | Loss: 2.3287\n",
      "Epoch: 5 | Step: 13200 | Loss: 2.3201\n",
      "Epoch: 5 | Step: 13250 | Loss: 2.3586\n",
      "Epoch: 5 | Step: 13300 | Loss: 2.2816\n",
      "Epoch: 5 | Step: 13350 | Loss: 2.3568\n",
      "Epoch: 5 | Step: 13400 | Loss: 2.2897\n",
      "Epoch: 5 | Step: 13450 | Loss: 2.3943\n",
      "Epoch: 5 | Step: 13500 | Loss: 2.2016\n",
      "Epoch: 5 | Step: 13550 | Loss: 2.3734\n",
      "Epoch: 5 | Step: 13600 | Loss: 2.3038\n",
      "Epoch: 5 | Step: 13650 | Loss: 2.2630\n",
      "Epoch: 5 | Step: 13700 | Loss: 2.2993\n",
      "Epoch: 5 | Step: 13750 | Loss: 2.4860\n",
      "Epoch: 5 | Step: 13800 | Loss: 2.4693\n",
      "Epoch: 5 | Step: 13850 | Loss: 2.3016\n",
      "Epoch: 5 | Step: 13900 | Loss: 2.2769\n",
      "Epoch: 5 | Step: 13950 | Loss: 2.3816\n",
      "Epoch: 5 | Step: 14000 | Loss: 2.3813\n",
      "Epoch: 5 | Step: 14050 | Loss: 2.3228\n",
      "Epoch: 5 | Step: 14100 | Loss: 2.3325\n",
      "Epoch: 5 | Step: 14150 | Loss: 2.3579\n",
      "Epoch: 5 | Step: 14200 | Loss: 2.3148\n",
      "Epoch: 5 | Step: 14250 | Loss: 2.2643\n",
      "Epoch: 5 | Step: 14300 | Loss: 2.3670\n",
      "Epoch: 5 | Step: 14350 | Loss: 2.2539\n",
      "Epoch: 5 | Step: 14400 | Loss: 2.3054\n",
      "Epoch: 5 | Step: 14450 | Loss: 2.3188\n",
      "Epoch: 5 | Step: 14500 | Loss: 2.4076\n",
      "Epoch: 5 | Step: 14550 | Loss: 2.3530\n",
      "Epoch: 5 | Step: 14600 | Loss: 2.3665\n",
      "Epoch: 5 | Step: 14650 | Loss: 2.4980\n",
      "Epoch: 5 | Step: 14700 | Loss: 2.1730\n",
      "Epoch: 5 | Step: 14750 | Loss: 2.4963\n",
      "Epoch: 5 | Step: 14800 | Loss: 2.4000\n",
      "Epoch: 5 | Step: 14850 | Loss: 2.2610\n",
      "Epoch: 5 | Step: 14900 | Loss: 2.3533\n",
      "Epoch: 5 | Step: 14950 | Loss: 2.2011\n",
      "Epoch: 5 | Step: 15000 | Loss: 2.4029\n",
      "Epoch: 5 | Step: 15050 | Loss: 2.2902\n",
      "Epoch: 5 | Step: 15100 | Loss: 2.4174\n",
      "Epoch: 5 | Step: 15150 | Loss: 2.3255\n",
      "Epoch: 5 | Step: 15200 | Loss: 2.3692\n",
      "Epoch: 5 | Step: 15250 | Loss: 2.3830\n",
      "Epoch: 5 | Step: 15300 | Loss: 2.4306\n",
      "Epoch: 5 | Step: 15350 | Loss: 2.3075\n",
      "Epoch: 5 | Step: 15400 | Loss: 2.3117\n",
      "Epoch: 5 | Step: 15450 | Loss: 2.4330\n",
      "Epoch: 5 | Step: 15500 | Loss: 2.3421\n",
      "Epoch: 5 | Step: 15550 | Loss: 2.4023\n",
      "Epoch: 5 | Step: 15600 | Loss: 2.3639\n"
     ]
    }
   ],
   "source": [
    "# 1. Initialize Tokenizer and Model\n",
    "tokenizer = tiktoken.get_encoding(\"gpt2\")\n",
    "model = NanoLLM(\n",
    "    maxlen=128, \n",
    "    vocab_size=tokenizer.n_vocab, \n",
    "    embed_dim=192, \n",
    "    num_heads=6, \n",
    "    feed_forward_dim=512, \n",
    "    num_transformer_blocks=6\n",
    ")\n",
    "\n",
    "# 2. Setup Optimizer (MLX equivalent of optax.adamw)\n",
    "# Note: For simplicity, we are using a static learning rate here, \n",
    "# but MLX supports lr schedules via optim.step_decay etc.\n",
    "optimizer = optim.AdamW(learning_rate=3e-4)\n",
    "\n",
    "# 3. Define the Loss Function\n",
    "def loss_fn(model, x, y):\n",
    "    logits = model(x)\n",
    "    # MLX has a built-in cross entropy loss that expects (logits, targets)\n",
    "    loss = nn.losses.cross_entropy(logits, y)\n",
    "    return mx.mean(loss)\n",
    "\n",
    "# 4. Compile the Value and Grad function\n",
    "# This is the exact MLX equivalent to nnx.value_and_grad\n",
    "loss_and_grad_fn = nn.value_and_grad(model, loss_fn)\n",
    "\n",
    "# 5. Run the Loop\n",
    "epochs = 5\n",
    "step = 0\n",
    "\n",
    "for epoch in range(epochs):\n",
    "    # Initialize our generator\n",
    "    dataloader = get_batches(\"TinyStories-100k.txt\", tokenizer, batch_size=32, maxlen=128)\n",
    "    \n",
    "    for x, y in dataloader:\n",
    "        # Calculate loss and gradients\n",
    "        loss, grads = loss_and_grad_fn(model, x, y)\n",
    "        \n",
    "        # Apply gradients\n",
    "        optimizer.update(model, grads)\n",
    "        \n",
    "        # CRITICAL: Force MLX to actually compute the updates on the Apple GPU\n",
    "        mx.eval(model.parameters(), optimizer.state, loss)\n",
    "        \n",
    "        if step % 50 == 0:\n",
    "            print(f\"Epoch: {epoch + 1} | Step: {step} | Loss: {loss.item():.4f}\")\n",
    "            \n",
    "        step += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "e1d66a2f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model saved successfully!\n"
     ]
    }
   ],
   "source": [
    "model.save_weights(\"small_checkpoint.safetensors\")\n",
    "\n",
    "print(\"Model saved successfully!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3bdbdc57",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "There lived kid named Tom and was a very good day. He had a big box of toys and he had a lot of toys. He was very happy and he wanted to play with his toys.\n",
      "\n",
      "One day, Tom was playing in the park. He saw a big tree with a big tree. He wanted to climb it, but he was too small. He tried to climb the tree, but he could not reach it.\n",
      "\n",
      "Tom was sad. He wanted to help, but he was not sure. He tried to climb the tree, but he was too high. He could not reach the tree. He was sad.\n",
      "\n",
      "Spot was sad. He tried to make his friends. He says he could not give him to the other animals.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "def generate_text(model, tokenizer, prompt, max_new_tokens=50, temperature=1.0):\n",
    "    # Encode the prompt and convert to an MLX array with a batch dimension\n",
    "    tokens = tokenizer.encode(prompt)\n",
    "    x = mx.array(tokens)[None, :] # Shape: (1, seq_len)\n",
    "    \n",
    "    end_token_id = tokenizer.encode('<|endoftext|>', allowed_special={'<|endoftext|>'})[0]\n",
    "    \n",
    "    print(prompt, end=\"\", flush=True)\n",
    "    \n",
    "    for _ in range(max_new_tokens):\n",
    "        # Truncate context if it exceeds the model's max length\n",
    "        if x.shape[1] > model.maxlen:\n",
    "            x = x[:, -model.maxlen:]\n",
    "            \n",
    "        # Forward pass\n",
    "        logits = model(x)\n",
    "        \n",
    "        # Grab the logits for the very last token in the sequence\n",
    "        next_token_logits = logits[0, -1, :] / temperature\n",
    "        \n",
    "        # Get the highest probability token (greedy decoding)\n",
    "        # Note: You can replace mx.argmax with mx.random.categorical for more creative outputs\n",
    "        next_token = mx.argmax(next_token_logits, axis=-1).item()\n",
    "        \n",
    "        if next_token == end_token_id:\n",
    "            break\n",
    "            \n",
    "        # Decode and print the single new token immediately\n",
    "        word = tokenizer.decode([next_token])\n",
    "        print(word, end=\"\", flush=True)\n",
    "        \n",
    "        # Append the new token to our input sequence for the next loop\n",
    "        x = mx.concatenate([x, mx.array([[next_token]])], axis=1)\n",
    "        \n",
    "    print(\"\\n\")\n",
    "\n",
    "# Run inference\n",
    "generate_text(model, tokenizer, \"The king of jungle\", max_new_tokens=1024, temperature=0.4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "e03eadd6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Weights loaded and ready for generation!\n"
     ]
    }
   ],
   "source": [
    "# 1. Initialize the empty architecture\n",
    "model = NanoLLM(\n",
    "    maxlen=128, \n",
    "    vocab_size=tokenizer.n_vocab, \n",
    "    embed_dim=192, \n",
    "    num_heads=6, \n",
    "    feed_forward_dim=512, \n",
    "    num_transformer_blocks=6\n",
    ")\n",
    "\n",
    "# 2. Load the trained weights directly into the layers\n",
    "model.load_weights(\"nanollm_tinystories.safetensors\")\n",
    "\n",
    "print(\"Weights loaded and ready for generation!\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}