{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6d6bc54f-2b16-4b0f-be69-957eed5d112f",
   "metadata": {},
   "source": [
    "<table style=\"width:100%\">\n",
    "<tr>\n",
    "<td style=\"vertical-align:middle; text-align:left;\">\n",
    "<font size=\"2\">\n",
    "Supplementary code for the <a href=\"http://mng.bz/orYv\">Build a Large Language Model From Scratch</a> book by <a href=\"https://sebastianraschka.com\">Sebastian Raschka</a><br>\n",
    "<br>Code repository: <a href=\"https://github.com/rasbt/LLMs-from-scratch\">https://github.com/rasbt/LLMs-from-scratch</a>\n",
    "</font>\n",
    "</td>\n",
    "<td style=\"vertical-align:middle; text-align:left;\">\n",
    "<a href=\"http://mng.bz/orYv\"><img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/cover-small.webp\" width=\"100px\"></a>\n",
    "</td>\n",
    "</tr>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72953590-5363-4398-85ce-54bde07f3d8a",
   "metadata": {},
   "source": [
    "# Bonus Code for Chapter 5"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a4ab5ee-e7b9-45d3-a82b-a12bcfc0945a",
   "metadata": {},
   "source": [
    "## Alternative Weight Loading from PyTorch state dicts"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2feea87-49f0-48b9-b925-b8f0dda4096f",
   "metadata": {},
   "source": [
    "- In the main chapter, we loaded the GPT model weights directly from OpenAI\n",
    "- This notebook provides alternative weight loading code to load the model weights from PyTorch state dict files that I created from the original TensorFlow files and uploaded to the [Hugging Face Model Hub](https://huggingface.co/docs/hub/en/models-the-hub) at [https://huggingface.co/rasbt/gpt2-from-scratch-pytorch](https://huggingface.co/rasbt/gpt2-from-scratch-pytorch)\n",
    "- This is conceptually the same as loading weights of a PyTorch model from via the state-dict method described in chapter 5:\n",
    "\n",
    "```python\n",
    "state_dict = torch.load(\"model_state_dict.pth\")\n",
    "model.load_state_dict(state_dict) \n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3f9fbb2-3e39-41ee-8a08-58ba0434a8f3",
   "metadata": {},
   "source": [
    "### Choose model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b0467eff-b43c-4a38-93e8-5ed87a5fc2b1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch version: 2.6.0\n"
     ]
    }
   ],
   "source": [
    "from importlib.metadata import version\n",
    "\n",
    "pkgs = [\"torch\"]\n",
    "for p in pkgs:\n",
    "    print(f\"{p} version: {version(p)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9ea9b1bc-7881-46ad-9555-27a9cf23faa7",
   "metadata": {},
   "outputs": [],
   "source": [
    "BASE_CONFIG = {\n",
    "    \"vocab_size\": 50257,    # Vocabulary size\n",
    "    \"context_length\": 1024, # Context length\n",
    "    \"drop_rate\": 0.0,       # Dropout rate\n",
    "    \"qkv_bias\": True        # Query-key-value bias\n",
    "}\n",
    "\n",
    "model_configs = {\n",
    "    \"gpt2-small (124M)\": {\"emb_dim\": 768, \"n_layers\": 12, \"n_heads\": 12},\n",
    "    \"gpt2-medium (355M)\": {\"emb_dim\": 1024, \"n_layers\": 24, \"n_heads\": 16},\n",
    "    \"gpt2-large (774M)\": {\"emb_dim\": 1280, \"n_layers\": 36, \"n_heads\": 20},\n",
    "    \"gpt2-xl (1558M)\": {\"emb_dim\": 1600, \"n_layers\": 48, \"n_heads\": 25},\n",
    "}\n",
    "\n",
    "\n",
    "CHOOSE_MODEL = \"gpt2-small (124M)\"\n",
    "BASE_CONFIG.update(model_configs[CHOOSE_MODEL])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d78fc2b0-ba27-4aff-8aa3-bc6e04fca69d",
   "metadata": {},
   "source": [
    "### Download file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ca224672-a0f7-4b39-9bc9-19ddde69487b",
   "metadata": {},
   "outputs": [],
   "source": [
    "file_name = \"gpt2-small-124M.pth\"\n",
    "# file_name = \"gpt2-medium-355M.pth\"\n",
    "# file_name = \"gpt2-large-774M.pth\"\n",
    "# file_name = \"gpt2-xl-1558M.pth\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e7b22375-6fac-4e90-9063-daa4de86c778",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloaded to gpt2-small-124M.pth\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import urllib.request\n",
    "\n",
    "url = f\"https://huggingface.co/rasbt/gpt2-from-scratch-pytorch/resolve/main/{file_name}\"\n",
    "\n",
    "if not os.path.exists(file_name):\n",
    "    urllib.request.urlretrieve(url, file_name)\n",
    "    print(f\"Downloaded to {file_name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e61f0990-74cf-4b6d-85e5-4c7d0554db32",
   "metadata": {},
   "source": [
    "### Load weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "cda44d37-92c0-4c19-a70a-15711513afce",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from llms_from_scratch.ch04 import GPTModel\n",
    "# For llms_from_scratch installation instructions, see:\n",
    "# https://github.com/rasbt/LLMs-from-scratch/tree/main/pkg\n",
    "\n",
    "\n",
    "gpt = GPTModel(BASE_CONFIG)\n",
    "gpt.load_state_dict(torch.load(file_name, weights_only=True))\n",
    "gpt.eval()\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "gpt.to(device);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0297fc4-11dc-4093-922f-dcaf85a75344",
   "metadata": {},
   "source": [
    "### Generate text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4ddd0d51-3ade-4890-9bab-d63f141d095f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output text:\n",
      " Every effort moves forward, but it's not enough.\n",
      "\n",
      "\"I'm not going to sit here and say, 'I'm not going to do this,'\n"
     ]
    }
   ],
   "source": [
    "import tiktoken\n",
    "from llms_from_scratch.ch05 import generate, text_to_token_ids, token_ids_to_text\n",
    "\n",
    "\n",
    "torch.manual_seed(123)\n",
    "\n",
    "tokenizer = tiktoken.get_encoding(\"gpt2\")\n",
    "\n",
    "token_ids = generate(\n",
    "    model=gpt.to(device),\n",
    "    idx=text_to_token_ids(\"Every effort moves\", tokenizer).to(device),\n",
    "    max_new_tokens=30,\n",
    "    context_size=BASE_CONFIG[\"context_length\"],\n",
    "    top_k=1,\n",
    "    temperature=1.0\n",
    ")\n",
    "\n",
    "print(\"Output text:\\n\", token_ids_to_text(token_ids, tokenizer))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa4a7912-ae51-4786-8ef4-42bd53682932",
   "metadata": {},
   "source": [
    "## Alternative safetensors file"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f774001-9cda-4b1f-88c5-ef99786a612b",
   "metadata": {},
   "source": [
    "- In addition, the [https://huggingface.co/rasbt/gpt2-from-scratch-pytorch](https://huggingface.co/rasbt/gpt2-from-scratch-pytorch) repository contains so-called `.safetensors` versions of the state dicts\n",
    "- The appeal of `.safetensors` files lies in their secure design, as they only store tensor data and avoid the execution of potentially malicious code during loading\n",
    "- In newer versions of PyTorch (e.g., 2.0 and newer), a `weights_only=True` argument can be used with `torch.load` (e.g., `torch.load(\"model_state_dict.pth\", weights_only=True)`) to improve safety by skipping the execution of code and loading only the weights (this is now enabled by default in PyTorch 2.6 and newer); so in that case loading the weights from the state dict files should not be a concern (anymore)\n",
    "- However, the code block below briefly shows how to load the model from these `.safetensor` files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c0a4fd86-4119-4a94-ae5e-13fb60d198bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "file_name = \"gpt2-small-124M.safetensors\"\n",
    "# file_name = \"gpt2-medium-355M.safetensors\"\n",
    "# file_name = \"gpt2-large-774M.safetensors\"\n",
    "# file_name = \"gpt2-xl-1558M.safetensors\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "20f96c2e-3469-47fb-bad3-e9173a1f1ba3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloaded to gpt2-small-124M.safetensors\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import urllib.request\n",
    "\n",
    "url = f\"https://huggingface.co/rasbt/gpt2-from-scratch-pytorch/resolve/main/{file_name}\"\n",
    "\n",
    "if not os.path.exists(file_name):\n",
    "    urllib.request.urlretrieve(url, file_name)\n",
    "    print(f\"Downloaded to {file_name}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d16a69b3-9bb4-42f8-8e4f-cc62a1a1a083",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load file\n",
    "\n",
    "from safetensors.torch import load_file\n",
    "\n",
    "gpt = GPTModel(BASE_CONFIG)\n",
    "gpt.load_state_dict(load_file(file_name))\n",
    "gpt.eval();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "352e57f7-8d82-4c12-900c-03e41bc9de58",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output text:\n",
      " Every effort moves forward, but it's not enough.\n",
      "\n",
      "\"I'm not going to sit here and say, 'I'm not going to do this,'\n"
     ]
    }
   ],
   "source": [
    "token_ids = generate(\n",
    "    model=gpt.to(device),\n",
    "    idx=text_to_token_ids(\"Every effort moves\", tokenizer).to(device),\n",
    "    max_new_tokens=30,\n",
    "    context_size=BASE_CONFIG[\"context_length\"],\n",
    "    top_k=1,\n",
    "    temperature=1.0\n",
    ")\n",
    "\n",
    "print(\"Output text:\\n\", token_ids_to_text(token_ids, tokenizer))"
   ]
  }
 ],
 "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.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}