gpt_dev.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wJpXpmjEYC_T"
   },
   "source": [
    "## Building a GPT\n",
    "\n",
    "Companion notebook to the [Zero To Hero](https://karpathy.ai/zero-to-hero.html) video on GPT."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "h5hjCcLDr2WC",
    "outputId": "ccc60f0c-fd78-4dbe-8598-0512d1036aad"
   },
   "outputs": [],
   "source": [
    "# We always start with a dataset to train on. Let's download the tiny shakespeare dataset\n",
    "#!wget https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt\n",
    "\n",
    "# I do this manually in windows via wsl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "print(os.getcwd())  # prints the current directory\n",
    "\n",
    "#%run \"../../datasets/kaggle/dataload.ipynb\"\n",
    "\n",
    "#df = load_data(\"C:/Users/Dasun/Data/NLP/datasets/kaggle/converted_data.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "O6medjfRsLD9"
   },
   "outputs": [],
   "source": [
    "# read it in to inspect it\n",
    "with open('input.txt', 'r', encoding='utf-8') as f:\n",
    "    text = f.read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "6xWI_VyAsN8F",
    "outputId": "ed819dd0-72e5-40a6-d2ed-928ff73bfda6"
   },
   "outputs": [],
   "source": [
    "print(\"length of dataset in characters: \", len(text))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "2c5V0FvqseE0",
    "outputId": "25ca7adc-b8c0-42d1-b08c-e0863c5c314e"
   },
   "outputs": [],
   "source": [
    "# let's look at the first 1000 characters\n",
    "print(text[:10],'...',text[-11:-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "0e-Rbyr8sfM8",
    "outputId": "f34e94a9-5b44-4cf3-885b-986731929109"
   },
   "outputs": [],
   "source": [
    "# here are all the unique characters that occur in this text\n",
    "chars = sorted(list(set(text)))\n",
    "vocab_size = len(chars)\n",
    "print(''.join(chars))\n",
    "print(vocab_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Yw1LKNCgwjj1",
    "outputId": "86fcc21c-2cf7-40d9-cd7b-b5a253da4459"
   },
   "outputs": [],
   "source": [
    "# create a mapping from characters to integers\n",
    "stoi = { ch:i for i,ch in enumerate(chars) }\n",
    "itos = { i:ch for i,ch in enumerate(chars) }\n",
    "encode = lambda s: [stoi[c] for c in s] # encoder: take a string, output a list of integers\n",
    "decode = lambda l: ''.join([itos[i] for i in l]) # decoder: take a list of integers, output a string\n",
    "\n",
    "print(encode(\"hii there\"))\n",
    "print(decode(encode(\"hii there\")))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "YJb0OXPwzvqg",
    "outputId": "db7297cc-36a9-4fae-e941-e7bb9e0e91d1"
   },
   "outputs": [],
   "source": [
    "# let's now encode the entire text dataset and store it into a torch.Tensor\n",
    "import torch # we use PyTorch: https://pytorch.org\n",
    "data = torch.tensor(encode(text), dtype=torch.long)\n",
    "print(data.shape, data.dtype) \n",
    "# the 20 characters we looked at earier will to the GPT look like this\n",
    "print((data[:10]),(data[-11:-1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "f_WIXqxz0lU5"
   },
   "outputs": [],
   "source": [
    "# Let's now split up the data into train and validation sets\n",
    "n = int(0.9*len(data)) # first 90% will be train, rest val\n",
    "train_data = data[:n]\n",
    "val_data = data[n:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "TD5Bj8Y6IAD4",
    "outputId": "bf23c586-1d33-4af1-b63d-ce6f90b0a528"
   },
   "outputs": [],
   "source": [
    "block_size = 8\n",
    "train_data[:block_size+1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "9HXDe8vGJCEn",
    "outputId": "588663aa-1de5-4ef7-aba0-4a96fe828353"
   },
   "outputs": [],
   "source": [
    "x = train_data[:block_size]   # Shift it to \n",
    "y = train_data[1:block_size+1]\n",
    "for t in range(block_size):\n",
    "    context = x[:t+1]\n",
    "    target = y[t]\n",
    "    print(f\"when input is {context} the target: {target}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Q3k1Czf7LuA9",
    "outputId": "4ea8e8a0-443c-49bb-b3bf-ba36e1712999"
   },
   "outputs": [],
   "source": [
    "torch.manual_seed(1337)\n",
    "batch_size = 4 # how many independent sequences will we process in parallel?\n",
    "block_size = 8 # what is the maximum context length for predictions?\n",
    "# I might need to change the CLength since we have two sentences\n",
    "\n",
    "def get_batch(split):\n",
    "    # generate a small batch of data of inputs x and targets y\n",
    "    data = train_data if split == 'train' else val_data\n",
    "    ix = torch.randint(len(data) - block_size, (batch_size,))\n",
    "    x = torch.stack([data[i:i+block_size] for i in ix])\n",
    "    y = torch.stack([data[i+1:i+block_size+1] for i in ix])\n",
    "    return x, y\n",
    "\n",
    "xb, yb = get_batch('train')\n",
    "print('inputs:')\n",
    "print(xb.shape)\n",
    "print(xb)\n",
    "print('targets:')\n",
    "print(yb.shape)\n",
    "print(yb)\n",
    "\n",
    "print('----')\n",
    "\n",
    "for b in range(batch_size): # batch dimension\n",
    "    for t in range(block_size): # time dimension\n",
    "        context = xb[b, :t+1]\n",
    "        target = yb[b,t]\n",
    "        print(f\"when input is {context.tolist()} the target: {target}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "qpyyAeIzQjlO",
    "outputId": "a650f8dc-da81-400b-bc59-0a595487fdb9"
   },
   "outputs": [],
   "source": [
    "print(xb) # our input to the transformer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "nql_1ER53oCf",
    "outputId": "5de90b1b-4603-428a-f571-fe4bd3c45436"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import functional as F\n",
    "torch.manual_seed(1337)\n",
    "\n",
    "class BigramLanguageModel(nn.Module):\n",
    "\n",
    "    def __init__(self, vocab_size):\n",
    "        super().__init__()\n",
    "        # each token directly reads off the logits for the next token from a lookup table\n",
    "        self.token_embedding_table = nn.Embedding(vocab_size, vocab_size)\n",
    "\n",
    "    def forward(self, idx, targets=None):\n",
    "\n",
    "        # idx and targets are both (B,T) tensor of integers\n",
    "        logits = self.token_embedding_table(idx) # (B,T,C)\n",
    "\n",
    "        if targets is None:\n",
    "            loss = None\n",
    "        else:\n",
    "            B, T, C = logits.shape\n",
    "            logits = logits.view(B*T, C)\n",
    "            targets = targets.view(B*T)\n",
    "            loss = F.cross_entropy(logits, targets)\n",
    "\n",
    "        return logits, loss\n",
    "\n",
    "    def generate(self, idx, max_new_tokens, temperature=1.0):\n",
    "        # idx is (B, T) array of indices in the current context\n",
    "        for _ in range(max_new_tokens):\n",
    "            # get the predictions\n",
    "            logits, loss = self(idx)\n",
    "            # focus only on the last time step\n",
    "            logits = logits[:, -1, :] # becomes (B, C)\n",
    "            # apply temperature\n",
    "            logits = logits/temperature\n",
    "            # apply softmax to get probabilities\n",
    "            probs = F.softmax(logits, dim=-1) # (B, C)\n",
    "            # sample from the distribution\n",
    "            idx_next = torch.multinomial(probs, num_samples=1) # (B, 1)\n",
    "            # append sampled index to the running sequence\n",
    "            idx = torch.cat((idx, idx_next), dim=1) # (B, T+1)\n",
    "        return idx\n",
    "\n",
    "m = BigramLanguageModel(vocab_size)\n",
    "logits, loss = m(xb, yb)\n",
    "print(logits.shape)\n",
    "print(loss)\n",
    "\n",
    "print(decode(m.generate(idx = torch.zeros((1, 1), dtype=torch.long), max_new_tokens=100)[0].tolist()))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "eTyJ8qAaDdiF"
   },
   "outputs": [],
   "source": [
    "# create a PyTorch optimizer\n",
    "optimizer = torch.optim.AdamW(m.parameters(), lr=1e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Hs4kI8YdEkQj",
    "outputId": "42ded55c-2983-4d91-c528-675b2edfa849"
   },
   "outputs": [],
   "source": [
    "batch_size = 32\n",
    "for steps in range(1000): # increase number of steps for good results...\n",
    "\n",
    "    # sample a batch of data\n",
    "    xb, yb = get_batch('train')\n",
    "\n",
    "    # evaluate the loss\n",
    "    logits, loss = m(xb, yb)\n",
    "    optimizer.zero_grad(set_to_none=True)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "print(loss.item())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "EcVIDWAZEtjN",
    "outputId": "0ad6f9d2-ad58-4498-a5f8-6f31407bb18b"
   },
   "outputs": [],
   "source": [
    "print(decode(m.generate(idx = torch.zeros((1, 1), dtype=torch.long), max_new_tokens=500)[0].tolist()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XinV8nmAnmKN"
   },
   "source": [
    "## The mathematical trick in self-attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "tukiH-NbRBhA",
    "outputId": "d981f6d4-ac08-4ec2-8284-82f5fa1e0815"
   },
   "outputs": [],
   "source": [
    "# toy example illustrating how matrix multiplication can be used for a \"weighted aggregation\"\n",
    "torch.manual_seed(42)\n",
    "a = torch.tril(torch.ones(3, 3))\n",
    "a = a / torch.sum(a, 1, keepdim=True)\n",
    "b = torch.randint(0,10,(3,2)).float()\n",
    "c = a @ b\n",
    "print('a=')\n",
    "print(a)\n",
    "print('--')\n",
    "print('b=')\n",
    "print(b)\n",
    "print('--')\n",
    "print('c=')\n",
    "print(c)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Hs_E24uRE8kr",
    "outputId": "8bf3ff5f-565e-48b8-de8e-7272706c8e12"
   },
   "outputs": [],
   "source": [
    "# consider the following toy example:\n",
    "\n",
    "torch.manual_seed(1337)\n",
    "B,T,C = 4,8,2 # batch, time, channels\n",
    "x = torch.randn(B,T,C)\n",
    "x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "86NuXX0fn7ps"
   },
   "outputs": [],
   "source": [
    "# We want x[b,t] = mean_{i<=t} x[b,i]\n",
    "xbow = torch.zeros((B,T,C))\n",
    "for b in range(B):\n",
    "    for t in range(T):\n",
    "        xprev = x[b,:t+1] # (t,C)\n",
    "        xbow[b,t] = torch.mean(xprev, 0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "yhdOAd6-wXkZ",
    "outputId": "eaf6ab61-dff1-4bb7-e623-47f692bad5f9"
   },
   "outputs": [],
   "source": [
    "# version 2: using matrix multiply for a weighted aggregation\n",
    "wei = torch.tril(torch.ones(T, T))\n",
    "wei = wei / wei.sum(1, keepdim=True)\n",
    "xbow2 = wei @ x # (B, T, T) @ (B, T, C) ----> (B, T, C)\n",
    "torch.allclose(xbow, xbow2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "wOURrfG-ysoL",
    "outputId": "080b500d-8110-4602-fcef-7d6f2ebfc6bc"
   },
   "outputs": [],
   "source": [
    "# version 3: use Softmax\n",
    "tril = torch.tril(torch.ones(T, T))\n",
    "wei = torch.zeros((T,T))\n",
    "wei = wei.masked_fill(tril == 0, float('-inf'))\n",
    "wei = F.softmax(wei, dim=-1)\n",
    "xbow3 = wei @ x\n",
    "torch.allclose(xbow, xbow3)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "EDarxEWIRMKq",
    "outputId": "07b587dd-a91c-4bb0-d7f1-e247cd5dacb5"
   },
   "outputs": [],
   "source": [
    "# version 4: self-attention!\n",
    "torch.manual_seed(1337)\n",
    "B,T,C = 4,8,32 # batch, time, channels\n",
    "x = torch.randn(B,T,C)\n",
    "\n",
    "# let's see a single Head perform self-attention\n",
    "head_size = 16\n",
    "key = nn.Linear(C, head_size, bias=False)\n",
    "query = nn.Linear(C, head_size, bias=False)\n",
    "value = nn.Linear(C, head_size, bias=False)\n",
    "k = key(x)   # (B, T, 16)\n",
    "q = query(x) # (B, T, 16)\n",
    "wei =  q @ k.transpose(-2, -1) # (B, T, 16) @ (B, 16, T) ---> (B, T, T)\n",
    "\n",
    "tril = torch.tril(torch.ones(T, T))\n",
    "#wei = torch.zeros((T,T))\n",
    "wei = wei.masked_fill(tril == 0, float('-inf'))\n",
    "wei = F.softmax(wei, dim=-1)\n",
    "\n",
    "v = value(x)\n",
    "out = wei @ v\n",
    "#out = wei @ x\n",
    "\n",
    "out.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "vT1hdtzXCjgL",
    "outputId": "6d2c569b-7922-451f-9934-0fc564678d17"
   },
   "outputs": [],
   "source": [
    "wei[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "M5CvobiQ0pLr"
   },
   "source": [
    "Notes:\n",
    "- Attention is a **communication mechanism**. Can be seen as nodes in a directed graph looking at each other and aggregating information with a weighted sum from all nodes that point to them, with data-dependent weights.\n",
    "- There is no notion of space. Attention simply acts over a set of vectors. This is why we need to positionally encode tokens.\n",
    "- Each example across batch dimension is of course processed completely independently and never \"talk\" to each other\n",
    "- In an \"encoder\" attention block just delete the single line that does masking with `tril`, allowing all tokens to communicate. This block here is called a \"decoder\" attention block because it has triangular masking, and is usually used in autoregressive settings, like language modeling.\n",
    "- \"self-attention\" just means that the keys and values are produced from the same source as queries. In \"cross-attention\", the queries still get produced from x, but the keys and values come from some other, external source (e.g. an encoder module)\n",
    "- \"Scaled\" attention additional divides `wei` by 1/sqrt(head_size). This makes it so when input Q,K are unit variance, wei will be unit variance too and Softmax will stay diffuse and not saturate too much. Illustration below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "4SNbLq5z3oBw"
   },
   "outputs": [],
   "source": [
    "k = torch.randn(B,T,head_size)\n",
    "q = torch.randn(B,T,head_size)\n",
    "wei = q @ k.transpose(-2, -1) * head_size**-0.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Nl6I9n9IRTSo",
    "outputId": "0c5b9cd0-af8a-4564-fbad-41d844e54822"
   },
   "outputs": [],
   "source": [
    "k.var()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "T1tQx7oeRvtc",
    "outputId": "3541ca1a-7447-4ef7-835e-81824aebc1b5"
   },
   "outputs": [],
   "source": [
    "q.var()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "MLb_odHU3iKM",
    "outputId": "a687a222-5a2c-4cdb-c1bf-17cd05b45b69"
   },
   "outputs": [],
   "source": [
    "wei.var()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "JB82yzt44REI",
    "outputId": "f07da2f1-10bb-4a7a-bcaa-578587977d00"
   },
   "outputs": [],
   "source": [
    "torch.softmax(torch.tensor([0.1, -0.2, 0.3, -0.2, 0.5]), dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Mpt8569BB9_f",
    "outputId": "5d8b910a-6192-44ba-ebb2-497d88e0b629"
   },
   "outputs": [],
   "source": [
    "torch.softmax(torch.tensor([0.1, -0.2, 0.3, -0.2, 0.5])*8, dim=-1) # gets too peaky, converges to one-hot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "2Num7sX9CKOH",
    "outputId": "929ceb78-a639-41d6-aac7-12997b5c93f0"
   },
   "outputs": [],
   "source": [
    "class LayerNorm1d: # (used to be BatchNorm1d)\n",
    "\n",
    "  def __init__(self, dim, eps=1e-5, momentum=0.1):\n",
    "    self.eps = eps\n",
    "    self.gamma = torch.ones(dim)\n",
    "    self.beta = torch.zeros(dim)\n",
    "\n",
    "  def __call__(self, x):\n",
    "    # calculate the forward pass\n",
    "    xmean = x.mean(1, keepdim=True) # batch mean\n",
    "    xvar = x.var(1, keepdim=True) # batch variance\n",
    "    xhat = (x - xmean) / torch.sqrt(xvar + self.eps) # normalize to unit variance\n",
    "    self.out = self.gamma * xhat + self.beta\n",
    "    return self.out\n",
    "\n",
    "  def parameters(self):\n",
    "    return [self.gamma, self.beta]\n",
    "\n",
    "torch.manual_seed(1337)\n",
    "module = LayerNorm1d(100)\n",
    "x = torch.randn(32, 100) # batch size 32 of 100-dimensional vectors\n",
    "x = module(x)\n",
    "x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "633T2cmnW1uk",
    "outputId": "7720fa58-0478-4e8a-86a7-502d4cce9443"
   },
   "outputs": [],
   "source": [
    "x[:,0].mean(), x[:,0].std() # mean,std of one feature across all batch inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "LN9cK9BoXCYb",
    "outputId": "6368ece0-600e-417d-8a91-7c1e5d750ba8"
   },
   "outputs": [],
   "source": [
    "x[0,:].mean(), x[0,:].std() # mean,std of a single input from the batch, of its features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "dRJH6wM_XFfU"
   },
   "outputs": [],
   "source": [
    "# French to English translation example:\n",
    "\n",
    "# <--------- ENCODE ------------------><--------------- DECODE ----------------->\n",
    "# les réseaux de neurones sont géniaux! <START> neural networks are awesome!<END>\n",
    "\n",
    "\"\"\"Yes, you can certainly retrain your model to perform translation. \n",
    "What you are describing is a transition from a **Language Model** (predicting the next token) \n",
    "to a **Sequence-to-Sequence (Seq2Seq)** model using the \"Causal Language Modeling\" approach.\n",
    "\n",
    "By concatenating the source (Sinhala) and the target (English) into a single sequence,\n",
    " you are teaching the model that the English translation is the natural \"continuation\" of the Sinhala prompt.\n",
    "\n",
    "### How to adapt your Bigram + Attention model\n",
    "\n",
    "To make this work effectively, you should consider the following architectural and data adjustments:\n",
    "\n",
    "---\n",
    "\n",
    "### 1. The Data Format\n",
    "\n",
    "Your proposed format is exactly how modern models like GPT-3 were fine-tuned for tasks. \n",
    "You need a clear separator so the model knows where the \"context\" ends and the \"answer\" begins.\n",
    "\n",
    "* **Format:** `[Sinhala Sentence] <SEP> [English Translation] <END>`\n",
    "* **Example:** `niyural netwerks maru! <SEP> neural networks are awesome! <END>`\n",
    "\n",
    "### 2. The Loss Masking (Crucial Step)\n",
    "\n",
    "In a standard bigram/attention model, you calculate loss on every token. \n",
    "However, for translation, you don't necessarily want the model to be penalized for \n",
    "failing to predict the *Sinhala* part (since that is provided as input).\n",
    "\n",
    "* **Tip:** When calculating your cross-entropy loss, you can \"mask\" the Sinhala portion so \n",
    "the gradient only updates based on how well the model predicts the English tokens.\n",
    "\n",
    "### 3. Attention Mechanism: Causal Masking\n",
    "\n",
    "Since you are likely using a \"Bigram with Attention\" (similar to a Transformer Decoder), \n",
    "ensure you are using a **Look-ahead Mask**. \n",
    "This prevents the model from \"cheating\" by looking at the English words while \n",
    "it is still processing the Sinhala words during training.\n",
    "\n",
    "---\n",
    "\n",
    "### 4. Comparison of Approaches\n",
    "\n",
    "| Feature | Your Current Model (Generative) | Your Target Model (Translation) |\n",
    "| --- | --- | --- |\n",
    "| **Input** | A few Sinhala words | Full Sinhala sentence + `<SEP>` |\n",
    "| **Output** | More Sinhala words | Equivalent English meaning |\n",
    "| **Vocabulary** | Sinhala tokens only | Combined Sinhala + English tokens |\n",
    "| **Context** | Short-range (Bigram) | Long-range (Attention) |\n",
    "\n",
    "### 5. Potential Challenges\n",
    "\n",
    "* **Vocabulary Size:** Your embedding layer and final linear layer must now \n",
    "accommodate both Sinhala and English characters/tokens. \n",
    "If your current vocabulary is only Sinhala, you'll need to rebuild it.\n",
    "* **Bigram Limitations:** A pure bigram model (looking only at the previous word) is very weak for translation. \n",
    "The **Attention** mechanism will be doing 99% of the heavy lifting here to \n",
    "map the Sinhala \"source\" tokens to the English \"target\" tokens.\n",
    "\n",
    "---\n",
    "### Recommendations for Success\n",
    "\n",
    "1. **Use a SentencePiece or BPE Tokenizer:** Instead of character-level or word-level, use sub-word tokenization.\n",
    " This helps the model handle the complex morphology of Sinhala.\n",
    "2. **Increase Context Window:** Ensure your `block_size` (max tokens) is large enough to\n",
    " hold both the Sinhala sentence and its English translation combined.\n",
    "\n",
    "Breaking this down incrementally is a smart move. Transitioning from a simple next-token predictor to a translator involves moving from **unstructured generation** to **conditioned generation**.\n",
    "\n",
    "Here is the roadmap for your modifications, ranked from the most straightforward to the most complex:\n",
    "\n",
    "### Incremental Modification Roadmap\n",
    "\n",
    "| Phase | Modification | Difficulty | Why it's necessary |\n",
    "| --- | --- | --- | --- |\n",
    "| **1** | **Unified Vocabulary** |  Low | Your model must now recognize both Sinhala characters/words and English ones in the same embedding space. |\n",
    "| **2** | **Data Formatting** |  Low | You need to wrap your data in the `Source <SEP> Target <END>` format so the model learns the boundary. |\n",
    "| **3** | **Block Size Increase** |  Medium | The `block_size` (context window) must now be large enough to fit *both* sentences combined. |\n",
    "| **4** | **Inference Logic** |  Medium | You must change how you \"prompt\" the model. You feed it the Sinhala sentence + `<SEP>`, then let it auto-regressively generate until it hits `<END>`. |\n",
    "| **5** | **Loss Masking** |  High | To get high quality, you should tell the model *not* to learn/calculate loss on the Sinhala input, only on the English output. |\n",
    "| **6** | **Sub-word Tokenization** |  High | Moving from characters to BPE (Byte Pair Encoding) helps handle the \"mismatch\" in sentence lengths between the two languages. |\n",
    "\n",
    "---\n",
    "\n",
    "### Step-by-Step Implementation Strategy\n",
    "\n",
    "If you want to start today, I recommend following this order to see immediate results:\n",
    "\n",
    "#### Phase 1: The \"Lazy\" Translation Approach (Easiest)\n",
    "\n",
    "Don't change your model architecture yet. Just change your `train.txt`. Instead of just Sinhala text, feed it pairs.\n",
    "\n",
    "* **Action:** Update your tokenizer to include all English letters and special symbols like `<START>`, `<SEP>`, and `<END>`.\n",
    "* **Result:** The model will start treating English as a \"continuation\" of Sinhala.\n",
    "\n",
    "#### Phase 2: Refined Attention & Context\n",
    "\n",
    "Because translation requires \"looking back\" at the beginning of the sentence to decide the end of the translation, a simple Bigram won't cut it.\n",
    "\n",
    "* **Action:** Ensure your **Self-Attention** layer is robust. In a translation task, when the model is predicting the 5th English word, it needs to use the Attention weights to \"look\" at the 2nd Sinhala word.\n",
    "\n",
    "#### Phase 3: Optimizing the Objective (Loss Masking)\n",
    "\n",
    "In your current code, your loss function likely looks like this:\n",
    "`loss = F.cross_entropy(logits, targets)`\n",
    "\n",
    "To make it a true translator, you would modify the `targets` so that all positions corresponding to the Sinhala input are ignored (usually set to a value like `-100`). This forces the model's \"brain\" to focus entirely on the accuracy of the English translation.\n",
    "\n",
    "---\n",
    "\n",
    "**Would you like me to provide the Python code for the \"Phase 1\" data loader so you can start retraining with your Sinhala-English pairs immediately?**\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZcvKeBXoZFOY"
   },
   "source": [
    "### Full finished code, for reference\n",
    "\n",
    "You may want to refer directly to the git repo instead though."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Limitations of this style for Translation\n",
    "While GPT models are powerful, using a Decoder-only style for translation has one major drawback \n",
    "compared to an Encoder-Decoder style:\n",
    "\n",
    "Unidirectional Context: While writing the first word of the English translation, \n",
    " model can see the Sinhala sentence. However, while reading the first word of the Sinhala sentence,\n",
    " it cannot see the last word of the Sinhala sentence (because of the tril mask).\n",
    "\n",
    "Why that matters: In translation, the meaning of the first word often depends on the last word \n",
    "(especially in languages with different word orders like Sinhala).\n",
    "\n",
    "Can you make it better while keeping the GPT style?\n",
    "If you want to keep your current code structure but make it more powerful for translation, \n",
    "you can try \"PrefixLM\" masking. \n",
    "This involves allowing the tokens in the Sinhala part to see each other bidirectionally \n",
    "(removing the mask for the prefix only), but keeping the causal mask for the English part. \n",
    "However, this is quite complex to implement in your current Head class.\n",
    "\n",
    "Since you've seen a small improvement with loss masking, \n",
    "Try \"Scaling Up\" the hyperparameters (n_embd and n_layer) \n",
    "next to see if that pushes your BLEU score higher\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "hoelkOrFY8bN",
    "outputId": "961304cd-e379-40d4-dd56-8de0b91d2861",
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "parent_dir:c:\\Users\\Dasun\\Data\\NLP\n",
      "device:cuda\n",
      "loading ds...\n",
      "Selecting SinhalaToEnglish\n",
      "Loading data from: c:\\Users\\Dasun\\Data\\NLP\\datasets\\kaggle\\converted_data.csv\n",
      "f_path:c:\\Users\\Dasun\\Data\\NLP\\datasets\\kaggle\\converted_data.csv\n",
      "File found converted_data.csv\n",
      "DataFrame Length:34469\n",
      "Combined si<S>en<E>: 5235074,<class 'str'>\n",
      "<class 'str'>c esc ma tete මෙය මග\n",
      ",<class 'str'>is yes i am sir  <E>\n",
      "loaded ds:SinhalaToEnglish\n",
      "<class 'list'>:['J', 'G', 'e', 'O', 'ේ', 'è', 'ඔ', 'u', 'ථ', ',']...['f', 'x', 'Q', '<S>', '<E>']\n",
      "All Character Set:151{'J', 'G', 'e', 'O', 'ේ', 'è', 'ඔ', 'u', 'ථ', ',', 'ඡ', 'o', 'ඤ', 'ෝ', 'D', 'ඉ', 'U', 'd', 'ü', 'ම', 'ෙ', 'ට', 'ං', 'ඓ', 'ඥ', 'B', '8', 'ි', 'ය', '<E>', 'v', 'න', 'V', 'z', 'ක', 'l', 'ෂ', 'ෞ', 'ෲ', 'ප', 'බ', 'ඵ', 'H', 'ඌ', 'E', 'p', 'ධ', 'ච', 'ෑ', 'K', 'k', '3', 'ැ', '\"', 'é', 'ර', 'ඊ', 'ෘ', '<S>', 'i', '\\xad', 'c', 'ඒ', 'P', 'ළ', '4', 'ෆ', 'A', 'w', ' ', 'අ', 'y', '\\x97', '\\x96', '5', 'ග', 't', 'ඟ', 'ූ', '.', 'ඛ', '7', '0', 'ත', 'ණ', '්', 'N', 'ආ', 'ද', 'r', \"'\", 'ඨ', 'T', 'g', 'ඃ', 'ෛ', 'ඳ', 'ñ', '1', 'ඹ', '2', '9', 's', 'ො', 'W', 'ඕ', 'ව', 'ඬ', 'n', '£', 'ී', 'I', 'එ', 'Y', 'ඍ', 'ඈ', '\\u200b', 'ä', 'ඩ', 'F', 'හ', 'R', 'a', 'ඖ', 'S', 'භ', '\\u200d', 'M', 'ජ', 'ල', 'Z', 'b', '6', 'ු', 'q', 'h', 'C', 'L', 'උ', 'm', 'ඝ', '?', 'ඇ', 'j', 'ඪ', 'ශ', 'ා', 'ස', 'f', 'x', 'Q'}\n",
      "Vocab Size:151\n",
      "stoi {' ': 0, '\"': 1, \"'\": 2, ',': 3, '.': 4, '0': 5, '1': 6, '2': 7, '3': 8, '4': 9, '5': 10, '6': 11, '7': 12, '8': 13, '9': 14, '<E>': 15, '<S>': 16, '?': 17, 'A': 18, 'B': 19, 'C': 20, 'D': 21, 'E': 22, 'F': 23, 'G': 24, 'H': 25, 'I': 26, 'J': 27, 'K': 28, 'L': 29, 'M': 30, 'N': 31, 'O': 32, 'P': 33, 'Q': 34, 'R': 35, 'S': 36, 'T': 37, 'U': 38, 'V': 39, 'W': 40, 'Y': 41, 'Z': 42, 'a': 43, 'b': 44, 'c': 45, 'd': 46, 'e': 47, 'f': 48, 'g': 49, 'h': 50, 'i': 51, 'j': 52, 'k': 53, 'l': 54, 'm': 55, 'n': 56, 'o': 57, 'p': 58, 'q': 59, 'r': 60, 's': 61, 't': 62, 'u': 63, 'v': 64, 'w': 65, 'x': 66, 'y': 67, 'z': 68, '\\x96': 69, '\\x97': 70, '£': 71, '\\xad': 72, 'ä': 73, 'è': 74, 'é': 75, 'ñ': 76, 'ü': 77, 'ං': 78, 'ඃ': 79, 'අ': 80, 'ආ': 81, 'ඇ': 82, 'ඈ': 83, 'ඉ': 84, 'ඊ': 85, 'උ': 86, 'ඌ': 87, 'ඍ': 88, 'එ': 89, 'ඒ': 90, 'ඓ': 91, 'ඔ': 92, 'ඕ': 93, 'ඖ': 94, 'ක': 95, 'ඛ': 96, 'ග': 97, 'ඝ': 98, 'ඟ': 99, 'ච': 100, 'ඡ': 101, 'ජ': 102, 'ඤ': 103, 'ඥ': 104, 'ට': 105, 'ඨ': 106, 'ඩ': 107, 'ඪ': 108, 'ණ': 109, 'ඬ': 110, 'ත': 111, 'ථ': 112, 'ද': 113, 'ධ': 114, 'න': 115, 'ඳ': 116, 'ප': 117, 'ඵ': 118, 'බ': 119, 'භ': 120, 'ම': 121, 'ඹ': 122, 'ය': 123, 'ර': 124, 'ල': 125, 'ව': 126, 'ශ': 127, 'ෂ': 128, 'ස': 129, 'හ': 130, 'ළ': 131, 'ෆ': 132, '්': 133, 'ා': 134, 'ැ': 135, 'ෑ': 136, 'ි': 137, 'ී': 138, 'ු': 139, 'ූ': 140, 'ෘ': 141, 'ෙ': 142, 'ේ': 143, 'ෛ': 144, 'ො': 145, 'ෝ': 146, 'ෞ': 147, 'ෲ': 148, '\\u200b': 149, '\\u200d': 150}\n",
      "itos {0: ' ', 1: '\"', 2: \"'\", 3: ',', 4: '.', 5: '0', 6: '1', 7: '2', 8: '3', 9: '4', 10: '5', 11: '6', 12: '7', 13: '8', 14: '9', 15: '<E>', 16: '<S>', 17: '?', 18: 'A', 19: 'B', 20: 'C', 21: 'D', 22: 'E', 23: 'F', 24: 'G', 25: 'H', 26: 'I', 27: 'J', 28: 'K', 29: 'L', 30: 'M', 31: 'N', 32: 'O', 33: 'P', 34: 'Q', 35: 'R', 36: 'S', 37: 'T', 38: 'U', 39: 'V', 40: 'W', 41: 'Y', 42: 'Z', 43: 'a', 44: 'b', 45: 'c', 46: 'd', 47: 'e', 48: 'f', 49: 'g', 50: 'h', 51: 'i', 52: 'j', 53: 'k', 54: 'l', 55: 'm', 56: 'n', 57: 'o', 58: 'p', 59: 'q', 60: 'r', 61: 's', 62: 't', 63: 'u', 64: 'v', 65: 'w', 66: 'x', 67: 'y', 68: 'z', 69: '\\x96', 70: '\\x97', 71: '£', 72: '\\xad', 73: 'ä', 74: 'è', 75: 'é', 76: 'ñ', 77: 'ü', 78: 'ං', 79: 'ඃ', 80: 'අ', 81: 'ආ', 82: 'ඇ', 83: 'ඈ', 84: 'ඉ', 85: 'ඊ', 86: 'උ', 87: 'ඌ', 88: 'ඍ', 89: 'එ', 90: 'ඒ', 91: 'ඓ', 92: 'ඔ', 93: 'ඕ', 94: 'ඖ', 95: 'ක', 96: 'ඛ', 97: 'ග', 98: 'ඝ', 99: 'ඟ', 100: 'ච', 101: 'ඡ', 102: 'ජ', 103: 'ඤ', 104: 'ඥ', 105: 'ට', 106: 'ඨ', 107: 'ඩ', 108: 'ඪ', 109: 'ණ', 110: 'ඬ', 111: 'ත', 112: 'ථ', 113: 'ද', 114: 'ධ', 115: 'න', 116: 'ඳ', 117: 'ප', 118: 'ඵ', 119: 'බ', 120: 'භ', 121: 'ම', 122: 'ඹ', 123: 'ය', 124: 'ර', 125: 'ල', 126: 'ව', 127: 'ශ', 128: 'ෂ', 129: 'ස', 130: 'හ', 131: 'ළ', 132: 'ෆ', 133: '්', 134: 'ා', 135: 'ැ', 136: 'ෑ', 137: 'ි', 138: 'ී', 139: 'ු', 140: 'ූ', 141: 'ෘ', 142: 'ෙ', 143: 'ේ', 144: 'ෛ', 145: 'ො', 146: 'ෝ', 147: 'ෞ', 148: 'ෲ', 149: '\\u200b', 150: '\\u200d'}\n",
      "text[:80] :c esc ma tete මෙය මගේ ප්‍රධාන අයිතියයි ඔබ ප්‍රශ්නාවලිය සඳහා සූදානම් ද බලන්න <S> \n",
      "encoded   :[45, 0, 47, 61, 45, 0, 55, 43, 0, 62, 47, 62, 47, 0, 121, 142, 123, 0, 121, 97, 143, 0, 117, 133, 150, 124, 114, 134, 115, 0, 80, 123, 137, 111, 137, 123, 123, 137, 0, 92, 119, 0, 117, 133, 150, 124, 127, 133, 115, 134, 126, 125, 137, 123, 0, 129, 116, 130, 134, 0, 129, 140, 113, 134, 115, 121, 133, 0, 113, 0, 119, 125, 115, 133, 115, 0, 16, 0, 45, 0, 47, 61, 45, 0, 55, 43, 0, 62, 47, 62, 47, 0, 62, 50, 51, 61, 0, 51, 61, 0, 55, 67, 0, 50, 47, 43, 46, 60, 51, 49, 50, 62, 0, 61, 47, 47, 0, 67, 57, 63, 0, 60, 47, 0, 60, 47, 43, 46, 67, 0, 48, 57, 60, 0, 62, 50, 47, 0, 59, 63, 51, 68, 0, 0, 15, 115, 136, 0, 115, 136, 0, 90, 95, 0, 121, 97, 143, 0, 126, 124, 113, 95, 133, 0, 80, 117, 137, 105, 0, 130, 124, 137, 0, 130, 121, 115, 133, 0, 95, 135, 121, 124, 145, 115, 133, 0, 95, 142, 115, 142, 95, 133, 0, 130, 137, 105, 137, 123, 142, 0, 115, 136, 0, 16, 0, 56, 57, 0, 56, 57, 0, 51, 62, 0, 61, 0, 55, 67, 0, 48, 43, 63, 54, 62, 0, 65, 47, 0, 46, 51, 46, 56, 0, 62, 0, 50, 43, 64, 47, 0, 43, 0, 58, 60, 57, 58, 47, 60, 0, 51]\n",
      "decoded   :c esc ma tete මෙය මගේ ප්‍රධාන අයිතියයි ඔබ ප්‍රශ්නාවලිය සඳහා සූදානම් ද බලන්න <S> c esc ma tete this is my headright see you re ready for the quiz  <E>නෑ නෑ ඒක මගේ වරදක් අපිට හරි හමන් කැමරොන් කෙනෙක් හිටියෙ නෑ <S> no no it s my fault we didn t have a proper i\n",
      "train_data needs clean:nks  <E>වෝනර්ට ඔහුගේම රෙදි සෝදන්නවත් බැහැ, මම දන්නවා ඔහු එය එවා ඇති බව \n",
      "clean train_data: anks \n",
      "val_data start needs clean:out  <E>මම දන්නවා ඔහු එය එවා ඇති බව ඔබ දන්නවාද ඔහුගේ පියාට ඇමතුමක් ගැනීමට ඔහු ඉල්ලුම් කළ විට ඔහු පොරොත්තු ලේඛනයට පත් වූ බව ඔබ දන්නවාද? \n",
      "clean val_data:...මම දන්නවා \n"
     ]
    },
    {
     "ename": "Exception",
     "evalue": "Execution Halt",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mException\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[33], line 219\u001b[0m\n\u001b[0;32m    216\u001b[0m     test_data \u001b[38;5;241m=\u001b[39m data[n1:]\n\u001b[0;32m    217\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m train_data,val_data,test_data\n\u001b[1;32m--> 219\u001b[0m train_data,val_data,test_data \u001b[38;5;241m=\u001b[39m \u001b[43msplit_train_val_data\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    220\u001b[0m halt_if(\u001b[38;5;28;01mTrue\u001b[39;00m)\n\u001b[0;32m    221\u001b[0m \u001b[38;5;66;03m# Convert to tensors\u001b[39;00m\n",
      "Cell \u001b[1;32mIn[33], line 208\u001b[0m, in \u001b[0;36msplit_train_val_data\u001b[1;34m(data)\u001b[0m\n\u001b[0;32m    206\u001b[0m     val_data \u001b[38;5;241m=\u001b[39m val_data[(eci\u001b[38;5;241m+\u001b[39m\u001b[38;5;241m3\u001b[39m):]\n\u001b[0;32m    207\u001b[0m     \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mclean val_data:...\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mval_data[\u001b[38;5;241m0\u001b[39m:\u001b[38;5;241m10\u001b[39m]\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m--> 208\u001b[0m \u001b[43mhalt_if\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m)\u001b[49m\n\u001b[0;32m    209\u001b[0m eci \u001b[38;5;241m=\u001b[39m val_data\u001b[38;5;241m.\u001b[39mrfind(oparam\u001b[38;5;241m.\u001b[39mending_char)\n\u001b[0;32m    210\u001b[0m sci \u001b[38;5;241m=\u001b[39m val_data\u001b[38;5;241m.\u001b[39mrfind(oparam\u001b[38;5;241m.\u001b[39mseperator_char)\n",
      "Cell \u001b[1;32mIn[33], line 39\u001b[0m, in \u001b[0;36mhalt_if\u001b[1;34m(flag)\u001b[0m\n\u001b[0;32m     35\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[0;32m     36\u001b[0m \u001b[38;5;124;03mhalt execution if flag is true\u001b[39;00m\n\u001b[0;32m     37\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[0;32m     38\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m flag:\n\u001b[1;32m---> 39\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mExecution Halt\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
      "\u001b[1;31mException\u001b[0m: Execution Halt"
     ]
    }
   ],
   "source": [
    "from IPython.display import display, HTML\n",
    "\n",
    "def create_input(label, type=\"text\", value=\"\"):# text or number\n",
    "    \"\"\"Creates a dynamic HTML input field.\"\"\"\n",
    "    html = f\"\"\"\n",
    "    <input type=\"{type}\" value=\"{value}\" placeholder=\"{label}\" />\n",
    "    \"\"\"\n",
    "    return input()\n",
    "\n",
    "# Input for a number\n",
    "# x = create_input(\"Enter a number:\", \"number\")\n",
    "#----------------------------------------\n",
    "import sys\n",
    "import os\n",
    "\n",
    "os.environ['CUDA_LAUNCH_BLOCKING'] = '1'\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import functional as F\n",
    "\n",
    "# Enable TensorFloat32 for faster matmul on Ampere+ GPUs\n",
    "torch.set_float32_matmul_precision('high')\n",
    "\n",
    "\n",
    "# Find path to the dataset\n",
    "parent_dir = os.path.abspath(os.path.join(os.getcwd(), \"../..\"))\n",
    "print(f\"parent_dir:{parent_dir}\")\n",
    "kaggle_ds_path = parent_dir+\"\\\\datasets\\\\kaggle\\\\\"\n",
    "if kaggle_ds_path not in sys.path:\n",
    "    sys.path.append(kaggle_ds_path)\n",
    "    print(f\"kaggle_ds_path:{kaggle_ds_path}\")\n",
    "\n",
    "def halt_if(flag=True):\n",
    "    \"\"\"\n",
    "    halt execution if flag is true\n",
    "    \"\"\"\n",
    "    if flag:\n",
    "        raise Exception(\"Execution Halt\")\n",
    "    \n",
    "import dataload\n",
    "\n",
    "avail_datasets = {0:\"TinyShakespear\",1:\"SinhalaOnly\",2:\"SinhalaToEnglish\"}\n",
    "selected_dataset = 2\n",
    "\n",
    "base_dir = \"C:/Users/Dasun/Data/NLP/\"\n",
    "\n",
    "class oparam: # Other Parameters\n",
    "    default_model_name = (\"bigram\",\"singit_1\")[0]\n",
    "    seperator_char = \"<S>\"\n",
    "    ending_char = \"<E>\"\n",
    "\n",
    "def select_dataset():\n",
    "    global selected_dataset\n",
    "    global avail_datasets\n",
    "    global base_dir\n",
    "    print(f\"Selecting {avail_datasets[selected_dataset]}\")\n",
    "    if(selected_dataset == 0):\n",
    "        with open('input.txt', 'r', encoding='utf-8') as f:\n",
    "            text = f.read()\n",
    "            print(f\"{len(text)},{type(text)}\")\n",
    "            return text,avail_datasets[selected_dataset]\n",
    "    elif selected_dataset == 1:\n",
    "        df = dataload.load_data(kaggle_ds_path)\n",
    "        text = \"\".join(df['Sinhala'].dropna().astype(str))\n",
    "        print(f\"Sinhala final text type: {type(text)}\") # Should be class 'str'\n",
    "        print(f\"{len(text)},{type(text)}\")\n",
    "        print(f\"{type(text[0])}{(text[:20])}\\n,{type(text[-1])}{(text[-20:])}\")\n",
    "        return text,avail_datasets[selected_dataset]\n",
    "    elif selected_dataset == 2:\n",
    "        df = dataload.load_data(kaggle_ds_path)\n",
    "        # Format of the data: \"Sinhala <S> English <E>\"\n",
    "        # We use <S> to tell the model to start translating and <E> to stop.\n",
    "        df['combined'] = df.iloc[:, 0] + \" \"+oparam.seperator_char+\" \" + df.iloc[:, 1] + \" \"+oparam.ending_char+\"\"\n",
    "        text = \"\".join(df['combined'].dropna().astype(str).tolist()) \n",
    "        print(f\"Combined si<S>en<E>: {len(text)},{type(text)}\")\n",
    "        print(f\"{type(text[0])}{(text[:20])}\\n,{type(text[-1])}{(text[-20:])}\")\n",
    "        return text,avail_datasets[selected_dataset]\n",
    "    else:\n",
    "        raise Exception(\"NotImplementedYet\")\n",
    "\n",
    "#block_size = ? must be large enough for (Sinhala + <S> + English)\n",
    "\n",
    "# hyperparameters\n",
    "class hp:\n",
    "    batch_size = 384# 256 # how many independent sequences will we process in parallel?\n",
    "    block_size = 128 #32 # what is the maximum context length for predictions?\n",
    "    #Block size change after saving the model does not work\n",
    "    max_iters = 2**9 #5000\n",
    "    eval_interval = 40 #20 # number of iteration to train before estimating the loss\n",
    "    learning_rate = 1e-3 # karpathy:3e-4 # use a smaller lr than original 1e-3\n",
    "    device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "    print(f\"device:{device}\")\n",
    "    eval_iters = 40 # number of times(for samples) the loss is aggregated and averaged\n",
    "    n_embd = 128 \n",
    "    \"\"\"\n",
    "1. Characters have no \"meaning\" on their own. In a word-level model, \n",
    "the word \"Apple\" might have an embedding that represents \"fruit,\" \"red,\" or \"sweet.\" \n",
    "In your character-level model, the character 'a' means nothing by itself. \n",
    "It could be part of \"apple,\" \"angry,\" or \"at.\"\n",
    "The n_embd (64 in your code) acts as a \"Working Memory\" or \"Space for Context.\" \n",
    "It gives the model 64 different \"channels\" to store information about that character in its specific position.\n",
    "Channel 1: Is this character at the start of a word?\n",
    "Channel 2: Is this a Sinhala character or an English one?\n",
    "Channel 3: Is this character part of a vowel sound?\n",
    "Channel 4-64: Complex patterns the model discovers during training.\n",
    "2. The Role of n_embd as a \"Vector Space\"\n",
    "When you tokenized your text, you turned the character 'ම' into the integer 45. \n",
    "Computers cannot do \"math\" on the number 45 to understand language. They need a Vector.\n",
    "The n_embd parameter defines the size of this vector. \n",
    "Instead of the model seeing \"45\", it sees a row of 64 numbers.\n",
    "3. How it enables Translation: In your Sinhala-to-English task, \n",
    "the n_embd space is where the \"translation\" actually happens.\n",
    "Imagine the 64-dimensional space as a giant map. \n",
    "During training, the model learns to move the vector for the Sinhala character ම \n",
    "and the English character I into regions of the map that represent \"First Person Subject.\"\n",
    " \"\"\"   \n",
    "    n_head = 8\n",
    "    n_layer = 4\n",
    "    dropout = 0.0 # not sure if useful for small models\n",
    "    temperature = 0.3#1.0 # Lower:more focused Higher:more creative\n",
    "    weight_decay = None#1e-2 # did not improve the metrics\n",
    "    use_hp_values_in_code = True # Make the values defined here override the model saved params\n",
    "\n",
    "batch_size = hp.batch_size \n",
    "block_size = hp.block_size\n",
    "max_iters = hp.max_iters\n",
    "eval_interval = hp.eval_interval\n",
    "learning_rate = hp.learning_rate\n",
    "device = hp.device\n",
    "eval_iters = hp.eval_iters\n",
    "n_embd = hp.n_embd\n",
    "n_head = hp.n_head\n",
    "n_layer = hp.n_layer\n",
    "dropout = hp.dropout\n",
    "temprtr = hp.temperature\n",
    "weight_decay = hp.weight_decay\n",
    "# -------------------------------------------\n",
    "torch.manual_seed(1337)\n",
    "#------------Data Loading Section------------\n",
    "\n",
    "print(f\"loading ds...\")\n",
    "text,dsn = select_dataset() # Access the selected dataset\n",
    "print(f\"loaded ds:{dsn}\")\n",
    "# list_set = list(set(text)) # Generalizing it to accommodate <S> and <E>\n",
    "char_list:list[str] = []\n",
    "text_copy = str(text)\n",
    "text_copy = text_copy.replace(oparam.seperator_char,\"\").replace(oparam.ending_char,\"\")\n",
    "char_list = list(set(text_copy))\n",
    "char_list.append(oparam.seperator_char)\n",
    "char_list.append(oparam.ending_char)\n",
    "print(f\"{type(char_list)}:{char_list[:10]}...{char_list[-5:]}\")\n",
    "all_char_set = set(char_list)\n",
    "print(f\"All Character Set:{len(all_char_set)}{all_char_set}\")\n",
    "# Here are all the unique characters that occur in this text\n",
    "chars = sorted(list(all_char_set))\n",
    "vocab_size = len(chars)\n",
    "print(f\"Vocab Size:{vocab_size}\")\n",
    "# Create a mapping from characters to integers\n",
    "stoi = { ch:i for i,ch in enumerate(chars) }\n",
    "itos = { i:ch for i,ch in enumerate(chars) }\n",
    "print(\"stoi\",stoi)\n",
    "print(\"itos\",itos)\n",
    "def encode(s):# encoder: take a string, output a list of integers\n",
    "    ret:list[int] = []\n",
    "    l = len(s)\n",
    "    i = 0\n",
    "    while i < l:\n",
    "        c = s[i]\n",
    "        if c == '<' and i+2<l: # Handle special control characters\n",
    "            t = \"\".join([s[i],s[i+1],s[i+2]]) #<E> or <S>\n",
    "            if t == oparam.ending_char or t == oparam.seperator_char:\n",
    "                ret.append(stoi[t])\n",
    "                i+=3\n",
    "            else: # May be there are '<' s independent from control chars\n",
    "                ret.append(stoi[c])\n",
    "                i+=1\n",
    "        else:\n",
    "            ret.append(stoi[c])\n",
    "            i+=1\n",
    "    return ret\n",
    "\n",
    "decode = lambda l: ''.join([itos[i] for i in l]) # decoder: take a list of integers, output a string\n",
    "print(f\"text[:80] :{text[:80]}\")\n",
    "t = encode(text[:256])\n",
    "print(f\"encoded   :{t}\")\n",
    "print(f\"decoded   :{decode(t)}\")\n",
    "\n",
    "# Train and test splits\n",
    "def split_train_val_data(data=text):\n",
    "    n0 = int(0.8*len(data)) # first 80% will be train\n",
    "    train_data = data[:n0]\n",
    "    eci = train_data.rfind(oparam.ending_char)\n",
    "    sci = train_data.rfind(oparam.seperator_char)\n",
    "    if sci > eci:# Clean the wrong cut where last sentence has no translation\n",
    "        print(f\"train_data needs clean:{train_data[eci-5:sci]}\")\n",
    "        train_data = train_data[:eci] # cut until ec\n",
    "        print(f\"clean train_data: {train_data[-6:-1]}\")\n",
    "    \n",
    "    n1 = int(0.9*len(data)) # Rest val(10%)+test(10%)\n",
    "    val_data = data[n0:n1]\n",
    "    eci = val_data.find(oparam.ending_char)\n",
    "    sci = val_data.find(oparam.seperator_char)\n",
    "    if eci < sci: #\n",
    "        print(f\"val_data start needs clean:{val_data[eci-5:sci]}\")\n",
    "        val_data = val_data[(eci+3):]\n",
    "        print(f\"clean val_data:{val_data[0:10]}\")\n",
    "    halt_if(False)\n",
    "    eci = val_data.rfind(oparam.ending_char)\n",
    "    sci = val_data.rfind(oparam.seperator_char)\n",
    "    if eci < sci :# Clean the wrong cut where last sentence has no translation\n",
    "        print(f\"val_data end needs clean:{val_data[eci-5:sci]}\")\n",
    "        val_data = val_data[:eci] # cut until ec\n",
    "        print(f\"clean val_data:{val_data[-6:-1]}\")\n",
    "\n",
    "    test_data = data[n1:]\n",
    "    return train_data,val_data,test_data\n",
    "\n",
    "train_data,val_data,test_data = split_train_val_data()\n",
    "halt_if(False)\n",
    "# Convert to tensors\n",
    "train_data = torch.tensor(encode(train_data), dtype=torch.long)\n",
    "val_data = torch.tensor(encode(val_data), dtype=torch.long)\n",
    "\n",
    "# data loading\n",
    "def get_batch(split):\n",
    "    # generate a small batch of data of inputs x and targets y\n",
    "    data = train_data if split == 'train' else val_data\n",
    "    if len(data) - block_size <= 0:\n",
    "        print(f\"randint({len(data)} - {block_size}, {(batch_size,)})!!!\")\n",
    "    ix = torch.randint(len(data) - block_size, (batch_size,))\n",
    "    x = torch.stack([data[i:i+block_size] for i in ix])\n",
    "    y = torch.stack([data[i+1:i+block_size+1] for i in ix])\n",
    "    x, y = x.to(device), y.to(device)\n",
    "    return x, y\n",
    "\n",
    "@torch.no_grad()\n",
    "def estimate_loss():\n",
    "    out = {}\n",
    "    model.eval()\n",
    "    for split in ['train', 'val']:\n",
    "        losses = torch.zeros(eval_iters)\n",
    "        for k in range(eval_iters):\n",
    "            X, Y = get_batch(split)\n",
    "            # Pass the <S> token ID here so validation loss is also masked\n",
    "            logits, loss = model(X, Y, sep_id=stoi['<S>'])\n",
    "\n",
    "            losses[k] = loss.item()\n",
    "        out[split] = losses.mean()\n",
    "    model.train()\n",
    "    return out\n",
    "\n",
    "class Head(nn.Module):\n",
    "    \"\"\" one head of self-attention \"\"\"\n",
    "\n",
    "    def __init__(self, head_size):\n",
    "        super().__init__()\n",
    "        self.key = nn.Linear(n_embd, head_size, bias=False)\n",
    "        self.query = nn.Linear(n_embd, head_size, bias=False)\n",
    "        self.value = nn.Linear(n_embd, head_size, bias=False)\n",
    "        self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))\n",
    "\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x):\n",
    "        B,T,C = x.shape\n",
    "        k = self.key(x)   # (B,T,C)\n",
    "        q = self.query(x) # (B,T,C)\n",
    "        # compute attention scores (\"affinities\")\n",
    "        wei = q @ k.transpose(-2,-1) * C**-0.5 # (B, T, C) @ (B, C, T) -> (B, T, T)\n",
    "        wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf')) # (B, T, T)\n",
    "        wei = F.softmax(wei, dim=-1) # (B, T, T)\n",
    "        wei = self.dropout(wei)\n",
    "        # perform the weighted aggregation of the values\n",
    "        v = self.value(x) # (B,T,C)\n",
    "        out = wei @ v # (B, T, T) @ (B, T, C) -> (B, T, C)\n",
    "        return out\n",
    "\n",
    "class MultiHeadAttention(nn.Module):\n",
    "    \"\"\" multiple heads of self-attention in parallel \"\"\"\n",
    "\n",
    "    def __init__(self, num_heads, head_size):\n",
    "        super().__init__()\n",
    "        self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])\n",
    "        self.proj = nn.Linear(n_embd, n_embd)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = torch.cat([h(x) for h in self.heads], dim=-1)\n",
    "        out = self.dropout(self.proj(out))\n",
    "        return out\n",
    "\n",
    "class FeedFoward(nn.Module):\n",
    "    \"\"\" a simple linear layer followed by a non-linearity \"\"\"\n",
    "\n",
    "    def __init__(self, n_embd):\n",
    "        super().__init__()\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Linear(n_embd, 4 * n_embd),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(4 * n_embd, n_embd),\n",
    "            nn.Dropout(dropout),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.net(x)\n",
    "\n",
    "class Block(nn.Module):\n",
    "    \"\"\" Transformer block: communication followed by computation \"\"\"\n",
    "\n",
    "    def __init__(self, n_embd, n_head):\n",
    "        # n_embd: embedding dimension, n_head: the number of heads we'd like\n",
    "        super().__init__()\n",
    "        head_size = n_embd // n_head\n",
    "        self.sa = MultiHeadAttention(n_head, head_size)\n",
    "        self.ffwd = FeedFoward(n_embd)\n",
    "        self.ln1 = nn.LayerNorm(n_embd)\n",
    "        self.ln2 = nn.LayerNorm(n_embd)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x + self.sa(self.ln1(x))\n",
    "        x = x + self.ffwd(self.ln2(x))\n",
    "        return x\n",
    "\n",
    "# super simple bigram model\n",
    "class BigramLanguageModel(nn.Module):\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        print(f\"Initializing Model: Vocab={vocab_size}, Embd={n_embd}\")\n",
    "        # each token directly reads off the logits for the next token from a lookup table\n",
    "        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)\n",
    "        self.position_embedding_table = nn.Embedding(block_size, n_embd)\n",
    "        self.blocks = nn.Sequential(*[Block(n_embd, n_head=n_head) for _ in range(n_layer)])\n",
    "        self.ln_f = nn.LayerNorm(n_embd) # final layer norm\n",
    "        self.lm_head = nn.Linear(n_embd, vocab_size)\n",
    "\n",
    "    def forward(self, idx, targets=None,sep_id=None):\n",
    "        B, T = idx.shape\n",
    "\n",
    "        # idx and targets are both (B,T) tensor of integers\n",
    "        tok_emb = self.token_embedding_table(idx) # (B,T,C)\n",
    "        pos_emb = self.position_embedding_table(torch.arange(T, device=device)) # (T,C)\n",
    "        x = tok_emb + pos_emb # (B,T,C)\n",
    "        x = self.blocks(x) # (B,T,C)\n",
    "        x = self.ln_f(x) # (B,T,C)\n",
    "        logits = self.lm_head(x) # (B,T,vocab_size)\n",
    "\n",
    "        if targets is None:\n",
    "            loss = None\n",
    "        else:\n",
    "            B, T, C = logits.shape\n",
    "            logits = logits.view(B*T, C)\n",
    "\n",
    "            # Loss Masking\n",
    "            if sep_id is not None:\n",
    "                targets_masked = targets.clone()\n",
    "                for b in range(B):\n",
    "                    # Find first occurrence of [SEP] token\n",
    "                    sep_pos = (targets[b] == sep_id).nonzero(as_tuple=True)[0]\n",
    "                    if len(sep_pos) > 0:\n",
    "                        # Mask everything before and including [SEP]\n",
    "                        targets_masked[b, :sep_pos[0].item() + 1] = -100\n",
    "                targets = targets_masked\n",
    "\n",
    "            targets = targets.view(B*T)\n",
    "            loss = F.cross_entropy(logits, targets,ignore_index=-100) # ignore_index=-100 ensures we don't calculate loss on Sinhala input\n",
    "\n",
    "        return logits, loss\n",
    "\n",
    "    def generate(self, idx, max_new_tokens, temperature=temprtr):\n",
    "        # idx is (B, T) array of indices in the current context\n",
    "        for _ in range(max_new_tokens):\n",
    "            # crop idx to the last block_size tokens\n",
    "            idx_cond = idx[:, -block_size:]\n",
    "            # get the predictions\n",
    "            logits, loss = self(idx_cond)\n",
    "            # focus only on the last time step\n",
    "            logits = logits[:, -1, :] # becomes (B, C)\n",
    "            # apply temperature\n",
    "            logits = logits/temperature\n",
    "            # apply softmax to get probabilities\n",
    "            probs = F.softmax(logits, dim=-1) # (B, C)\n",
    "            # sample from the distribution\n",
    "            idx_next = torch.multinomial(probs, num_samples=1) # (B, 1)\n",
    "            # append sampled index to the running sequence\n",
    "            idx = torch.cat((idx, idx_next), dim=1) # (B, T+1)\n",
    "        return idx\n",
    "\n",
    "#---save and load methods---\n",
    "name_tmpl =  'name.pth'\n",
    "best_val_loss = float('inf')  # Any first loss is better\n",
    "def save_model(mod,optm,tot_iter,vali_loss,name=oparam.default_model_name):\n",
    "    #import json\n",
    "    \n",
    "    checkpoint = {\n",
    "        'model_state_dict': mod.state_dict(),\n",
    "        'optimizer_state_dict': optm.state_dict(),\n",
    "        'config': {\n",
    "            'n_embd': hp.n_embd,\n",
    "            'n_head': hp.n_head,\n",
    "            'n_layer': hp.n_layer,\n",
    "            'block_size': hp.block_size,\n",
    "            'vocab_size': vocab_size,\n",
    "        },\n",
    "        'vocab': {\n",
    "            'stoi': stoi,\n",
    "            'itos': itos\n",
    "        },\n",
    "        'iters_completed': tot_iter,\n",
    "        'vali_loss':vali_loss\n",
    "    }\n",
    "    name = name_tmpl.replace('name',name)\n",
    "    print(f\"Model File:{name}\")\n",
    "    torch.save(checkpoint, name)\n",
    "    print(\"Successfully saved model and vocabulary mappings.\")\n",
    "# ------------\n",
    "def load_model(dev,name=oparam.default_model_name,info_only=False):\n",
    "    name = name_tmpl.replace('name',name)\n",
    "    print(f\"Trying to load model:{name}\")\n",
    "\n",
    "    checkpoint = torch.load(name, map_location=dev)\n",
    "    conf = checkpoint['config']\n",
    "\n",
    "    if info_only:\n",
    "        print(f\"conf['n_embd']{conf['n_embd']}\")\n",
    "        print(f\"conf['n_head']{conf['n_head']}\")\n",
    "        print(f\"conf['n_layer']{conf['n_layer']}\")\n",
    "        print(f\"conf['block_size']{conf['block_size']}\")\n",
    "        print(f\"conf['vocab_size']{conf['vocab_size']}\")\n",
    "        print(f\"iters_completed:{checkpoint['iters_completed']}\")\n",
    "        print(f\"vali_loss:{checkpoint['vali_loss']}\")\n",
    "        return # exit\n",
    "\n",
    "    conf = checkpoint['config']    \n",
    "    # Restore vocabulary and mappings\n",
    "    stoi = checkpoint['vocab']['stoi']\n",
    "    itos = {int(k): v for k, v in checkpoint['vocab']['itos'].items()} # JSON keys are strings\n",
    "    print(f\"stoi:{stoi}\\nitos:{itos}\")\n",
    "\n",
    "    if not hp.use_hp_values_in_code:# if not overridden, use the model saved values\n",
    "        hp.n_embd = conf['n_embd']\n",
    "        hp.n_head = conf['n_head']\n",
    "        hp.n_layer = conf['n_layer']\n",
    "        hp.block_size = conf['block_size']\n",
    "    \n",
    "    hp.vocab_size = conf['vocab_size']\n",
    "    global vocab_size\n",
    "    vocab_size = hp.vocab_size\n",
    "    \n",
    "    print(f\"Hyperparameters restored: n_embd={n_embd}, n_head={n_head}, block_size={block_size}\")\n",
    "    \n",
    "    model = BigramLanguageModel() \n",
    "    model.load_state_dict(checkpoint['model_state_dict'])\n",
    "    model.to(device)\n",
    "    \n",
    "    # Restore optimizer (only if you want to keep training)\n",
    "    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)\n",
    "    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n",
    "    iter_complete = int(checkpoint['iters_completed'])\n",
    "    try:\n",
    "        vali_loss = float(checkpoint['vali_loss'])\n",
    "    except:\n",
    "        vali_loss = float('inf')\n",
    "    print(f\"Loaded model from step {iter_complete}\")\n",
    "    return model,optimizer,iter_complete,vali_loss\n",
    "#---save and load methods---\n",
    "import pandas as pd\n",
    "from datetime import datetime\n",
    "\n",
    "def update_model_registry(m_id, final_loss, csv_path='model_ids.csv'):\n",
    "    try:\n",
    "        df = pd.read_csv(csv_path)\n",
    "        df.columns = df.columns.str.strip()\n",
    "        # Check if the ID exists\n",
    "        if m_id in df['model_id'].values:\n",
    "            timestamp = datetime.now().strftime(\"%Y-%m-%d %H:%M\")\n",
    "            # Update the model_info column with latest stats\n",
    "            new_info = f\"Loss: {final_loss:.4f} | Updated: {timestamp}\"\n",
    "            df.loc[df['model_id'] == m_id, 'model_info'] = new_info\n",
    "            \n",
    "            df.to_csv(csv_path, index=False)\n",
    "            print(f\"Registry updated for {m_id}\")\n",
    "        else:\n",
    "            print(f\"Could not update: ID {m_id} not found.\")\n",
    "            \n",
    "    except Exception as e:\n",
    "        print(f\"Registry update failed: {e}\")\n",
    "\n",
    "def get_model_filename(m_id, csv_path='model_ids.csv'):\n",
    "    df = pd.read_csv(csv_path)\n",
    "    df.columns = df.columns.str.strip()\n",
    "    result = df[df['model_id'] == m_id]\n",
    "    return result.iloc[0]['model_f_name'] if not result.empty else None\n",
    "\n",
    "do_inference_only = False\n",
    "tot_iter = 0\n",
    "max_tokens = 10\n",
    "do_translate_only = False\n",
    "sinh_sentence = None\n",
    "do_evaluate_only = False\n",
    "# Input from user\n",
    "while True:\n",
    "    print(\"-------------------------------- MENU --------------------------------\")\n",
    "    print(\"Menu: 0:load; 1:train; 2:info; 3:gener; 4:trans; 5:evaluate -1:Exit:\",flush=True)\n",
    "    ui = int(create_input(\"?\", \"number\"))\n",
    "    print(\"Choice:\",ui)\n",
    "    if ui == 0:\n",
    "        print(f\"Load model from disk to train\")\n",
    "        model,optimizer,tot_iter,best_val_loss = load_model(dev=device)\n",
    "        m = model.to(device)\n",
    "        break\n",
    "    elif ui == 1 :\n",
    "        print(f\"Train the model from scratch\")\n",
    "        print(\"Model ID?\",flush=True)\n",
    "        ui1 = str(create_input(\"?\", \"string\"))\n",
    "        m_id = get_model_filename(ui1)\n",
    "        model = BigramLanguageModel()\n",
    "        m = model.to(device)\n",
    "        # Create a PyTorch optimizer\n",
    "        if hp.weight_decay is not None:\n",
    "            optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=hp.weight_decay)\n",
    "        else:\n",
    "            optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)\n",
    "        break\n",
    "    elif ui == 2:\n",
    "        print(f\"Show saved model information\")\n",
    "        print(\"File Name?\",flush=True)\n",
    "        ui1 = str(create_input(\"File Name?\", \"string\"))\n",
    "        print(\"Reading:\",ui1)\n",
    "        load_model(dev=device,name=ui1,info_only=True)\n",
    "    elif ui == 3:\n",
    "        print(f\"Generate tokens\")\n",
    "        flag = False\n",
    "        if flag:\n",
    "            print(\"File Name?\",flush=True)\n",
    "            ui1 = str(create_input(\"?\", \"string\"))\n",
    "            model,optimizer,tot_iter,best_val_loss = load_model(dev=device,name=ui1)\n",
    "        ui1 = 'bigram'\n",
    "        print(f\"Loading {ui1} to generate tokens:\")\n",
    "        model,optimizer,tot_iter,best_val_loss = load_model(dev=device)\n",
    "        print(\"Max Tokens?\",end='',flush=True)\n",
    "        ui1 = int(create_input(\"?\", \"number\"))\n",
    "        print(ui1)\n",
    "        max_tokens = ui1\n",
    "        m = model.to(device)\n",
    "        do_inference_only = True\n",
    "        break\n",
    "    elif ui == 4:\n",
    "        print(f\"Translate a sentence\")\n",
    "        flag = False # skip querying the user for a file name\n",
    "        if flag:\n",
    "            print(\"File Name?\",flush=True)\n",
    "            ui1 = str(create_input(\"?\", \"string\"))\n",
    "        else:    \n",
    "            ui1 = \"bigram\"\n",
    "        print(f\"Loading {ui1} to generate tokens:\")\n",
    "        model,optimizer,tot_iter,best_val_loss = load_model(dev=device,name=ui1)\n",
    "        print(\"Input a Sinhala sentence to translate:\",end='',flush=True)\n",
    "        ui1 = str(create_input(\"?\", \"str\"))\n",
    "        print(\"Sinhala:\",ui1)\n",
    "        sinh_sentence = ui1\n",
    "        m = model.to(device)\n",
    "        do_translate_only = True\n",
    "        break\n",
    "    elif ui == 5:\n",
    "        print(f\"Evaluate the model\")\n",
    "        flag = False\n",
    "        if flag:\n",
    "            print(\"File Name?\",flush=True)\n",
    "            ui1 = str(create_input(\"?\", \"string\"))\n",
    "        else:    \n",
    "            ui1 = \"bigram\"\n",
    "        print(f\"Loading {ui1} to evaluate\")\n",
    "        model,optimizer,tot_iter,best_val_loss = load_model(dev=device,name=ui1)\n",
    "        m = model.to(device)\n",
    "        do_evaluate_only = True\n",
    "        break\n",
    "    elif ui == -1:\n",
    "        raise Exception(\"Ending Cell\")\n",
    "\n",
    "# print the number of parameters in the model\n",
    "print(sum(p.numel() for p in m.parameters())/1000, 'K parameter model')\n",
    "\n",
    "def translate_sin2eng(source_senten:str):\n",
    "    \"\"\"\n",
    "    Given a source sentence as input returns the translated sentence from the model\n",
    "    \"\"\"\n",
    "    if source_senten is not None:\n",
    "        source_senten = source_senten.strip() # remove any garbage\n",
    "    else:\n",
    "        return None\n",
    "    \n",
    "    if not source_senten.endswith(oparam.seperator_char):\n",
    "        source_senten+=oparam.seperator_char # Include the seperator character if missing\n",
    "    \n",
    "    context = torch.tensor(encode(source_senten), dtype=torch.long, device=device).unsqueeze(0)\n",
    "\n",
    "    #halt_if(False)\n",
    "    tokens_at_once = 5 # generate small number of tokens with looping, this avoids the long wait\n",
    "    target_senten = \"\"\n",
    "    MAX_TOK = 256\n",
    "    print(source_senten)\n",
    "    for _ in range(0,MAX_TOK,tokens_at_once):\n",
    "        output = m.generate(context, max_new_tokens=tokens_at_once)\n",
    "        #print(\"output.shape\",output.shape)     \n",
    "        new_tokens = output[0, context.shape[1]:]# 0th row, column: slice after the input context, passed to the 'generate'\n",
    "        #... since 'output' starts with the context\n",
    "        new_text = decode(new_tokens.tolist())\n",
    "        \n",
    "        print(new_text, flush=True, end='')\n",
    "        target_senten += new_text\n",
    "        end_pos = target_senten.find(oparam.ending_char)\n",
    "        if end_pos > -1:\n",
    "            return target_senten[0:end_pos] # end the translation\n",
    "        \n",
    "        context = output # Update the context\n",
    "\n",
    "    print('\\n')    \n",
    "    return target_senten\n",
    "\n",
    "def generate_tokens(max_tokens=500):\n",
    "    \"\"\"\"generate some tokens from the model\"\"\"\n",
    "    tokens_at_once = 5 # generate small number of tokens with looping to avoid the long wait\n",
    "    context = torch.zeros((1, 1), dtype=torch.long, device=device)\n",
    "    print(f\"Generating tokens starting from {context}\\n\")\n",
    "    generated_text = \"\"\n",
    "    for _ in range(0,max_tokens,tokens_at_once):\n",
    "        output = m.generate(context, max_new_tokens=tokens_at_once)#, stopping_criteria=your_stopping_criteria)  #TODO <E> stopping\n",
    "        \n",
    "        new_tokens = output[0, context.shape[1]:]\n",
    "        new_text = decode(new_tokens.tolist())\n",
    "        \n",
    "        print(new_text, flush=True, end=\"\")\n",
    "        generated_text += new_text\n",
    "        \n",
    "        context = output\n",
    "    \n",
    "    return generated_text\n",
    "\n",
    "def evaluate_translation(test_sources=None,test_references=None,full_test_set:str=None):\n",
    "\n",
    "    from translation_eval import compute_translation_metrics\n",
    "\n",
    "    if full_test_set is not None:\n",
    "        test_sources = []\n",
    "        test_references = []\n",
    "        sentence_list = full_test_set.split(sep=oparam.ending_char)\n",
    "        print(f\"sentence_list:{len(sentence_list)}\")\n",
    "        i = 0\n",
    "        for li in sentence_list:\n",
    "            #print(f\"{li}\")\n",
    "            i+=1\n",
    "            if li.find(oparam.seperator_char) > -1:\n",
    "                lii = li.split(oparam.seperator_char)\n",
    "                print(f\"{lii}\")\n",
    "                test_sources.append(lii[0])\n",
    "                test_references.append(lii[1])\n",
    "            else:\n",
    "                print(f\"ignored sentence:{li}\")\n",
    "            if i == 200:\n",
    "                break # stop for 10 and see\n",
    "        print(f\"{test_sources[0]}...{test_sources[-1]}\")\n",
    "        print(f\"{test_references[0]}...{test_references[-1]}\")\n",
    "\n",
    "    if test_sources is None: # dummy data\n",
    "        # sample test data \n",
    "        test_sources = [\n",
    "            \"මම ඔයාට ආදරෙයි\",\n",
    "            \"අද කාලගුණය ලස්සනයි\",\n",
    "            \"මුහුදු තීරයේ යන්න ආසද?\",\n",
    "            # ... add more Sinhala sentences\n",
    "        ]\n",
    "    if test_references is None:\n",
    "        # ... corresponding English references\n",
    "        test_references = [\n",
    "            \"I love you\",\n",
    "            \"The weather is beautiful today\",\n",
    "            \"Do you want to go to the beach?\",    \n",
    "        ] \n",
    "\n",
    "    # Run evaluation using existing 'translate_sin2eng' function\n",
    "    results = compute_translation_metrics(\n",
    "        sources=test_sources,\n",
    "        references=test_references,\n",
    "        translate_func=translate_sin2eng,\n",
    "        batch_size=4,\n",
    "        show_examples=10,\n",
    "        verbose=True\n",
    "    )\n",
    "\n",
    "    # Scores\n",
    "    print(f\"Final ChrF++ : {results['chrf++']}\")\n",
    "    print(f\"Final BLEU   : {results['bleu']}\")\n",
    "\n",
    "#batched evaluation script\n",
    "import json\n",
    "\n",
    "def evaluate_translation_segment(full_test_set, start_idx, batch_size):\n",
    "    # 1. Parse the full string into lists\n",
    "    all_sentences = [li for li in full_test_set.split(sep=oparam.ending_char) if oparam.seperator_char in li]\n",
    "    \n",
    "    # 2. Slice the specific segment\n",
    "    end_idx = min(start_idx + batch_size, len(all_sentences))\n",
    "    segment = all_sentences[start_idx:end_idx]\n",
    "    \n",
    "    if not segment:\n",
    "        print(\"No sentences found in this range.\")\n",
    "        return\n",
    "\n",
    "    test_sources = []\n",
    "    test_references = []\n",
    "    for li in segment:\n",
    "        parts = li.split(oparam.seperator_char)\n",
    "        test_sources.append(parts[0].strip())\n",
    "        test_references.append(parts[1].strip())\n",
    "\n",
    "    print(f\"Evaluating records {start_idx} to {end_idx}...\")\n",
    "\n",
    "    # 3. Run translation and get metrics\n",
    "    from translation_eval import compute_translation_metrics\n",
    "    results = compute_translation_metrics(\n",
    "        sources=test_sources,\n",
    "        references=test_references,\n",
    "        translate_func=translate_sin2eng,\n",
    "        batch_size=1, \n",
    "        show_examples=0,\n",
    "        verbose=False\n",
    "    )\n",
    "\n",
    "    # 4. Save to a unique file\n",
    "    filename = f\"eval_{start_idx+1}_{end_idx}.json\"\n",
    "    output_data = {\n",
    "        \"metadata\": {\"start\": start_idx + 1, \"end\": end_idx, \"count\": len(test_sources)},\n",
    "        \"metrics\": {\"bleu\": results['bleu'], \"chrf++\": results['chrf++']}\n",
    "    }\n",
    "    \n",
    "    with open(filename, \"w\") as f:\n",
    "        json.dump(output_data, f, indent=4)\n",
    "    \n",
    "    print(f\"Saved results to {filename}\")\n",
    "\n",
    "import csv\n",
    "import json\n",
    "import os\n",
    "from datetime import datetime\n",
    "\n",
    "class TrainLogger:\n",
    "    def __init__(self,mod_file_name,max_iter,hypp:hp=hp):\n",
    "        self.mod_file_name = mod_file_name\n",
    "        self.saved_model_path = name_tmpl.replace('name',mod_file_name)\n",
    "        self.max_iters = max_iter\n",
    "        clean_hypp = {k: v for k, v in hypp.__dict__.items() if not k.startswith('__')}\n",
    "        # This is a JSON-structured master dictionary\n",
    "        self.session_data = {\n",
    "            \"config\": {\n",
    "                \"hyper_params\":clean_hypp,\n",
    "                \"max_iterations\": max_iter,\n",
    "                \"model_name\": mod_file_name\n",
    "            },\n",
    "            \"logs\": {\n",
    "                \"iteration\":[],\n",
    "                \"train_loss\":[],\n",
    "                \"val_loss\":[]\n",
    "            }\n",
    "        }\n",
    "\n",
    "    def save_session(self,name=None,bvl=None):\n",
    "        \"\"\"Dumps everything into a final, readable CSV.\"\"\"\n",
    "        timestamp = datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n",
    "        \n",
    "        if name is None:\n",
    "            filename = f\"{self.mod_file_name}_log_{timestamp}.csv\"\n",
    "        else:\n",
    "            filename = filename.replace(\".csv\", name)\n",
    "        try:\n",
    "            with open(filename, 'w', newline='',encoding=\"UTF-8\") as f:\n",
    "                writer = csv.writer(f)\n",
    "                \n",
    "                # Write Config Header\n",
    "                writer.writerow([\"--- CONFIGURATION ---\"])\n",
    "                for key, val in self.session_data[\"config\"].items():\n",
    "                    writer.writerow([key, val])\n",
    "                \n",
    "                writer.writerow([]) # Spacer\n",
    "                \n",
    "                # Write Logs Header\n",
    "                writer.writerow([\"--- TRAINING LOGS ---\"])\n",
    "                headers = list(self.session_data[\"logs\"].keys())\n",
    "                writer.writerow(headers)\n",
    "                \n",
    "                # Transpose lists into rows\n",
    "                rows = zip(*[self.session_data[\"logs\"][h] for h in headers])\n",
    "                writer.writerows(rows)\n",
    "                \n",
    "                if bvl is not None:\n",
    "                    writer.writerow([f\"Best Validation Loss:{bvl}\"])\n",
    "                \n",
    "            print(f\"\\n[SUCCESS] Final report saved to {filename}\")\n",
    "            \n",
    "            # Clean up the temporary safety JSON if it exists\n",
    "            #if os.path.exists(\"safety_backup.json\"):\n",
    "            #    os.remove(\"safety_backup.json\")\n",
    "                \n",
    "        except Exception as e:\n",
    "            print(f\"Error during final save: {e}\")\n",
    "\n",
    "    def safety_save(self):\n",
    "        \"\"\"dumps the current dictionary to prevent data loss.\"\"\"\n",
    "        self.save_session(name=\".backup\")\n",
    "\n",
    "def training_loop_breaker():\n",
    "    if os.path.exists(\"training_loop_breaker\"):\n",
    "        print(\"\\nFound 'training_loop_breaker'!!!\")\n",
    "        try:\n",
    "            with open(\"training_loop_breaker\",mode=\"r\") as f:\n",
    "                s = f.read()\n",
    "                if s is not None:\n",
    "                    step = int(s.strip())\n",
    "                    print(f\"Training will stop at {step}\\n\")\n",
    "                    os.remove(\"training_loop_breaker\")\n",
    "                    return step\n",
    "        except Exception as e:\n",
    "            print(f\"\\nError 'training_loop_breaker':{e}\")\n",
    "    return int('inf')\n",
    "\n",
    "\n",
    "\n",
    "if do_inference_only:\n",
    "    generate_tokens(max_tokens=max_tokens)\n",
    "elif do_translate_only:\n",
    "    translate_sin2eng(sinh_sentence)\n",
    "elif do_evaluate_only:\n",
    "    print(\"Start_Index Batch_Size?\",flush=True)\n",
    "    ui2 = str(create_input(\"?\", \"string\"))\n",
    "    if ui2.find(\" \") > -1: #has a space\n",
    "        start_idx_batch_size = ui2.split(\" \")\n",
    "        sidx = int(start_idx_batch_size[0])\n",
    "        bsize = int(start_idx_batch_size[1])\n",
    "        evaluate_translation_segment(full_test_set=test_data,start_idx=sidx,batch_size=bsize)\n",
    "    else:\n",
    "        evaluate_translation(full_test_set=test_data) # Full test dataset\n",
    "\n",
    "else: # Train the model\n",
    "    tot_iter_at_start = tot_iter  # Keep track of starting point\n",
    "    print(f\"Training starting for {max_iters} iterations\")\n",
    "    print(f\"To pass the train set {len(train_data)/(hp.batch_size*hp.block_size)} steps\")\n",
    "    print(f\"Create 'training_loop_breaker' to stop at a step\")\n",
    "    tlb = int('inf')\n",
    "    anim = True\n",
    "    logger = TrainLogger('bigram',max_iters)\n",
    "    logger.session_data[\"logs\"][\"iteration\"].append(-1)\n",
    "    logger.session_data[\"logs\"][\"train_loss\"].append(float(-1))\n",
    "    logger.session_data[\"logs\"][\"val_loss\"].append(float(best_val_loss))\n",
    "    # Error : Triton installation is needed\n",
    "    # model = torch.compile(model)   # if version > pytorch 2.0\n",
    "    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1000, gamma=0.5)# half the lr by 1k steps\n",
    "    for iter in range(max_iters):\n",
    "        print('0',end='') if anim else print(f\"\\b\\b\\b\\b{iter}\\r\",end='') \n",
    "        # every once in a while evaluate the loss on train and val sets\n",
    "        if iter % eval_interval == 0 or iter == max_iters - 1:\n",
    "            print('0.0',end='',flush=True) if anim else print('\\b\\b\\b',end='')\n",
    "            losses = estimate_loss()\n",
    "            current_val_loss = losses['val']\n",
    "            current_step = tot_iter_at_start + iter\n",
    "            tloss = losses['train']\n",
    "            print(f\"\\nstep {iter}: tot {current_step} train loss {tloss:.4f}, val loss {current_val_loss:.4f}\")\n",
    "            \n",
    "            # Check if this is the best model we've seen so far\n",
    "            if current_val_loss < best_val_loss:\n",
    "                best_val_loss = current_val_loss\n",
    "                print(f\"--> Better model saving as {'bigram'}...\",end='')\n",
    "                save_model(model, optimizer, current_step, current_val_loss)\n",
    "            else:\n",
    "                print(f\"--> Val loss did not improve (Best: {best_val_loss:.4f})\",end='')\n",
    "            # logging structure\n",
    "            logger.session_data[\"logs\"][\"iteration\"].append(iter)\n",
    "            logger.session_data[\"logs\"][\"train_loss\"].append(float(tloss.item()))\n",
    "            logger.session_data[\"logs\"][\"val_loss\"].append(float(current_val_loss))\n",
    "            logger.safety_save()   \n",
    "            print('\\r')\n",
    "            tlb = training_loop_breaker()\n",
    "            if iter > tlb:\n",
    "                break\n",
    "\n",
    "        print('1',end='') if anim else print('\\b\\b',end='')\n",
    "        # sample a batch of data\n",
    "        xb, yb = get_batch('train')\n",
    "        print('2',end='') if anim else print('\\b\\b',end='')\n",
    "        # evaluate the loss\n",
    "        #logits, loss = model(xb, yb)\n",
    "        # Pass the sep_id so the loss ignores the Sinhala portion\n",
    "        logits, loss = model(xb, yb, sep_id=stoi['<S>'])\n",
    "\n",
    "        print('3',end='') if anim else print('\\b\\b',end='')\n",
    "        optimizer.zero_grad(set_to_none=True)\n",
    "        print('4',end='') if anim else print('\\b\\b',end='')\n",
    "        loss.backward()\n",
    "        print('5',end='') if anim else print('\\b\\b',end='')\n",
    "        optimizer.step()\n",
    "        print('6',end='') if anim else print('\\b\\b',end='')\n",
    "        scheduler.step()\n",
    "        print('7',end='',flush=True) if anim else print('\\b\\b\\b\\b\\b\\b',end='',flush=True)\n",
    "        anim=not(anim)\n",
    "        \n",
    "\n",
    "\n",
    "    tot_iter+=max_iters\n",
    "    logger.save_session(best_val_loss)\n",
    "    print(f\"\\nTraining End: steps tot {tot_iter}\") # train loss {losses['train']:.4f}, val loss {losses['val']:.4f}\")\n",
    "    if m_id in locals(): # assuming you stored the user's input ID in m_id\n",
    "        update_model_registry(m_id, best_val_loss)\n",
    "\n",
    "\"\"\"\n",
    "jupyter nbconvert --to python the_notebook.ipynb --no-prompt\n",
    "creates a the_notebook.py file, add --no-prompt for a clean script\n",
    "\n",
    "t=1.0:[මට දෙක වෙනකල් ගෙදර ඉන්න ඕන]=>[ mantarge get ssenay and bad time and you know mircant for a bar dog now die  ]<E>m\n",
    "t=0.1:[මට දෙක වෙනකල් ගෙදර ඉන්න ඕන]=>[ what s a good to be see the way i want to see to do the say  ]\n",
    "t=0.3:[ඔබ බ්රැන්ඩි ටිකක් කැමතිද ස්තුතියි]=>\n",
    "https://tinyllm.org/\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "1"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}