gpt.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras import layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Checking if gpu is being used\n",
    "#tf.debugging.set_log_device_placement(True)\n",
    "\n",
    "# Create some tensors\n",
    "#a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])\n",
    "#b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])\n",
    "#c = tf.matmul(a, b)\n",
    "\n",
    "#print(c)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# hyperparameters\n",
    "batch_size = 64\n",
    "block_size = 256\n",
    "max_iters = 5000\n",
    "eval_interval = 500\n",
    "learning_rate = 3e-4\n",
    "eval_iters = 200\n",
    "n_embd = 384\n",
    "n_head = 6\n",
    "n_layer = 6\n",
    "dropout = 0.2\n",
    "\n",
    "# Set random seed\n",
    "tf.random.set_seed(1337)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Loading data\n",
    "with open('input.txt', 'r', encoding='utf-8') as f:\n",
    "    text = f.read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create character mappings\n",
    "chars = sorted(list(set(text)))\n",
    "vocab_size = len(chars)\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]\n",
    "decode = lambda l: ''.join([itos[i] for i in l])\n",
    "\n",
    "# Train and test splits\n",
    "data = tf.constant(encode(text), dtype=tf.int64)\n",
    "n = int(0.9 * len(data))\n",
    "train_data = data[:n]\n",
    "val_data = data[n:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Data loading into batches\n",
    "def get_batch(split):\n",
    "    data_split = train_data if split == 'train' else val_data\n",
    "    ix = tf.random.uniform(shape=(batch_size,), maxval=len(data_split) - block_size, dtype=tf.int32)\n",
    "    x = tf.stack([data_split[i:i+block_size] for i in ix])\n",
    "    y = tf.stack([data_split[i+1:i+block_size+1] for i in ix])\n",
    "    return x, y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Calculating loss of the model\n",
    "def estimate_loss(model):\n",
    "    out = {}\n",
    "    model.trainable = False\n",
    "    for split in ['train', 'val']:\n",
    "        losses = tf.TensorArray(tf.float32, size=eval_iters)\n",
    "        for k in range(eval_iters):\n",
    "            X, Y = get_batch(split)\n",
    "            logits, loss = model(X, Y)\n",
    "            losses = losses.write(k, loss)\n",
    "        out[split] = losses.stack().numpy().mean()\n",
    "    model.trainable = True\n",
    "    return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Head(tf.keras.layers.Layer):\n",
    "    \"\"\" one head of self-attention \"\"\"\n",
    "\n",
    "    def __init__(self, head_size):\n",
    "        super(Head, self).__init__()\n",
    "        self.key = tf.keras.layers.Dense(head_size, use_bias=False)\n",
    "        self.query = tf.keras.layers.Dense(head_size, use_bias=False)\n",
    "        self.value = tf.keras.layers.Dense(head_size, use_bias=False)\n",
    "\n",
    "        tril = tf.linalg.band_part(tf.ones((block_size, block_size)), -1, 0)\n",
    "        self.tril = tf.constant(tril)\n",
    "\n",
    "        self.dropout = tf.keras.layers.Dropout(dropout)\n",
    "\n",
    "    def call(self, x):\n",
    "        # input of size (batch, time-step, channels)\n",
    "        # output of size (batch, time-step, head size)\n",
    "        B, T, C = x.shape\n",
    "        k = self.key(x)   # (B, T, hs)\n",
    "        q = self.query(x) # (B, T, hs)\n",
    "\n",
    "        # compute attention scores (\"affinities\")\n",
    "        wei = tf.matmul(q, tf.transpose(k, perm=[0, 2, 1])) * tf.math.rsqrt(tf.cast(k.shape[-1], tf.float32))  # (B, T, T)\n",
    "        wei = tf.where(self.tril[:T, :T] == 0, float('-inf'), wei)  # (B, T, T)\n",
    "        wei = tf.nn.softmax(wei, axis=-1)  # (B, T, T)\n",
    "        wei = self.dropout(wei)\n",
    "\n",
    "        # perform the weighted aggregation of the values\n",
    "        v = self.value(x)  # (B, T, hs)\n",
    "        out = tf.matmul(wei, v)  # (B, T, T) @ (B, T, hs) -> (B, T, hs)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadAttention(tf.keras.layers.Layer):\n",
    "    \"\"\" multiple heads of self-attention in parallel \"\"\"\n",
    "\n",
    "    def __init__(self, num_heads, head_size):\n",
    "        super(MultiHeadAttention, self).__init__()\n",
    "        self.heads = [Head(head_size) for _ in range(num_heads)]\n",
    "        self.proj = tf.keras.layers.Dense(n_embd)\n",
    "        self.dropout = tf.keras.layers.Dropout(dropout)\n",
    "\n",
    "    def call(self, x):\n",
    "        out = tf.concat([h(x) for h in self.heads], axis=-1)\n",
    "        out = self.dropout(self.proj(out))\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FeedForward(tf.keras.layers.Layer):\n",
    "    \"\"\" a simple linear layer followed by a non-linearity \"\"\"\n",
    "\n",
    "    def __init__(self, n_embd):\n",
    "        super(FeedForward, self).__init__()\n",
    "        self.net = tf.keras.Sequential([\n",
    "            tf.keras.layers.Dense(4 * n_embd),\n",
    "            tf.keras.layers.ReLU(),\n",
    "            tf.keras.layers.Dense(n_embd),\n",
    "            tf.keras.layers.Dropout(dropout),\n",
    "        ])\n",
    "\n",
    "    def call(self, x):\n",
    "        return self.net(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Block(tf.keras.layers.Layer):\n",
    "    \"\"\" Transformer block: communication followed by computation \"\"\"\n",
    "\n",
    "    def __init__(self, n_embd, n_head):\n",
    "        super(Block, self).__init__()\n",
    "        head_size = n_embd // n_head\n",
    "        self.sa = MultiHeadAttention(n_head, head_size)\n",
    "        self.ffwd = FeedForward(n_embd)\n",
    "        self.ln1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
    "        self.ln2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
    "\n",
    "    def call(self, x):\n",
    "        x = x + self.sa(self.ln1(x))\n",
    "        x = x + self.ffwd(self.ln2(x))\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GPTLanguageModel(tf.keras.Model):\n",
    "\n",
    "    def __init__(self):\n",
    "        super(GPTLanguageModel, self).__init__()\n",
    "        # each token directly reads off the logits for the next token from a lookup table\n",
    "        self.token_embedding_table = tf.keras.layers.Embedding(vocab_size, n_embd)\n",
    "        self.position_embedding_table = tf.keras.layers.Embedding(block_size, n_embd)\n",
    "        self.blocks = tf.keras.Sequential([Block(n_embd, n_head=n_head) for _ in range(n_layer)])\n",
    "        self.ln_f = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
    "        self.lm_head = tf.keras.layers.Dense(vocab_size, kernel_initializer='normal', bias_initializer='zeros')\n",
    "\n",
    "    def call(self, idx, targets=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(tf.range(T, dtype=tf.float32))  # (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 = tf.reshape(logits, (B * T, C))\n",
    "            targets = tf.reshape(targets, (B * T,))\n",
    "            loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)(targets, logits)\n",
    "\n",
    "        return logits, loss\n",
    "\n",
    "    def generate(self, idx, max_new_tokens):\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 softmax to get probabilities\n",
    "            probs = tf.nn.softmax(logits, axis=-1)  # (B, C)\n",
    "            # sample from the distribution\n",
    "            idx_next = tf.random.categorical(tf.math.log(probs), num_samples=1, dtype=tf.int64)  # (B, 1)\n",
    "            # append sampled index to the running sequence\n",
    "            idx = tf.concat([idx, idx_next], axis=1)  # (B, T+1)\n",
    "        return idx\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initializing model with pre-trained weights. Use this if you don't want to re-train the model.\n",
    "#model = GPTLanguageModel()\n",
    "#dummy_input = tf.constant([[0]], dtype=tf.int32)  # Example input, adjust shape as needed\n",
    "#model(dummy_input)\n",
    "#model.load_weights('gpt_model_weights.h5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "step 0: train loss 4.5158, val loss 4.5177\n",
      "step 500: train loss 1.9006, val loss 2.0083\n",
      "step 1000: train loss 1.4417, val loss 1.6584\n",
      "step 1500: train loss 1.2854, val loss 1.5992\n",
      "step 2000: train loss 1.1676, val loss 1.5936\n",
      "step 2500: train loss 1.0419, val loss 1.6674\n",
      "step 3000: train loss 0.9076, val loss 1.8094\n",
      "step 3500: train loss 0.7525, val loss 2.0218\n",
      "step 4000: train loss 0.6012, val loss 2.3162\n",
      "step 4500: train loss 0.4598, val loss 2.6565\n",
      "step 4999: train loss 0.3497, val loss 2.9876\n"
     ]
    }
   ],
   "source": [
    "#Training the model. GPU is recommended for training.\n",
    "\n",
    "model = GPTLanguageModel()\n",
    "optimizer = tf.keras.optimizers.Adam(learning_rate)\n",
    "\n",
    "for iter in range(max_iters):\n",
    "\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",
    "        losses = estimate_loss(model)\n",
    "        print(f\"step {iter}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}\")\n",
    "\n",
    "    # sample a batch of data\n",
    "    xb, yb = get_batch('train')\n",
    "\n",
    "    # evaluate the loss\n",
    "    with tf.GradientTape() as tape:\n",
    "        logits, loss = model(xb, yb)\n",
    "\n",
    "    grads = tape.gradient(loss, model.trainable_variables)\n",
    "    optimizer.apply_gradients(zip(grads, model.trainable_variables))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Now keeps.\n",
      "Can I know should thee were trans--I protest,\n",
      "To betwixt the Samart's the mutine.\n",
      "\n",
      "CAMILLO:\n",
      "Ha, madam!\n",
      "Sir, you!\n",
      "You pitiff now, but you are worth aboards,\n",
      "Betwixt the right of your ox adversaries,\n",
      "Or let our suddenly in all severaltius free\n",
      "Than Bolingbroke to England. Mercutio,\n",
      "Ever justice with his praisence, he was proud;\n",
      "When she departed by his fortune like a greer,\n",
      "And in the gentle king fair hateful man.\n",
      "Farewell; so old Cominius, away; I rather,\n",
      "To you are therefore be behold\n"
     ]
    }
   ],
   "source": [
    "# generate from the model\n",
    "context = tf.zeros((1, 1), dtype=tf.int64)\n",
    "generated_sequence = model.generate(context, max_new_tokens=500).numpy()\n",
    "print(decode(generated_sequence[0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save_weights('gpt_model_weights.h5')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}