scratch.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "050b1a81-1747-4a6b-beac-1b4120143a15",
   "metadata": {},
   "outputs": [],
   "source": [
    "# https://github.com/pytorch/pytorch/blob/3cbe7a53a9a1cea2ef2a042f1ab6f7758f7e4d74/torch/csrc/api/include/torch/nn/functional/activation.h#L643"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eed0e657-8ca7-46f8-b98e-ea641cd96b90",
   "metadata": {},
   "outputs": [],
   "source": [
    "# https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1bc23e8-3f2b-4aa2-84bc-f034ef1099e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# https://youtu.be/kCc8FmEb1nY?si=AEN-_b8nxN1nxvDf&t=5248"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6cae412b-f8ff-4a87-8cd6-c06a245a3f63",
   "metadata": {},
   "outputs": [],
   "source": [
    "# https://github.com/karpathy/nanoGPT/issues/399"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aafc433e-2bc5-48b1-9a59-cf0044b01193",
   "metadata": {},
   "outputs": [],
   "source": [
    "# sentencepiece\n",
    "# tiktoken"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8bc012f-2957-4291-9902-5f9862b8ff97",
   "metadata": {},
   "outputs": [],
   "source": [
    "## deprecated\n",
    "\n",
    "# @kfp.dsl.component(base_image=PIPELINE_IMAGE_URI) # TODO: replace with lightweight base_image\n",
    "# def get_configs(bucket_name: str, path_model_config: str, path_tiktoken_config: str) -> dict:\n",
    "#     import pickle\n",
    "    \n",
    "#     from google.cloud import storage\n",
    "\n",
    "\n",
    "#     storage_client = storage.Client()\n",
    "#     bucket = storage_client.get_bucket(bucket_name)\n",
    "    \n",
    "#     blob_model_config = bucket.blob(path_model_config)\n",
    "#     blob_tiktoken_config = bucket.blob(path_tiktoken_config)\n",
    "    \n",
    "#     with blob_model_config.open(mode='rb') as f:\n",
    "#         model_config = pickle.load(f)\n",
    "#     with blob_tiktoken_config.open(mode='rb') as f:\n",
    "#         tiktoken_config = pickle.load(f)\n",
    "\n",
    "#     assert model_config['embed_dim'] % model_config['num_heads'] == 0\n",
    "\n",
    "#     return {'model_config': model_config, 'tiktoken_config': tiktoken_config}\n",
    "\n",
    "\n",
    "# @kfp.dsl.component(base_image=PIPELINE_IMAGE_URI)\n",
    "# def get_data(bucket_name: str, data_path: str) -> str:\n",
    "#     from google.cloud import storage\n",
    "\n",
    "\n",
    "#     bucket = storage_client.get_bucket(bucket_name)\n",
    "#     blob = bucket.blob(storage)\n",
    "\n",
    "#     return blob_data.download_as_text()\n",
    "\n",
    "# https://cloud.google.com/vertex-ai/docs/training/exporting-model-artifacts\n",
    "\n",
    "# @kfp.dsl.container_component\n",
    "# def prep_deployment(artifact_model: Input[Model], artifact_model_mar: Output[Model]):\n",
    "#     # Note: python code glitchy in this container_component\n",
    "#     return kfp.dsl.ContainerSpec(\n",
    "#         image=PIPELINE_IMAGE_URI,\n",
    "#         command=[\n",
    "#             'sh', '-c', '''\n",
    "#             echo $1\\\n",
    "#             && echo $2\\\n",
    "#             && mkdir -p $1\\\n",
    "#             && echo hello world\\ \n",
    "#             && torch-model-archiver --model-name model --version 0.1 --serialized-file $2 --handler deployment_handler.py --export-path $1\n",
    "#             '''\n",
    "#         ],\n",
    "#         args=[str(artifact_model_mar.path), str(artifact_model.path) + '/model.pth']\n",
    "#     )\n",
    "\n",
    "# @kfp.dsl.component(base_image=PIPELINE_IMAGE_URI)\n",
    "# def prep_deployment(artifact_model: Input[Model], artifact_model_mar: Output[Model]):\n",
    "#     import subprocess\n",
    "#     import os\n",
    "#     os.makedirs(artifact_model_mar.path, exist_ok=True)\n",
    "#     cmd = f\"torch-model-archiver\\\n",
    "#         --model-name model\\\n",
    "#         --version 0.1\\\n",
    "#         --serialized-file {str(artifact_model.path) + '/model.pth'}\\\n",
    "#         --handler deployment_handler.py\\\n",
    "#         --export-path {str(artifact_model_mar.path)}\"\n",
    "#     subprocess.run([cmd], shell=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f7fb02d-1855-4cc4-808c-0d74d5014577",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import tiktoken\n",
    "# enc = tiktoken.get_encoding('gpt2')\n",
    "# enc.encode('hi my name is linsu!')\n",
    "# enc.decode([5303, 616, 1438, 318, 300, 1040, 84, 0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff060a1c-0799-43e8-96b2-d4ea570e49b7",
   "metadata": {},
   "source": [
    "![](diagram.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0cf8c377-cd94-49bf-8dbc-7c24a3e5bddc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "554699f6-13ac-4473-bc0c-f66868949498",
   "metadata": {},
   "outputs": [],
   "source": [
    "B, T, E = 4, 8, 2\n",
    "x = torch.randn(B, T, E)\n",
    "xbow = torch.zeros((B, T, E))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "640b2c36-b359-4a16-9325-1fe5b5fef903",
   "metadata": {},
   "outputs": [],
   "source": [
    "# method 1\n",
    "for b in range(B):\n",
    "    for t in range(T):\n",
    "        xprev = x[b,:t+1] # (t, D)\n",
    "        xbow[b, t] = torch.mean(xprev, 0) # (D,)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb26b63d-e04c-4f1b-9891-2e4226a8bdb5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# method 2\n",
    "wei = torch.tril(torch.ones((T, T)))\n",
    "wei = wei / wei.sum(1, keepdim=True)\n",
    "xbow2 = wei @ x # (-, T, T) @ (B, T, C) = (B, T, C)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "adc56176-43db-4baf-b9d8-9e1afcea5818",
   "metadata": {},
   "outputs": [],
   "source": [
    "# method 3\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 # (-, T, T) @ (B, T, C) = (B, T, C)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1bebeeb3-60bb-4875-8c5e-e8cf4f527f30",
   "metadata": {},
   "outputs": [],
   "source": [
    "# method 4: Attention Mechanism\n",
    "head_dim = 16\n",
    "query = nn.Linear(E, head_dim, bias=False)\n",
    "key = nn.Linear(E, head_dim, bias=False)\n",
    "value = nn.Linear(E, head_dim, bias=False)\n",
    "q = query(x) # q = Wx (B, T, 16)\n",
    "k = key(x) # k = Wx (B, T, 16)\n",
    "v = value(x) # v = Wx (B, T, 16)\n",
    "\n",
    "# dot product between two vectors measures similarity\n",
    "# this matrix multiplication dots each vector to every other vector\n",
    "wei = q @ k.transpose(-2, -1) # (B, T, 16) @ (8, 16, T) = (B, T, T)\n",
    "wei = wei * head_dim**(-0.5)\n",
    "\n",
    "tril = torch.tril(torch.ones(T, T)) # optional (for decoder)\n",
    "wei = wei.masked_fill(tril == 0, float('-inf')) # optional (for decoder)\n",
    "wei = F.softmax(wei, dim=-1)\n",
    "out = wei @ v # (B, T, T) @ (B, T, H) = (B, T, H)"
   ]
  }
 ],
 "metadata": {
  "environment": {
   "kernel": "python3",
   "name": ".m113",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/:m113"
  },
  "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}