experiments.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Python311\\Lib\\site-packages\\tqdm\\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "using device: cuda\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "from transformers import AutoTokenizer\n",
    "# from transformers import GPTNeoXForCausalLM\n",
    "from gpt_neox_mlkv import GPTNeoXConfig, GPTNeoXForCausalLM\n",
    "# from transformers import GPTNeoXConfig, GPTNeoXForCausalLM\n",
    "# check for cuda\n",
    "import torch\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(\"using device:\", device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "GPTNeoXForCausalLM(\n",
       "  (gpt_neox): GPTNeoXModel(\n",
       "    (embed_in): Embedding(50304, 768)\n",
       "    (layers): ModuleList(\n",
       "      (0): GPTNeoXLayer(\n",
       "        (input_layernorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (post_attention_layernorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (attention): GPTNeoXAttention(\n",
       "          (rotary_emb): RotaryEmbedding()\n",
       "          (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (key): KeyValueMLP(\n",
       "            (dense_in): Linear(in_features=768, out_features=4259, bias=True)\n",
       "            (dense_out): Linear(in_features=4259, out_features=64, bias=True)\n",
       "            (act): GELUActivation()\n",
       "          )\n",
       "          (value): KeyValueMLP(\n",
       "            (dense_in): Linear(in_features=768, out_features=4259, bias=True)\n",
       "            (dense_out): Linear(in_features=4259, out_features=64, bias=True)\n",
       "            (act): GELUActivation()\n",
       "          )\n",
       "          (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "        )\n",
       "        (mlp): GPTNeoXMLP(\n",
       "          (dense_h_to_4h): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          (dense_4h_to_h): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          (act): GELUActivation()\n",
       "        )\n",
       "      )\n",
       "      (1-5): 5 x GPTNeoXLayer(\n",
       "        (input_layernorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (post_attention_layernorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (attention): GPTNeoXAttention(\n",
       "          (rotary_emb): RotaryEmbedding()\n",
       "          (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "        )\n",
       "        (mlp): GPTNeoXMLP(\n",
       "          (dense_h_to_4h): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          (dense_4h_to_h): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          (act): GELUActivation()\n",
       "        )\n",
       "      )\n",
       "      (6): GPTNeoXLayer(\n",
       "        (input_layernorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (post_attention_layernorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (attention): GPTNeoXAttention(\n",
       "          (rotary_emb): RotaryEmbedding()\n",
       "          (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (key): KeyValueMLP(\n",
       "            (dense_in): Linear(in_features=768, out_features=4259, bias=True)\n",
       "            (dense_out): Linear(in_features=4259, out_features=64, bias=True)\n",
       "            (act): GELUActivation()\n",
       "          )\n",
       "          (value): KeyValueMLP(\n",
       "            (dense_in): Linear(in_features=768, out_features=4259, bias=True)\n",
       "            (dense_out): Linear(in_features=4259, out_features=64, bias=True)\n",
       "            (act): GELUActivation()\n",
       "          )\n",
       "          (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "        )\n",
       "        (mlp): GPTNeoXMLP(\n",
       "          (dense_h_to_4h): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          (dense_4h_to_h): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          (act): GELUActivation()\n",
       "        )\n",
       "      )\n",
       "      (7-11): 5 x GPTNeoXLayer(\n",
       "        (input_layernorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (post_attention_layernorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "        (attention): GPTNeoXAttention(\n",
       "          (rotary_emb): RotaryEmbedding()\n",
       "          (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "        )\n",
       "        (mlp): GPTNeoXMLP(\n",
       "          (dense_h_to_4h): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          (dense_4h_to_h): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          (act): GELUActivation()\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (final_layer_norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "  )\n",
       "  (embed_out): Linear(in_features=768, out_features=50304, bias=False)\n",
       ")"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# load model\n",
    "# tokenizer = AutoTokenizer.from_pretrained('../pythia-160m-deduped_mlkv_12_12')\n",
    "tokenizer = AutoTokenizer.from_pretrained('../pythia-70m-deduped')\n",
    "# config = GPTNeoXConfig.from_pretrained('../pythia-160m-deduped')\n",
    "# model = GPTNeoXForCausalLM.from_pretrained('../pythia-160m-deduped_mlkv_mlp_4_3')\n",
    "# model = GPTNeoXForCausalLM.from_pretrained('../pythia-70m-deduped')\n",
    "# tokenizer = AutoTokenizer.from_pretrained('EleutherAI/pythia-160m-deduped')\n",
    "# model = GPTNeoXForCausalLM.from_pretrained('../pythia-160m-deduped')\n",
    "model = GPTNeoXForCausalLM.from_pretrained('faridlazuarda/pythia-160m-mlkv-mlp-2-b12-g2-v1', revision='master')\n",
    "# model = GPTNeoXForCausalLM.from_pretrained('test')\n",
    "# model = GPTNeoXForCausalLM.from_pretrained('pythia-70m-fixed/checkpoint-12000')\n",
    "# model = GPTNeoXForCausalLM.from_pretrained('pythia-160m-mlkv-48/checkpoint-10536')\n",
    "model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
      "Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GreedySearchDecoderOnlyOutput(sequences=tensor([[  510, 13642,   273,  3448,  3210,   281,  7114,  1014,  8307,   273,\n",
      "         25866,   273,  3602,  1160,  1896,   407,   253,  4122,  7529, 12729,\n",
      "         39707, 10336,   556,  3977,   281,   253,  6054,   273,  1781,  3448,\n",
      "          3210,   313,  2293, 12822,    10,   347,   253,   268,  2966,  6929,\n",
      "           273,  3626,  3448,  5162,   285,  3448, 17032,   393,  4904,   445,\n",
      "         41766,    92,    81,  1714,   938,  1423, 20246,   314,  7165, 43308,\n",
      "           432, 17170,  1375,    14,  1171,    14,   783,    14,   435,   275,\n",
      "          2710, 49602, 12319,   432,  4028,  2127,   281,  1846,  3282,   285,\n",
      "         14720,    13, 21708, 12822,   452,   671,   644, 18329, 48994,   387,\n",
      "          4311,   323,  4893,   824,   347,  7489, 18128,   434, 24180,    40,\n",
      "          5736,   390,   616,  8990,   273,  2710,  3210,   393,  4904,   445,\n",
      "         41766,    92,  5758,  2284,   938,  1508,    72,   431,    21,  7165,\n",
      "          1292,  9417,  6494,    14,  7527,  3210,   281,  9790,   273,  6383,\n",
      "          4916,   247,  2201,  5691,    13, 10568,  7863, 11897,   281, 42514,\n",
      "           253,  3058, 22667, 10806,    15, 43308,   432,   247,  1355,  2408,\n",
      "         10827,   323, 13782,    13,  3541,   310, 45891,  3058,   323,   767,\n",
      "          2022,  1841,    15, 35346,   253,  1566, 13461,    13,   534,   476,\n",
      "           320,  1781,   533,  3297,  4229,  2378, 18329,    15,   380,   625,\n",
      "          8664,  3541, 10630,   310,   253,  2234,    14,  2877, 11556,   313,\n",
      "            44,    55, 11556,   582,   534,  4428,   253,  7141,  1396,   569,\n",
      "           273,  4116, 11911,   387,   512,  2045,   673,  5018,    15,   831,\n",
      "           310,   253,   906,   273, 39793,  1006,   800, 17032,    27, 21761,\n",
      "           403,  4561,   581,   407,   581,    13,  6753,    14,  1747,   560,\n",
      "          1242,    15,  1916,  3693, 15279,   294,    14,   681, 10340,   273,\n",
      "           841,  1396,   569,    13,   597,   403,  7141,   275,  3541,   281,\n",
      "           320,   908,   323, 11365,  2852, 21761,    15,  2058,  1355, 11498,\n",
      "            13,   611,    55, 11556,   476,   320,   440, 37411,    13,   285,\n",
      "          3541,   310,  6571,  3058,   323,   253, 13461,    15,  1292,  3192,\n",
      "         37692,    14, 14840,    35,   347,   247,  1781,  4311,  1650,    13,\n",
      "           697, 20105,  6494,  3602,  2430, 28325,  6920,   273,  3541,    15,\n",
      "          9110,   672,   275, 46382,   247,  3425,  2978,   273,  1384,  2385,\n",
      "           285, 12672,   247, 14604,   273, 12842, 14649,   387,  2378,    13,\n",
      "           253,   611,    55, 11556,   476,  1379,   598,   281, 49305,  6920,\n",
      "           273,  3541,   393,  4904,   445, 41766,    92,   965,    86,   938,\n",
      "          1508,  1026,   739,   263, 43442,  7165,   496,   436,  1083,    13,\n",
      "           253, 11556,   310,  1264,  2069,  4067,   685,   253,  1566, 13461,\n",
      "            13,  2568,  1014,   840,    13,   611,    55, 11556,   476,  1335,\n",
      "          1756, 39450,   275,  5886,   281,   253,  3425,  2978,    13, 14604,\n",
      "          1979,    13,   285,  1566,  7877,    15, 30688,    28,  4976,    27,\n",
      "           253, 13642,   273,  3448,  3210,   281,  7114,  1014,  8307,   273,\n",
      "          9790,   273,  3602,    15,   187,   187,   510, 13642,   273,  3448,\n",
      "          3210,   281,  7114,  1014,  8307,   273,  9790,   273,  3602,   310,\n",
      "           247,  2234,  4735,   273,   253,  1566,    15,   187,   187,   510,\n",
      "         13642,   273,  3448,  3210,   281,  7114,  1014,  8307,   273,  9790,\n",
      "           273,  3602,   310,   247,  2234,  4735,   273,   253,  1566,    15,\n",
      "           187,   187,   510, 13642,   273,  3448,  3210,   281,  7114,  1014,\n",
      "          8307,   273,  9790,   273,  3602,   310,   247,  2234,  4735,   273,\n",
      "           253,  1566,    15,   187,   187,   510, 13642,   273,  3448,  3210,\n",
      "           281,  7114,  1014,  8307,   273,  9790,   273,  3602,   310,   247,\n",
      "          2234,  4735,   273,   253,  1566,    15,   187,   187,   510, 13642]],\n",
      "       device='cuda:0'), scores=None, attentions=None, hidden_states=None)\n",
      "Output:\n",
      "----------------------------------------------------------------------------------------------------\n",
      "The scaling of language models to tens even hundreds of billions of parameters made possible by the highly parallelizable transformer architecture has led to the rise of large language models (LLMs) as the pinnacle of natural language processing and language inference \\parencite{pope2022efficiently}. Aside from achieving state-of-the-art in various benchmarks ranging from writing code to common sense and reasoning, LLMs have also been deployed massively at scale for applications such as OpenAI's ChatGPT or their API of various models \\parencite{openai2023gpt4}. But serving billion-scale models to millions of customers becomes a major challenge, requiring massive compute to fulfil the needed latency constraints. Aside from a small amount reserved for computation, memory is excessively needed for two main things. Firstly the model weights, which can be large but stay fixed once deployed. The more concerning memory consumer is the key-value cache (KV cache), which contains the stored activations of attention modules at all previous time steps. This is the result of optimizing generative inference: tokens are generated one by one, auto-regressively. To avoid unnecessary re-computation of these activations, they are stored in memory to be used for generating future tokens. At small scales, KV cache can be unassuming, and memory is mostly needed for the weights. But taking OPT-175B as a large scale example, its 175 billion parameters require 325GB of memory. Yet when inferencing a sequence length of 2048 and computing a batch of 128 generations at once, the KV cache can take up to 950GB of memory \\parencite{liu2023scissorhands}. In this case, the cache is three times larger than the model weights, yet even then, KV cache can still grow indefinitely in relation to the sequence length, batch size, and model dimension. TL;DR: the scaling of language models to tens even hundreds of millions of parameters.\n",
      "\n",
      "The scaling of language models to tens even hundreds of millions of parameters is a key feature of the model.\n",
      "\n",
      "The scaling of language models to tens even hundreds of millions of parameters is a key feature of the model.\n",
      "\n",
      "The scaling of language models to tens even hundreds of millions of parameters is a key feature of the model.\n",
      "\n",
      "The scaling of language models to tens even hundreds of millions of parameters is a key feature of the model.\n",
      "\n",
      "The scaling\n"
     ]
    }
   ],
   "source": [
    "prompt = \"\"\"The scaling of language models to tens even hundreds of billions of parameters made possible by the highly parallelizable transformer architecture has led to the rise of large language models (LLMs) as the pinnacle of natural language processing and language inference \\parencite{pope2022efficiently}. Aside from achieving state-of-the-art in various benchmarks ranging from writing code to common sense and reasoning, LLMs have also been deployed massively at scale for applications such as OpenAI's ChatGPT or their API of various models \\parencite{openai2023gpt4}. But serving billion-scale models to millions of customers becomes a major challenge, requiring massive compute to fulfil the needed latency constraints. Aside from a small amount reserved for computation, memory is excessively needed for two main things. Firstly the model weights, which can be large but stay fixed once deployed. The more concerning memory consumer is the key-value cache (KV cache), which contains the stored activations of attention modules at all previous time steps. This is the result of optimizing generative inference: tokens are generated one by one, auto-regressively. To avoid unnecessary re-computation of these activations, they are stored in memory to be used for generating future tokens. At small scales, KV cache can be unassuming, and memory is mostly needed for the weights. But taking OPT-175B as a large scale example, its 175 billion parameters require 325GB of memory. Yet when inferencing a sequence length of 2048 and computing a batch of 128 generations at once, the KV cache can take up to 950GB of memory \\parencite{liu2023scissorhands}. In this case, the cache is three times larger than the model weights, yet even then, KV cache can still grow indefinitely in relation to the sequence length, batch size, and model dimension. TL;DR:\"\"\"\n",
    "# prompt = \"\"\"Wherever you are, I'll always make you smile. Wherever you are, I'm always by your side. Whatever you say, I\"\"\"\n",
    "# prompt = \"\"\"Whether you're in the main venue at Montevideo, at one of our several satellite locations, or participating remotely from anywhere in the world, you and your team will leverage technology and data to\"\"\"\n",
    "input_ids = tokenizer.encode(prompt, return_tensors='pt')\n",
    "greedy_output = model.generate(input_ids.to(device), max_length=500, return_dict_in_generate=True)\n",
    "print(greedy_output)\n",
    "print(\"Output:\\n\" + 100 * '-')\n",
    "print(tokenizer.decode(greedy_output.sequences[0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "parameter count: 162339980\n"
     ]
    }
   ],
   "source": [
    "# get parameter count\n",
    "print(\"parameter count:\", sum(p.numel() for p in model.parameters() if p.requires_grad))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "memory usage: 1762 MiB\n"
     ]
    }
   ],
   "source": [
    "import subprocess\n",
    "import re\n",
    "smi = subprocess.check_output('nvidia-smi')\n",
    "# extract the memory usage with regex. example:\" 2467MiB / 12288MiB \"\n",
    "mem = int(re.search(r'(\\d+)MiB / (\\d+)MiB', str(smi)).group(1))\n",
    "print(\"memory usage:\", mem, \"MiB\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "embed_in.weight torch.Size([50304, 768])\n",
      "layers.0.input_layernorm.weight torch.Size([768])\n",
      "layers.0.input_layernorm.bias torch.Size([768])\n",
      "layers.0.post_attention_layernorm.weight torch.Size([768])\n",
      "layers.0.post_attention_layernorm.bias torch.Size([768])\n",
      "layers.0.attention.query.weight torch.Size([768, 768])\n",
      "layers.0.attention.query.bias torch.Size([768])\n",
      "layers.0.attention.key.dense_in.weight torch.Size([4259, 768])\n",
      "layers.0.attention.key.dense_in.bias torch.Size([4259])\n",
      "layers.0.attention.key.dense_out.weight torch.Size([64, 4259])\n",
      "layers.0.attention.key.dense_out.bias torch.Size([64])\n",
      "layers.0.attention.value.dense_in.weight torch.Size([4259, 768])\n",
      "layers.0.attention.value.dense_in.bias torch.Size([4259])\n",
      "layers.0.attention.value.dense_out.weight torch.Size([64, 4259])\n",
      "layers.0.attention.value.dense_out.bias torch.Size([64])\n",
      "layers.0.attention.dense.weight torch.Size([768, 768])\n",
      "layers.0.attention.dense.bias torch.Size([768])\n",
      "layers.0.mlp.dense_h_to_4h.weight torch.Size([3072, 768])\n",
      "layers.0.mlp.dense_h_to_4h.bias torch.Size([3072])\n",
      "layers.0.mlp.dense_4h_to_h.weight torch.Size([768, 3072])\n",
      "layers.0.mlp.dense_4h_to_h.bias torch.Size([768])\n",
      "layers.1.input_layernorm.weight torch.Size([768])\n",
      "layers.1.input_layernorm.bias torch.Size([768])\n",
      "layers.1.post_attention_layernorm.weight torch.Size([768])\n",
      "layers.1.post_attention_layernorm.bias torch.Size([768])\n",
      "layers.1.attention.query.weight torch.Size([768, 768])\n",
      "layers.1.attention.query.bias torch.Size([768])\n",
      "layers.1.attention.dense.weight torch.Size([768, 768])\n",
      "layers.1.attention.dense.bias torch.Size([768])\n",
      "layers.1.mlp.dense_h_to_4h.weight torch.Size([3072, 768])\n",
      "layers.1.mlp.dense_h_to_4h.bias torch.Size([3072])\n",
      "layers.1.mlp.dense_4h_to_h.weight torch.Size([768, 3072])\n",
      "layers.1.mlp.dense_4h_to_h.bias torch.Size([768])\n",
      "layers.2.input_layernorm.weight torch.Size([768])\n",
      "layers.2.input_layernorm.bias torch.Size([768])\n",
      "layers.2.post_attention_layernorm.weight torch.Size([768])\n",
      "layers.2.post_attention_layernorm.bias torch.Size([768])\n",
      "layers.2.attention.query.weight torch.Size([768, 768])\n",
      "layers.2.attention.query.bias torch.Size([768])\n",
      "layers.2.attention.dense.weight torch.Size([768, 768])\n",
      "layers.2.attention.dense.bias torch.Size([768])\n",
      "layers.2.mlp.dense_h_to_4h.weight torch.Size([3072, 768])\n",
      "layers.2.mlp.dense_h_to_4h.bias torch.Size([3072])\n",
      "layers.2.mlp.dense_4h_to_h.weight torch.Size([768, 3072])\n",
      "layers.2.mlp.dense_4h_to_h.bias torch.Size([768])\n",
      "layers.3.input_layernorm.weight torch.Size([768])\n",
      "layers.3.input_layernorm.bias torch.Size([768])\n",
      "layers.3.post_attention_layernorm.weight torch.Size([768])\n",
      "layers.3.post_attention_layernorm.bias torch.Size([768])\n",
      "layers.3.attention.query.weight torch.Size([768, 768])\n",
      "layers.3.attention.query.bias torch.Size([768])\n",
      "layers.3.attention.dense.weight torch.Size([768, 768])\n",
      "layers.3.attention.dense.bias torch.Size([768])\n",
      "layers.3.mlp.dense_h_to_4h.weight torch.Size([3072, 768])\n",
      "layers.3.mlp.dense_h_to_4h.bias torch.Size([3072])\n",
      "layers.3.mlp.dense_4h_to_h.weight torch.Size([768, 3072])\n",
      "layers.3.mlp.dense_4h_to_h.bias torch.Size([768])\n",
      "layers.4.input_layernorm.weight torch.Size([768])\n",
      "layers.4.input_layernorm.bias torch.Size([768])\n",
      "layers.4.post_attention_layernorm.weight torch.Size([768])\n",
      "layers.4.post_attention_layernorm.bias torch.Size([768])\n",
      "layers.4.attention.query.weight torch.Size([768, 768])\n",
      "layers.4.attention.query.bias torch.Size([768])\n",
      "layers.4.attention.dense.weight torch.Size([768, 768])\n",
      "layers.4.attention.dense.bias torch.Size([768])\n",
      "layers.4.mlp.dense_h_to_4h.weight torch.Size([3072, 768])\n",
      "layers.4.mlp.dense_h_to_4h.bias torch.Size([3072])\n",
      "layers.4.mlp.dense_4h_to_h.weight torch.Size([768, 3072])\n",
      "layers.4.mlp.dense_4h_to_h.bias torch.Size([768])\n",
      "layers.5.input_layernorm.weight torch.Size([768])\n",
      "layers.5.input_layernorm.bias torch.Size([768])\n",
      "layers.5.post_attention_layernorm.weight torch.Size([768])\n",
      "layers.5.post_attention_layernorm.bias torch.Size([768])\n",
      "layers.5.attention.query.weight torch.Size([768, 768])\n",
      "layers.5.attention.query.bias torch.Size([768])\n",
      "layers.5.attention.dense.weight torch.Size([768, 768])\n",
      "layers.5.attention.dense.bias torch.Size([768])\n",
      "layers.5.mlp.dense_h_to_4h.weight torch.Size([3072, 768])\n",
      "layers.5.mlp.dense_h_to_4h.bias torch.Size([3072])\n",
      "layers.5.mlp.dense_4h_to_h.weight torch.Size([768, 3072])\n",
      "layers.5.mlp.dense_4h_to_h.bias torch.Size([768])\n",
      "layers.6.input_layernorm.weight torch.Size([768])\n",
      "layers.6.input_layernorm.bias torch.Size([768])\n",
      "layers.6.post_attention_layernorm.weight torch.Size([768])\n",
      "layers.6.post_attention_layernorm.bias torch.Size([768])\n",
      "layers.6.attention.query.weight torch.Size([768, 768])\n",
      "layers.6.attention.query.bias torch.Size([768])\n",
      "layers.6.attention.key.dense_in.weight torch.Size([4259, 768])\n",
      "layers.6.attention.key.dense_in.bias torch.Size([4259])\n",
      "layers.6.attention.key.dense_out.weight torch.Size([64, 4259])\n",
      "layers.6.attention.key.dense_out.bias torch.Size([64])\n",
      "layers.6.attention.value.dense_in.weight torch.Size([4259, 768])\n",
      "layers.6.attention.value.dense_in.bias torch.Size([4259])\n",
      "layers.6.attention.value.dense_out.weight torch.Size([64, 4259])\n",
      "layers.6.attention.value.dense_out.bias torch.Size([64])\n",
      "layers.6.attention.dense.weight torch.Size([768, 768])\n",
      "layers.6.attention.dense.bias torch.Size([768])\n",
      "layers.6.mlp.dense_h_to_4h.weight torch.Size([3072, 768])\n",
      "layers.6.mlp.dense_h_to_4h.bias torch.Size([3072])\n",
      "layers.6.mlp.dense_4h_to_h.weight torch.Size([768, 3072])\n",
      "layers.6.mlp.dense_4h_to_h.bias torch.Size([768])\n",
      "layers.7.input_layernorm.weight torch.Size([768])\n",
      "layers.7.input_layernorm.bias torch.Size([768])\n",
      "layers.7.post_attention_layernorm.weight torch.Size([768])\n",
      "layers.7.post_attention_layernorm.bias torch.Size([768])\n",
      "layers.7.attention.query.weight torch.Size([768, 768])\n",
      "layers.7.attention.query.bias torch.Size([768])\n",
      "layers.7.attention.dense.weight torch.Size([768, 768])\n",
      "layers.7.attention.dense.bias torch.Size([768])\n",
      "layers.7.mlp.dense_h_to_4h.weight torch.Size([3072, 768])\n",
      "layers.7.mlp.dense_h_to_4h.bias torch.Size([3072])\n",
      "layers.7.mlp.dense_4h_to_h.weight torch.Size([768, 3072])\n",
      "layers.7.mlp.dense_4h_to_h.bias torch.Size([768])\n",
      "layers.8.input_layernorm.weight torch.Size([768])\n",
      "layers.8.input_layernorm.bias torch.Size([768])\n",
      "layers.8.post_attention_layernorm.weight torch.Size([768])\n",
      "layers.8.post_attention_layernorm.bias torch.Size([768])\n",
      "layers.8.attention.query.weight torch.Size([768, 768])\n",
      "layers.8.attention.query.bias torch.Size([768])\n",
      "layers.8.attention.dense.weight torch.Size([768, 768])\n",
      "layers.8.attention.dense.bias torch.Size([768])\n",
      "layers.8.mlp.dense_h_to_4h.weight torch.Size([3072, 768])\n",
      "layers.8.mlp.dense_h_to_4h.bias torch.Size([3072])\n",
      "layers.8.mlp.dense_4h_to_h.weight torch.Size([768, 3072])\n",
      "layers.8.mlp.dense_4h_to_h.bias torch.Size([768])\n",
      "layers.9.input_layernorm.weight torch.Size([768])\n",
      "layers.9.input_layernorm.bias torch.Size([768])\n",
      "layers.9.post_attention_layernorm.weight torch.Size([768])\n",
      "layers.9.post_attention_layernorm.bias torch.Size([768])\n",
      "layers.9.attention.query.weight torch.Size([768, 768])\n",
      "layers.9.attention.query.bias torch.Size([768])\n",
      "layers.9.attention.dense.weight torch.Size([768, 768])\n",
      "layers.9.attention.dense.bias torch.Size([768])\n",
      "layers.9.mlp.dense_h_to_4h.weight torch.Size([3072, 768])\n",
      "layers.9.mlp.dense_h_to_4h.bias torch.Size([3072])\n",
      "layers.9.mlp.dense_4h_to_h.weight torch.Size([768, 3072])\n",
      "layers.9.mlp.dense_4h_to_h.bias torch.Size([768])\n",
      "layers.10.input_layernorm.weight torch.Size([768])\n",
      "layers.10.input_layernorm.bias torch.Size([768])\n",
      "layers.10.post_attention_layernorm.weight torch.Size([768])\n",
      "layers.10.post_attention_layernorm.bias torch.Size([768])\n",
      "layers.10.attention.query.weight torch.Size([768, 768])\n",
      "layers.10.attention.query.bias torch.Size([768])\n",
      "layers.10.attention.dense.weight torch.Size([768, 768])\n",
      "layers.10.attention.dense.bias torch.Size([768])\n",
      "layers.10.mlp.dense_h_to_4h.weight torch.Size([3072, 768])\n",
      "layers.10.mlp.dense_h_to_4h.bias torch.Size([3072])\n",
      "layers.10.mlp.dense_4h_to_h.weight torch.Size([768, 3072])\n",
      "layers.10.mlp.dense_4h_to_h.bias torch.Size([768])\n",
      "layers.11.input_layernorm.weight torch.Size([768])\n",
      "layers.11.input_layernorm.bias torch.Size([768])\n",
      "layers.11.post_attention_layernorm.weight torch.Size([768])\n",
      "layers.11.post_attention_layernorm.bias torch.Size([768])\n",
      "layers.11.attention.query.weight torch.Size([768, 768])\n",
      "layers.11.attention.query.bias torch.Size([768])\n",
      "layers.11.attention.dense.weight torch.Size([768, 768])\n",
      "layers.11.attention.dense.bias torch.Size([768])\n",
      "layers.11.mlp.dense_h_to_4h.weight torch.Size([3072, 768])\n",
      "layers.11.mlp.dense_h_to_4h.bias torch.Size([3072])\n",
      "layers.11.mlp.dense_4h_to_h.weight torch.Size([768, 3072])\n",
      "layers.11.mlp.dense_4h_to_h.bias torch.Size([768])\n",
      "final_layer_norm.weight torch.Size([768])\n",
      "final_layer_norm.bias torch.Size([768])\n"
     ]
    }
   ],
   "source": [
    "for name, param in model.base_model.named_parameters():\n",
    "    print(name, param.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "gpt_neox.embed_in.weight\n",
      "gpt_neox.layers.0.input_layernorm.weight\n",
      "gpt_neox.layers.0.input_layernorm.bias\n",
      "gpt_neox.layers.0.post_attention_layernorm.weight\n",
      "gpt_neox.layers.0.post_attention_layernorm.bias\n",
      "gpt_neox.layers.0.attention.bias\n",
      "gpt_neox.layers.0.attention.masked_bias\n",
      "gpt_neox.layers.0.attention.rotary_emb.inv_freq\n",
      "gpt_neox.layers.0.attention.query_key_value.weight\n",
      "gpt_neox.layers.0.attention.query_key_value.bias\n",
      "gpt_neox.layers.0.attention.dense.weight\n",
      "gpt_neox.layers.0.attention.dense.bias\n",
      "gpt_neox.layers.0.mlp.dense_h_to_4h.weight\n",
      "gpt_neox.layers.0.mlp.dense_h_to_4h.bias\n",
      "gpt_neox.layers.0.mlp.dense_4h_to_h.weight\n",
      "gpt_neox.layers.0.mlp.dense_4h_to_h.bias\n",
      "gpt_neox.layers.1.input_layernorm.weight\n",
      "gpt_neox.layers.1.input_layernorm.bias\n",
      "gpt_neox.layers.1.post_attention_layernorm.weight\n",
      "gpt_neox.layers.1.post_attention_layernorm.bias\n",
      "gpt_neox.layers.1.attention.bias\n",
      "gpt_neox.layers.1.attention.masked_bias\n",
      "gpt_neox.layers.1.attention.rotary_emb.inv_freq\n",
      "gpt_neox.layers.1.attention.query_key_value.weight\n",
      "gpt_neox.layers.1.attention.query_key_value.bias\n",
      "gpt_neox.layers.1.attention.dense.weight\n",
      "gpt_neox.layers.1.attention.dense.bias\n",
      "gpt_neox.layers.1.mlp.dense_h_to_4h.weight\n",
      "gpt_neox.layers.1.mlp.dense_h_to_4h.bias\n",
      "gpt_neox.layers.1.mlp.dense_4h_to_h.weight\n",
      "gpt_neox.layers.1.mlp.dense_4h_to_h.bias\n",
      "gpt_neox.layers.2.input_layernorm.weight\n",
      "gpt_neox.layers.2.input_layernorm.bias\n",
      "gpt_neox.layers.2.post_attention_layernorm.weight\n",
      "gpt_neox.layers.2.post_attention_layernorm.bias\n",
      "gpt_neox.layers.2.attention.bias\n",
      "gpt_neox.layers.2.attention.masked_bias\n",
      "gpt_neox.layers.2.attention.rotary_emb.inv_freq\n",
      "gpt_neox.layers.2.attention.query_key_value.weight\n",
      "gpt_neox.layers.2.attention.query_key_value.bias\n",
      "gpt_neox.layers.2.attention.dense.weight\n",
      "gpt_neox.layers.2.attention.dense.bias\n",
      "gpt_neox.layers.2.mlp.dense_h_to_4h.weight\n",
      "gpt_neox.layers.2.mlp.dense_h_to_4h.bias\n",
      "gpt_neox.layers.2.mlp.dense_4h_to_h.weight\n",
      "gpt_neox.layers.2.mlp.dense_4h_to_h.bias\n",
      "gpt_neox.layers.3.input_layernorm.weight\n",
      "gpt_neox.layers.3.input_layernorm.bias\n",
      "gpt_neox.layers.3.post_attention_layernorm.weight\n",
      "gpt_neox.layers.3.post_attention_layernorm.bias\n",
      "gpt_neox.layers.3.attention.bias\n",
      "gpt_neox.layers.3.attention.masked_bias\n",
      "gpt_neox.layers.3.attention.rotary_emb.inv_freq\n",
      "gpt_neox.layers.3.attention.query_key_value.weight\n",
      "gpt_neox.layers.3.attention.query_key_value.bias\n",
      "gpt_neox.layers.3.attention.dense.weight\n",
      "gpt_neox.layers.3.attention.dense.bias\n",
      "gpt_neox.layers.3.mlp.dense_h_to_4h.weight\n",
      "gpt_neox.layers.3.mlp.dense_h_to_4h.bias\n",
      "gpt_neox.layers.3.mlp.dense_4h_to_h.weight\n",
      "gpt_neox.layers.3.mlp.dense_4h_to_h.bias\n",
      "gpt_neox.layers.4.input_layernorm.weight\n",
      "gpt_neox.layers.4.input_layernorm.bias\n",
      "gpt_neox.layers.4.post_attention_layernorm.weight\n",
      "gpt_neox.layers.4.post_attention_layernorm.bias\n",
      "gpt_neox.layers.4.attention.bias\n",
      "gpt_neox.layers.4.attention.masked_bias\n",
      "gpt_neox.layers.4.attention.rotary_emb.inv_freq\n",
      "gpt_neox.layers.4.attention.query_key_value.weight\n",
      "gpt_neox.layers.4.attention.query_key_value.bias\n",
      "gpt_neox.layers.4.attention.dense.weight\n",
      "gpt_neox.layers.4.attention.dense.bias\n",
      "gpt_neox.layers.4.mlp.dense_h_to_4h.weight\n",
      "gpt_neox.layers.4.mlp.dense_h_to_4h.bias\n",
      "gpt_neox.layers.4.mlp.dense_4h_to_h.weight\n",
      "gpt_neox.layers.4.mlp.dense_4h_to_h.bias\n",
      "gpt_neox.layers.5.input_layernorm.weight\n",
      "gpt_neox.layers.5.input_layernorm.bias\n",
      "gpt_neox.layers.5.post_attention_layernorm.weight\n",
      "gpt_neox.layers.5.post_attention_layernorm.bias\n",
      "gpt_neox.layers.5.attention.bias\n",
      "gpt_neox.layers.5.attention.masked_bias\n",
      "gpt_neox.layers.5.attention.rotary_emb.inv_freq\n",
      "gpt_neox.layers.5.attention.query_key_value.weight\n",
      "gpt_neox.layers.5.attention.query_key_value.bias\n",
      "gpt_neox.layers.5.attention.dense.weight\n",
      "gpt_neox.layers.5.attention.dense.bias\n",
      "gpt_neox.layers.5.mlp.dense_h_to_4h.weight\n",
      "gpt_neox.layers.5.mlp.dense_h_to_4h.bias\n",
      "gpt_neox.layers.5.mlp.dense_4h_to_h.weight\n",
      "gpt_neox.layers.5.mlp.dense_4h_to_h.bias\n",
      "gpt_neox.layers.6.input_layernorm.weight\n",
      "gpt_neox.layers.6.input_layernorm.bias\n",
      "gpt_neox.layers.6.post_attention_layernorm.weight\n",
      "gpt_neox.layers.6.post_attention_layernorm.bias\n",
      "gpt_neox.layers.6.attention.bias\n",
      "gpt_neox.layers.6.attention.masked_bias\n",
      "gpt_neox.layers.6.attention.rotary_emb.inv_freq\n",
      "gpt_neox.layers.6.attention.query_key_value.weight\n",
      "gpt_neox.layers.6.attention.query_key_value.bias\n",
      "gpt_neox.layers.6.attention.dense.weight\n",
      "gpt_neox.layers.6.attention.dense.bias\n",
      "gpt_neox.layers.6.mlp.dense_h_to_4h.weight\n",
      "gpt_neox.layers.6.mlp.dense_h_to_4h.bias\n",
      "gpt_neox.layers.6.mlp.dense_4h_to_h.weight\n",
      "gpt_neox.layers.6.mlp.dense_4h_to_h.bias\n",
      "gpt_neox.layers.7.input_layernorm.weight\n",
      "gpt_neox.layers.7.input_layernorm.bias\n",
      "gpt_neox.layers.7.post_attention_layernorm.weight\n",
      "gpt_neox.layers.7.post_attention_layernorm.bias\n",
      "gpt_neox.layers.7.attention.bias\n",
      "gpt_neox.layers.7.attention.masked_bias\n",
      "gpt_neox.layers.7.attention.rotary_emb.inv_freq\n",
      "gpt_neox.layers.7.attention.query_key_value.weight\n",
      "gpt_neox.layers.7.attention.query_key_value.bias\n",
      "gpt_neox.layers.7.attention.dense.weight\n",
      "gpt_neox.layers.7.attention.dense.bias\n",
      "gpt_neox.layers.7.mlp.dense_h_to_4h.weight\n",
      "gpt_neox.layers.7.mlp.dense_h_to_4h.bias\n",
      "gpt_neox.layers.7.mlp.dense_4h_to_h.weight\n",
      "gpt_neox.layers.7.mlp.dense_4h_to_h.bias\n",
      "gpt_neox.layers.8.input_layernorm.weight\n",
      "gpt_neox.layers.8.input_layernorm.bias\n",
      "gpt_neox.layers.8.post_attention_layernorm.weight\n",
      "gpt_neox.layers.8.post_attention_layernorm.bias\n",
      "gpt_neox.layers.8.attention.bias\n",
      "gpt_neox.layers.8.attention.masked_bias\n",
      "gpt_neox.layers.8.attention.rotary_emb.inv_freq\n",
      "gpt_neox.layers.8.attention.query_key_value.weight\n",
      "gpt_neox.layers.8.attention.query_key_value.bias\n",
      "gpt_neox.layers.8.attention.dense.weight\n",
      "gpt_neox.layers.8.attention.dense.bias\n",
      "gpt_neox.layers.8.mlp.dense_h_to_4h.weight\n",
      "gpt_neox.layers.8.mlp.dense_h_to_4h.bias\n",
      "gpt_neox.layers.8.mlp.dense_4h_to_h.weight\n",
      "gpt_neox.layers.8.mlp.dense_4h_to_h.bias\n",
      "gpt_neox.layers.9.input_layernorm.weight\n",
      "gpt_neox.layers.9.input_layernorm.bias\n",
      "gpt_neox.layers.9.post_attention_layernorm.weight\n",
      "gpt_neox.layers.9.post_attention_layernorm.bias\n",
      "gpt_neox.layers.9.attention.bias\n",
      "gpt_neox.layers.9.attention.masked_bias\n",
      "gpt_neox.layers.9.attention.rotary_emb.inv_freq\n",
      "gpt_neox.layers.9.attention.query_key_value.weight\n",
      "gpt_neox.layers.9.attention.query_key_value.bias\n",
      "gpt_neox.layers.9.attention.dense.weight\n",
      "gpt_neox.layers.9.attention.dense.bias\n",
      "gpt_neox.layers.9.mlp.dense_h_to_4h.weight\n",
      "gpt_neox.layers.9.mlp.dense_h_to_4h.bias\n",
      "gpt_neox.layers.9.mlp.dense_4h_to_h.weight\n",
      "gpt_neox.layers.9.mlp.dense_4h_to_h.bias\n",
      "gpt_neox.layers.10.input_layernorm.weight\n",
      "gpt_neox.layers.10.input_layernorm.bias\n",
      "gpt_neox.layers.10.post_attention_layernorm.weight\n",
      "gpt_neox.layers.10.post_attention_layernorm.bias\n",
      "gpt_neox.layers.10.attention.bias\n",
      "gpt_neox.layers.10.attention.masked_bias\n",
      "gpt_neox.layers.10.attention.rotary_emb.inv_freq\n",
      "gpt_neox.layers.10.attention.query_key_value.weight\n",
      "gpt_neox.layers.10.attention.query_key_value.bias\n",
      "gpt_neox.layers.10.attention.dense.weight\n",
      "gpt_neox.layers.10.attention.dense.bias\n",
      "gpt_neox.layers.10.mlp.dense_h_to_4h.weight\n",
      "gpt_neox.layers.10.mlp.dense_h_to_4h.bias\n",
      "gpt_neox.layers.10.mlp.dense_4h_to_h.weight\n",
      "gpt_neox.layers.10.mlp.dense_4h_to_h.bias\n",
      "gpt_neox.layers.11.input_layernorm.weight\n",
      "gpt_neox.layers.11.input_layernorm.bias\n",
      "gpt_neox.layers.11.post_attention_layernorm.weight\n",
      "gpt_neox.layers.11.post_attention_layernorm.bias\n",
      "gpt_neox.layers.11.attention.bias\n",
      "gpt_neox.layers.11.attention.masked_bias\n",
      "gpt_neox.layers.11.attention.rotary_emb.inv_freq\n",
      "gpt_neox.layers.11.attention.query_key_value.weight\n",
      "gpt_neox.layers.11.attention.query_key_value.bias\n",
      "gpt_neox.layers.11.attention.dense.weight\n",
      "gpt_neox.layers.11.attention.dense.bias\n",
      "gpt_neox.layers.11.mlp.dense_h_to_4h.weight\n",
      "gpt_neox.layers.11.mlp.dense_h_to_4h.bias\n",
      "gpt_neox.layers.11.mlp.dense_4h_to_h.weight\n",
      "gpt_neox.layers.11.mlp.dense_4h_to_h.bias\n",
      "gpt_neox.final_layer_norm.weight\n",
      "gpt_neox.final_layer_norm.bias\n",
      "embed_out.weight\n"
     ]
    }
   ],
   "source": [
    "# inspect pytorch_model.bin and see all the keys\n",
    "state_dict = torch.load('../pythia-160m-deduped/pytorch_model.bin', map_location='cpu')\n",
    "for k in state_dict.keys():\n",
    "    print(k)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2304, 768])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "state_dict['gpt_neox.layers.11.attention.query_key_value.weight'].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "Parent directory ../pythia-160m-altered does not exist.",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[9], line 6\u001b[0m\n\u001b[0;32m      4\u001b[0m state_dict[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mgpt_neox.layers.11.attention.query_key_value.weight\u001b[39m\u001b[38;5;124m'\u001b[39m] \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mfrom_numpy(new_weights)\u001b[38;5;241m.\u001b[39mfloat()\n\u001b[0;32m      5\u001b[0m \u001b[38;5;66;03m# save the new weights\u001b[39;00m\n\u001b[1;32m----> 6\u001b[0m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msave\u001b[49m\u001b[43m(\u001b[49m\u001b[43mstate_dict\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43m../pythia-160m-altered/pytorch_model.bin\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32mc:\\Python311\\Lib\\site-packages\\torch\\serialization.py:504\u001b[0m, in \u001b[0;36msave\u001b[1;34m(obj, f, pickle_module, pickle_protocol, _use_new_zipfile_serialization, _disable_byteorder_record)\u001b[0m\n\u001b[0;32m    501\u001b[0m _check_save_filelike(f)\n\u001b[0;32m    503\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m _use_new_zipfile_serialization:\n\u001b[1;32m--> 504\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m \u001b[43m_open_zipfile_writer\u001b[49m\u001b[43m(\u001b[49m\u001b[43mf\u001b[49m\u001b[43m)\u001b[49m \u001b[38;5;28;01mas\u001b[39;00m opened_zipfile:\n\u001b[0;32m    505\u001b[0m         _save(obj, opened_zipfile, pickle_module, pickle_protocol, _disable_byteorder_record)\n\u001b[0;32m    506\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m\n",
      "File \u001b[1;32mc:\\Python311\\Lib\\site-packages\\torch\\serialization.py:378\u001b[0m, in \u001b[0;36m_open_zipfile_writer\u001b[1;34m(name_or_buffer)\u001b[0m\n\u001b[0;32m    376\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m    377\u001b[0m     container \u001b[38;5;241m=\u001b[39m _open_zipfile_writer_buffer\n\u001b[1;32m--> 378\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mcontainer\u001b[49m\u001b[43m(\u001b[49m\u001b[43mname_or_buffer\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32mc:\\Python311\\Lib\\site-packages\\torch\\serialization.py:349\u001b[0m, in \u001b[0;36m_open_zipfile_writer_file.__init__\u001b[1;34m(self, name)\u001b[0m\n\u001b[0;32m    347\u001b[0m     \u001b[38;5;28msuper\u001b[39m()\u001b[38;5;241m.\u001b[39m\u001b[38;5;21m__init__\u001b[39m(torch\u001b[38;5;241m.\u001b[39m_C\u001b[38;5;241m.\u001b[39mPyTorchFileWriter(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfile_stream))\n\u001b[0;32m    348\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m--> 349\u001b[0m     \u001b[38;5;28msuper\u001b[39m()\u001b[38;5;241m.\u001b[39m\u001b[38;5;21m__init__\u001b[39m(\u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_C\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mPyTorchFileWriter\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mname\u001b[49m\u001b[43m)\u001b[49m)\n",
      "\u001b[1;31mRuntimeError\u001b[0m: Parent directory ../pythia-160m-altered does not exist."
     ]
    }
   ],
   "source": [
    "# want to randomize one of the layers\n",
    "shape = state_dict['gpt_neox.layers.11.attention.query_key_value.weight'].shape\n",
    "new_weights = np.random.normal(size=shape)\n",
    "state_dict['gpt_neox.layers.11.attention.query_key_value.weight'] = torch.from_numpy(new_weights).float()\n",
    "# save the new weights\n",
    "torch.save(state_dict, '../pythia-160m-altered/pytorch_model.bin')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# want to randomize one of the layers\n",
    "shape = state_dict['gpt_neox.layers.0.mlp.dense_4h_to_h.weight'].shape\n",
    "new_weights = np.random.normal(size=shape)\n",
    "state_dict['gpt_neox.layers.0.mlp.dense_4h_to_h.weight'] = torch.from_numpy(new_weights).float()\n",
    "# save the new weights\n",
    "torch.save(state_dict, '../pythia-70m-altered/pytorch_model.bin')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# want to randomize one of the layers\n",
    "shape = state_dict['embed_out.weight'].shape\n",
    "new_weights = np.random.normal(size=shape)\n",
    "state_dict['embed_out.weight'] = torch.from_numpy(new_weights).float()\n",
    "# save the new weights\n",
    "torch.save(state_dict, '../pythia-70m-altered/pytorch_model.bin')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "gpt_neox.embed_in.weight True\n",
      "gpt_neox.layers.0.input_layernorm.weight True\n",
      "gpt_neox.layers.0.input_layernorm.bias True\n",
      "gpt_neox.layers.0.post_attention_layernorm.weight True\n",
      "gpt_neox.layers.0.post_attention_layernorm.bias True\n",
      "gpt_neox.layers.0.attention.bias True\n",
      "gpt_neox.layers.0.attention.masked_bias True\n",
      "gpt_neox.layers.0.attention.rotary_emb.inv_freq True\n",
      "gpt_neox.layers.0.attention.query_key_value.weight True\n",
      "gpt_neox.layers.0.attention.query_key_value.bias True\n",
      "gpt_neox.layers.0.attention.dense.weight True\n",
      "gpt_neox.layers.0.attention.dense.bias True\n",
      "gpt_neox.layers.0.mlp.dense_h_to_4h.weight True\n",
      "gpt_neox.layers.0.mlp.dense_h_to_4h.bias True\n",
      "gpt_neox.layers.0.mlp.dense_4h_to_h.weight True\n",
      "gpt_neox.layers.0.mlp.dense_4h_to_h.bias True\n",
      "gpt_neox.layers.1.input_layernorm.weight True\n",
      "gpt_neox.layers.1.input_layernorm.bias True\n",
      "gpt_neox.layers.1.post_attention_layernorm.weight True\n",
      "gpt_neox.layers.1.post_attention_layernorm.bias True\n",
      "gpt_neox.layers.1.attention.bias True\n",
      "gpt_neox.layers.1.attention.masked_bias True\n",
      "gpt_neox.layers.1.attention.rotary_emb.inv_freq True\n",
      "gpt_neox.layers.1.attention.query_key_value.weight True\n",
      "gpt_neox.layers.1.attention.query_key_value.bias True\n",
      "gpt_neox.layers.1.attention.dense.weight True\n",
      "gpt_neox.layers.1.attention.dense.bias True\n",
      "gpt_neox.layers.1.mlp.dense_h_to_4h.weight True\n",
      "gpt_neox.layers.1.mlp.dense_h_to_4h.bias True\n",
      "gpt_neox.layers.1.mlp.dense_4h_to_h.weight True\n",
      "gpt_neox.layers.1.mlp.dense_4h_to_h.bias True\n",
      "gpt_neox.layers.2.input_layernorm.weight True\n",
      "gpt_neox.layers.2.input_layernorm.bias True\n",
      "gpt_neox.layers.2.post_attention_layernorm.weight True\n",
      "gpt_neox.layers.2.post_attention_layernorm.bias True\n",
      "gpt_neox.layers.2.attention.bias True\n",
      "gpt_neox.layers.2.attention.masked_bias True\n",
      "gpt_neox.layers.2.attention.rotary_emb.inv_freq True\n",
      "gpt_neox.layers.2.attention.query_key_value.weight True\n",
      "gpt_neox.layers.2.attention.query_key_value.bias True\n",
      "gpt_neox.layers.2.attention.dense.weight True\n",
      "gpt_neox.layers.2.attention.dense.bias True\n",
      "gpt_neox.layers.2.mlp.dense_h_to_4h.weight True\n",
      "gpt_neox.layers.2.mlp.dense_h_to_4h.bias True\n",
      "gpt_neox.layers.2.mlp.dense_4h_to_h.weight True\n",
      "gpt_neox.layers.2.mlp.dense_4h_to_h.bias True\n",
      "gpt_neox.layers.3.input_layernorm.weight True\n",
      "gpt_neox.layers.3.input_layernorm.bias True\n",
      "gpt_neox.layers.3.post_attention_layernorm.weight True\n",
      "gpt_neox.layers.3.post_attention_layernorm.bias True\n",
      "gpt_neox.layers.3.attention.bias True\n",
      "gpt_neox.layers.3.attention.masked_bias True\n",
      "gpt_neox.layers.3.attention.rotary_emb.inv_freq True\n",
      "gpt_neox.layers.3.attention.query_key_value.weight True\n",
      "gpt_neox.layers.3.attention.query_key_value.bias True\n",
      "gpt_neox.layers.3.attention.dense.weight True\n",
      "gpt_neox.layers.3.attention.dense.bias True\n",
      "gpt_neox.layers.3.mlp.dense_h_to_4h.weight True\n",
      "gpt_neox.layers.3.mlp.dense_h_to_4h.bias True\n",
      "gpt_neox.layers.3.mlp.dense_4h_to_h.weight True\n",
      "gpt_neox.layers.3.mlp.dense_4h_to_h.bias True\n",
      "gpt_neox.layers.4.input_layernorm.weight True\n",
      "gpt_neox.layers.4.input_layernorm.bias True\n",
      "gpt_neox.layers.4.post_attention_layernorm.weight True\n",
      "gpt_neox.layers.4.post_attention_layernorm.bias True\n",
      "gpt_neox.layers.4.attention.bias True\n",
      "gpt_neox.layers.4.attention.masked_bias True\n",
      "gpt_neox.layers.4.attention.rotary_emb.inv_freq True\n",
      "gpt_neox.layers.4.attention.query_key_value.weight True\n",
      "gpt_neox.layers.4.attention.query_key_value.bias True\n",
      "gpt_neox.layers.4.attention.dense.weight True\n",
      "gpt_neox.layers.4.attention.dense.bias True\n",
      "gpt_neox.layers.4.mlp.dense_h_to_4h.weight True\n",
      "gpt_neox.layers.4.mlp.dense_h_to_4h.bias True\n",
      "gpt_neox.layers.4.mlp.dense_4h_to_h.weight True\n",
      "gpt_neox.layers.4.mlp.dense_4h_to_h.bias True\n",
      "gpt_neox.layers.5.input_layernorm.weight True\n",
      "gpt_neox.layers.5.input_layernorm.bias True\n",
      "gpt_neox.layers.5.post_attention_layernorm.weight True\n",
      "gpt_neox.layers.5.post_attention_layernorm.bias True\n",
      "gpt_neox.layers.5.attention.bias True\n",
      "gpt_neox.layers.5.attention.masked_bias True\n",
      "gpt_neox.layers.5.attention.rotary_emb.inv_freq True\n",
      "gpt_neox.layers.5.attention.query_key_value.weight True\n",
      "gpt_neox.layers.5.attention.query_key_value.bias True\n",
      "gpt_neox.layers.5.attention.dense.weight True\n",
      "gpt_neox.layers.5.attention.dense.bias True\n",
      "gpt_neox.layers.5.mlp.dense_h_to_4h.weight True\n",
      "gpt_neox.layers.5.mlp.dense_h_to_4h.bias True\n",
      "gpt_neox.layers.5.mlp.dense_4h_to_h.weight True\n",
      "gpt_neox.layers.5.mlp.dense_4h_to_h.bias True\n",
      "gpt_neox.layers.6.input_layernorm.weight True\n",
      "gpt_neox.layers.6.input_layernorm.bias True\n",
      "gpt_neox.layers.6.post_attention_layernorm.weight True\n",
      "gpt_neox.layers.6.post_attention_layernorm.bias True\n",
      "gpt_neox.layers.6.attention.bias True\n",
      "gpt_neox.layers.6.attention.masked_bias True\n",
      "gpt_neox.layers.6.attention.rotary_emb.inv_freq True\n",
      "gpt_neox.layers.6.attention.query_key_value.weight True\n",
      "gpt_neox.layers.6.attention.query_key_value.bias True\n",
      "gpt_neox.layers.6.attention.dense.weight True\n",
      "gpt_neox.layers.6.attention.dense.bias True\n",
      "gpt_neox.layers.6.mlp.dense_h_to_4h.weight True\n",
      "gpt_neox.layers.6.mlp.dense_h_to_4h.bias True\n",
      "gpt_neox.layers.6.mlp.dense_4h_to_h.weight True\n",
      "gpt_neox.layers.6.mlp.dense_4h_to_h.bias True\n",
      "gpt_neox.layers.7.input_layernorm.weight True\n",
      "gpt_neox.layers.7.input_layernorm.bias True\n",
      "gpt_neox.layers.7.post_attention_layernorm.weight True\n",
      "gpt_neox.layers.7.post_attention_layernorm.bias True\n",
      "gpt_neox.layers.7.attention.bias True\n",
      "gpt_neox.layers.7.attention.masked_bias True\n",
      "gpt_neox.layers.7.attention.rotary_emb.inv_freq True\n",
      "gpt_neox.layers.7.attention.query_key_value.weight True\n",
      "gpt_neox.layers.7.attention.query_key_value.bias True\n",
      "gpt_neox.layers.7.attention.dense.weight True\n",
      "gpt_neox.layers.7.attention.dense.bias True\n",
      "gpt_neox.layers.7.mlp.dense_h_to_4h.weight True\n",
      "gpt_neox.layers.7.mlp.dense_h_to_4h.bias True\n",
      "gpt_neox.layers.7.mlp.dense_4h_to_h.weight True\n",
      "gpt_neox.layers.7.mlp.dense_4h_to_h.bias True\n",
      "gpt_neox.layers.8.input_layernorm.weight True\n",
      "gpt_neox.layers.8.input_layernorm.bias True\n",
      "gpt_neox.layers.8.post_attention_layernorm.weight True\n",
      "gpt_neox.layers.8.post_attention_layernorm.bias True\n",
      "gpt_neox.layers.8.attention.bias True\n",
      "gpt_neox.layers.8.attention.masked_bias True\n",
      "gpt_neox.layers.8.attention.rotary_emb.inv_freq True\n",
      "gpt_neox.layers.8.attention.query_key_value.weight True\n",
      "gpt_neox.layers.8.attention.query_key_value.bias True\n",
      "gpt_neox.layers.8.attention.dense.weight True\n",
      "gpt_neox.layers.8.attention.dense.bias True\n",
      "gpt_neox.layers.8.mlp.dense_h_to_4h.weight True\n",
      "gpt_neox.layers.8.mlp.dense_h_to_4h.bias True\n",
      "gpt_neox.layers.8.mlp.dense_4h_to_h.weight True\n",
      "gpt_neox.layers.8.mlp.dense_4h_to_h.bias True\n",
      "gpt_neox.layers.9.input_layernorm.weight True\n",
      "gpt_neox.layers.9.input_layernorm.bias True\n",
      "gpt_neox.layers.9.post_attention_layernorm.weight True\n",
      "gpt_neox.layers.9.post_attention_layernorm.bias True\n",
      "gpt_neox.layers.9.attention.bias True\n",
      "gpt_neox.layers.9.attention.masked_bias True\n",
      "gpt_neox.layers.9.attention.rotary_emb.inv_freq True\n",
      "gpt_neox.layers.9.attention.query_key_value.weight True\n",
      "gpt_neox.layers.9.attention.query_key_value.bias True\n",
      "gpt_neox.layers.9.attention.dense.weight True\n",
      "gpt_neox.layers.9.attention.dense.bias True\n",
      "gpt_neox.layers.9.mlp.dense_h_to_4h.weight True\n",
      "gpt_neox.layers.9.mlp.dense_h_to_4h.bias True\n",
      "gpt_neox.layers.9.mlp.dense_4h_to_h.weight True\n",
      "gpt_neox.layers.9.mlp.dense_4h_to_h.bias True\n",
      "gpt_neox.layers.10.input_layernorm.weight True\n",
      "gpt_neox.layers.10.input_layernorm.bias True\n",
      "gpt_neox.layers.10.post_attention_layernorm.weight True\n",
      "gpt_neox.layers.10.post_attention_layernorm.bias True\n",
      "gpt_neox.layers.10.attention.bias True\n",
      "gpt_neox.layers.10.attention.masked_bias True\n",
      "gpt_neox.layers.10.attention.rotary_emb.inv_freq True\n",
      "gpt_neox.layers.10.attention.query_key_value.weight True\n",
      "gpt_neox.layers.10.attention.query_key_value.bias True\n",
      "gpt_neox.layers.10.attention.dense.weight True\n",
      "gpt_neox.layers.10.attention.dense.bias True\n",
      "gpt_neox.layers.10.mlp.dense_h_to_4h.weight True\n",
      "gpt_neox.layers.10.mlp.dense_h_to_4h.bias True\n",
      "gpt_neox.layers.10.mlp.dense_4h_to_h.weight True\n",
      "gpt_neox.layers.10.mlp.dense_4h_to_h.bias True\n",
      "gpt_neox.layers.11.input_layernorm.weight True\n",
      "gpt_neox.layers.11.input_layernorm.bias True\n",
      "gpt_neox.layers.11.post_attention_layernorm.weight True\n",
      "gpt_neox.layers.11.post_attention_layernorm.bias True\n",
      "gpt_neox.layers.11.attention.bias True\n",
      "gpt_neox.layers.11.attention.masked_bias True\n",
      "gpt_neox.layers.11.attention.rotary_emb.inv_freq True\n",
      "gpt_neox.layers.11.attention.query_key_value.weight False\n",
      "gpt_neox.layers.11.attention.query_key_value.bias True\n",
      "gpt_neox.layers.11.attention.dense.weight True\n",
      "gpt_neox.layers.11.attention.dense.bias True\n",
      "gpt_neox.layers.11.mlp.dense_h_to_4h.weight True\n",
      "gpt_neox.layers.11.mlp.dense_h_to_4h.bias True\n",
      "gpt_neox.layers.11.mlp.dense_4h_to_h.weight True\n",
      "gpt_neox.layers.11.mlp.dense_4h_to_h.bias True\n",
      "gpt_neox.final_layer_norm.weight True\n",
      "gpt_neox.final_layer_norm.bias True\n",
      "embed_out.weight True\n"
     ]
    }
   ],
   "source": [
    "# compare old and new\n",
    "old_state_dict = torch.load('../pythia-160m-deduped/pytorch_model.bin', map_location='cpu')\n",
    "new_state_dict = torch.load('../pythia-160m-altered/pytorch_model.bin', map_location='cpu')\n",
    "for k in old_state_dict.keys():\n",
    "    print(k, np.allclose(old_state_dict[k].numpy(), new_state_dict[k].numpy()))\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n",
      "C:\\Users\\Khalid Zuhri\\AppData\\Roaming\\Python\\Python311\\site-packages\\transformers\\generation\\configuration_utils.py:389: UserWarning: `do_sample` is set to `False`. However, `temperature` is set to `0.5` -- this flag is only used in sample-based generation modes. You should set `do_sample=True` or unset `temperature`.\n",
      "  warnings.warn(\n",
      "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
      "Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output:\n",
      "----------------------------------------------------------------------------------------------------\n",
      "The scaling of language models to tens even hundreds of billions of parameters made possible by the highly parallelizable transformer architecture has led to the rise of large language models (LLMs) as the pinnacle of natural language processing and language inference \\parencite{pope2022efficiently}. Aside from achieving state-of-the-art in various benchmarks ranging from writing code to common sense and reasoning, LLMs have also been deployed massively at scale for applications such as OpenAI's ChatGPT or their API of various models \\parencite{openai2023gpt4}. But serving billion-scale models to millions of customers becomes a major challenge, requiring massive compute to fulfil the needed latency constraints. Aside from a small amount reserved for computation, memory is excessively needed for two main things. Firstly the model weights, which can be large but stay fixed once deployed. The more concerning memory consumer is the key-value cache (KV cache), which contains the stored activations of attention modules at all previous time steps. This is the result of optimizing generative inference: tokens are generated one by one, auto-regressively. To avoid unnecessary re-computation of these activations, they are stored in memory to be used for generating future tokens. At small scales, KV cache can be unassuming, and memory is mostly needed for the weights. But taking OPT-175B as a large scale example, its 175 billion parameters require 325GB of memory. Yet when inferencing a sequence length of 2048 and computing a batch of 128 generations at once, the KV cache can take up to 950GB of memory \\parencite{liu2023scissorhands}. In this case, the cache is three times larger than the model weights, yet even then, KV cache can still grow indefinitely in relation to the sequence length, batch size, and model dimension. TL;DR: reason equity grace grace reason reason grace reason reason grace reason reason Somal Football graceropic reasonable reasonencephal reasonefulropiceful Somal reel graceestly reason Tru lag equity feelicableropicift love love feel electric grace reason loveefulencephal feelousiftser behindrails feel feel lag behind behindropic behindicable hateser behindefulicable love grace reason love love feel feel reason behindropic favorable behind electric lag again graceefulropic loveropic grace EEG againcels equity reasonefulsible loveropic feelicable reasonencephalropicencephalefulropicropicserropicropicserefulefulefulencephalropicropicserefulefulefulefulefulefulefulencephal electric West behind Notice love behindoietic again loveencephalroll due reason reason behindidelines behind again love behindicul behindousefuloietic Orientestly behindropic favorable favorable electricropic Somalefuloieticoieticicableefulencephalencephalsibleefulencephalencephalencephalencephal EEGoieticefulencephalencephalencephalefuleful EEGencephal EEGsibleefulencephalencephalencephalencephalencephalencephalsibleefulencephalencephalencephalencephalencephalencephalefulefulencephalencephalefulencephalencephalencephalencephalencephalceptiveencephalencephalencephaleful lag equity stomach love lag lag behindeful EEG reason behindropicefulENTIAL love behindencephal reasonropicropicropicropicoieticropicencephalefulencephalencephalencephalencephalencephalencephalceptiveencephal EEG future love behindropiceful behindencephaleful reason EEGropic EEG EEGencephalsibleencephalsibleencephalencephalencephalencephalencephalencephalencephalencephalencephalencephalencephal EEGoieticropic knowestlyropicencephalencephalencephalencephalencephalencephalencephalencephalceptiveeful knoweful loveeful loveeful behindencephal reasonropicoieticropic loveropic loveencephal loveeful hateefulencephal Eastern loveeful lageful behind EEGropicropic loveeful behind electricous Somalropicefulefulefulropicropic Somalousoietic electric loveoietic appeals behindencephal reasoneful reasonropic Western behindoieticropic love attractive love reason attractive attractiveeful behindariantestly attractiveefulencephalroelectric reasonserefulefulencephal EEGropicropicefulefulencephal electricropicropicropicousefuleful EEGsibleropicefulencephalencephal electric attractiveropic love Somal attractive EEG impression lovericular reasoneful behind behind west behindropic EEGift behind behindropicropic Somalropic behindropicicable loveeful behind EEGencephal EEGropicefulropicropic behindefulefulefulefulefulefulefulefulropic loveeful again attractiveefulefuleful opportunity opportunityefuleful blues grace Somal behind attractiveeful appealingefulefulefulefulefulefulencephalefulricular loveoieticefulropic behind EEG EEGropic behindefulefulefulefuleful EEG reasoneful attractiveefulefulropic loveefulrollropic loveeful behindefulropic love unfavorable behindefuleful loveeful loveeful loveitions again behindropicefulilateral west behindefulencephal behindefulencephalwokeoieticeful loveropic reasoneful behind EEGicableefulropic behindefuleful EEG loveeful behindencephal behindsible behindefulefulefulefuleful Orienteful unfavorableableeful loveencephal reasonefulropic behindropicefuleful attractiveropicropicserefulefulefulefulefulefulefulefulefulefulropic reason electricprototype behind behindefulilateral knoweful reason reason behindeful behindefulefulefulefulropic reason reasonrollefulefulefulefulefulencephalencephal reason behindefulencephal Around loveropic reason attractive\n"
     ]
    }
   ],
   "source": [
    "# load model\n",
    "model = None\n",
    "tokenizer = AutoTokenizer.from_pretrained('../pythia-70m-deduped')\n",
    "prompt = \"\"\"The scaling of language models to tens even hundreds of billions of parameters made possible by the highly parallelizable transformer architecture has led to the rise of large language models (LLMs) as the pinnacle of natural language processing and language inference \\parencite{pope2022efficiently}. Aside from achieving state-of-the-art in various benchmarks ranging from writing code to common sense and reasoning, LLMs have also been deployed massively at scale for applications such as OpenAI's ChatGPT or their API of various models \\parencite{openai2023gpt4}. But serving billion-scale models to millions of customers becomes a major challenge, requiring massive compute to fulfil the needed latency constraints. Aside from a small amount reserved for computation, memory is excessively needed for two main things. Firstly the model weights, which can be large but stay fixed once deployed. The more concerning memory consumer is the key-value cache (KV cache), which contains the stored activations of attention modules at all previous time steps. This is the result of optimizing generative inference: tokens are generated one by one, auto-regressively. To avoid unnecessary re-computation of these activations, they are stored in memory to be used for generating future tokens. At small scales, KV cache can be unassuming, and memory is mostly needed for the weights. But taking OPT-175B as a large scale example, its 175 billion parameters require 325GB of memory. Yet when inferencing a sequence length of 2048 and computing a batch of 128 generations at once, the KV cache can take up to 950GB of memory \\parencite{liu2023scissorhands}. In this case, the cache is three times larger than the model weights, yet even then, KV cache can still grow indefinitely in relation to the sequence length, batch size, and model dimension. TL;DR:\"\"\"\n",
    "input_ids = tokenizer.encode(prompt, return_tensors='pt')\n",
    "new_model = GPTNeoXForCausalLM.from_pretrained('../pythia-160m-altered')\n",
    "new_model.to(device)\n",
    "greedy_output = new_model.generate(input_ids.to(device), max_length=1000, temperature=0.5)\n",
    "print(\"Output:\\n\" + 100 * '-')\n",
    "print(tokenizer.decode(greedy_output[0], skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Auburn deputy mayor Salim Mehajer is back in the spotlight after yet another run in with police.\\n\\nControversial Auburn deputy mayor Salim Mehajer wants an apology from police after he was stopped twice by them in less than half an hour while driving his Ferrari.\\n\\nMehajer, who has rarely been out of the headlines since his extravagant traffic-stopping August wedding in western Sydney, is to face court for allegedly driving an unregistered Ferrari without a licence in Punchbowl at about 10.50pm on Saturday.\\n\\nJust 25 minutes later he was stopped again after ignoring police instructions not to drive and was issued with a field court attendance notice, police said.\\n\\nThe 29-year-old said the vehicle did not show up on the police database because it was registered online and his driver\\'s licence was in the car.\\n\\nHe said no charges have been laid.\\n\\n\"It\\'s called media attention,\" Mehajer told AAP on Monday via email.\\n\\n\"I want nothing but an apology.\"\\n\\nAfter Saturday\\'s incident the flashy councillor took his frustration to social media asking police to use taxpayers\\' money more wisely.\\n\\nHe says police were in the wrong a few weeks ago when questioning whether or not his Mercedes G63 was legal.\\n\\n\"My wife was stranded on the road for five hours - as `apparently the vehicle was illegal\\',\" Mehajer posted on Facebook.\\n\\n\"After they wasted three weeks of our lives and trying to tarnish our names; The vehicle was indeed legal.\"\\n\\n\"And now, They are targeting my `Ferrari\\'?\"\\n\\nMehajer urged police to \"go out there and target real criminals\" claiming a vehicle next to him when stopped on Saturday ran a red light.\\n\\n\"The other vehicle right beside me `runs the red light\\' and almost runs down a residential home... But instead they ignore that and stop me,\" he posted on Facebook.\\n\\nHe said his role as Auburn\\'s deputy mayor would not be affected by the recent spate of allegations.\\n\\n\"My online vehicle clearance has showed everyone that I had no wrong doings,\" Mehajer told AAP.\\n\\nLess than two weeks ago the deputy mayor faced Burwood Local court on charges of intimidation for allegedly making threats against the father of a Lindt Cafe siege survivor.\\n\\nFor his latest traffic offence, Mehajer is due to appear at Bankstown local court on November 18.\\n\\nIn August he was fined $220 for shutting down streets in Lidcombe without authorisation for his extravagant wedding, in which he hired four helicopters that landed in a local park, while his fiancee Aysha travelled with a motorcade of motorbikes and luxury cars worth a $50 million.'"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "dataset = load_dataset(\"zaydzuhri/the_pile_tokenized_6k\")\n",
    "dataset = dataset['train']\n",
    "# test decode tokenizer\n",
    "tokenizer.decode(dataset[0]['input_ids'])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}