rama.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModelForSeq2SeqLM\n",
    "import torch\n",
    "torch.set_default_device('mps')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "llama_tok = AutoTokenizer.from_pretrained(\"meta-llama/Llama-2-7b-hf\")\n",
    "llama_tok.pad_token_id = 891"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2cb321a70cd64dd5b3b7feba51aa3545",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "LlamaForCausalLM(\n",
       "  (model): LlamaModel(\n",
       "    (embed_tokens): Embedding(32000, 4096, padding_idx=0)\n",
       "    (layers): ModuleList(\n",
       "      (0-31): 32 x LlamaDecoderLayer(\n",
       "        (self_attn): LlamaAttention(\n",
       "          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (o_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (rotary_emb): LlamaRotaryEmbedding()\n",
       "        )\n",
       "        (mlp): LlamaMLP(\n",
       "          (gate_proj): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "          (down_proj): Linear(in_features=11008, out_features=4096, bias=False)\n",
       "          (up_proj): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "          (act_fn): SiLUActivation()\n",
       "        )\n",
       "        (input_layernorm): LlamaRMSNorm()\n",
       "        (post_attention_layernorm): LlamaRMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (norm): LlamaRMSNorm()\n",
       "  )\n",
       "  (lm_head): Linear(in_features=4096, out_features=32000, bias=False)\n",
       ")"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "llama_mod = AutoModelForCausalLM.from_pretrained(\"meta-llama/Llama-2-7b-hf\")\n",
    "llama_mod.to('mps')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sentence_transformers import SentenceTransformer\n",
    "emb_mod = SentenceTransformer('BAAI/bge-large-en')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# load link_pairs.json\n",
    "import json\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('link_pairs.json', 'r') as f:\n",
    "    link_pairs = json.load(f)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('link_pairs_emb.json', 'r') as f:\n",
    "    link_pairs = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# for each pair, add a third item, the embedding vector\n",
    "for pair in link_pairs:\n",
    "    link_embed = emb_mod.encode(pair['link']).tolist()\n",
    "    pair['lembed'] = link_embed\n",
    "    summary_embeds = []\n",
    "    for summary in pair['summaries']:\n",
    "        summary_embeds.append(emb_mod.encode(summary).tolist())\n",
    "    pair['sembeds'] = summary_embeds\n",
    "\n",
    "# save link_pairs.json\n",
    "with open('link_pairs_emb.json', 'w') as f:\n",
    "    json.dump(link_pairs, f)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import cosine similarity function\n",
    "from sklearn.metrics.pairwise import cosine_similarity\n",
    "\n",
    "# function, given a query and a context (link), return the top 7 most similar contexts plus the original context if it isnt in the top 7\n",
    "def get_top7(query_embed, context):\n",
    "    # get the cosine similarity between the query and all the other contexts\n",
    "    cos_sims = []\n",
    "    for pair in link_pairs:\n",
    "        if pair['link'] != context:\n",
    "            cos_sims.append(cosine_similarity([query_embed], [pair['lembed']]))\n",
    "    # get the top 7 most similar contexts\n",
    "    top7 = []\n",
    "    for i in range(6):\n",
    "        top7.append(link_pairs[cos_sims.index(max(cos_sims))]['link'])\n",
    "        cos_sims[cos_sims.index(max(cos_sims))] = -1\n",
    "    top7.append(context)\n",
    "    return top7\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "628\n",
      "['An experimental study questioning wave-particle duality in quantum mechanics.', ['obscillesk: https://www.tomshardware.com/news/quantum-computing-researchers-achieve-100-million-quantum-operations more quantum shenanigans\\nLink: \"Quantum Computing: Researchers Achieve 100 Million Quantum Operations\" Description: \"That\\'s a lot of processing power being used in five-second workloads.\"\\n\\n\\njcorvinus: I think I\\'m gonna need a bigger R&D budget lol\\n\\n\\nlauren0001: manhattan project for ASI is publicly underway though so idk @_@\\n\\n\\nlauren0001: absolutely\\n\\n\\njcorvinus: I wonder if they\\'ll designate it as ITAR restricted\\n\\n\\nlauren0001: idk about full singularity\\n\\n\\nlauren0001: life extension should become a button to push\\n\\n\\nlauren0001: if we can get to hard superintelligence in the next couple of years after human level, then assuming the org that does this is aligned with humanity personally, then global warming should be possible to solve via unreasonably advanced bioengineering of rapid growth tree species\\n\\n\\nlauren0001: to me\\n\\n\\nlauren0001: alignment is not looking, uh, apolitical\\n\\n\\nlauren0001: you should be worried about who gets this much power at least as much\\n\\n\\nlauren0001: that\\'s small potatoes\\n\\n\\nlauren0001: yeah well\\n\\n\\nlauren0001: which to be clear is not this trivial task\\n\\n\\njcorvinus: Would love to have some really powerful AI systems work on biocompatible material discovery\\n\\n\\nlauren0001: the components are there, putting them together at that scale without getting wasted capacity problems is the issue\\n\\n\\nlauren0001: if you can do many iterations of 100t scale models, though, then you can make it happen\\n\\n\\n', 'obscillesk: https://www.sciencealert.com/physics-breakthrough-as-ai-successfully-controls-plasma-in-nuclear-fusion-experiment/amp\\n\\n\\n', 'Deleted User: https://www.cs.ox.ac.uk/activities/probverifbio/ here\\nLink: \"Probabilistic Verification for Systems Biology\" Description: \"\"\\n\\n\\nDeleted User: But they can be useful still\\n\\n\\nDeleted User: tho id say such proofs are never perfect. In the real world there are always loopholes\\n\\n\\nDeleted User: i think i saw some work from someone here in Oxford about verification of biological systems\\n\\n\\nlauren0001: (probably)\\n\\n\\nlauren0001: not forbidden by physics after all\\n\\n\\nlauren0001: but it should be possible\\n\\n\\nlauren0001: it\\'s not a sub-ASI task\\n\\n\\njcorvinus: oh nice\\n\\n\\nlauren0001: I think you can make a security grade proof of correctness of bio interventions actually\\n\\n\\nlauren0001: I\\'ve been reading up on proving neural networks and proving physics\\n\\n\\njcorvinus: use it in space\\n\\n\\nlauren0001: and I think I can have one\\n\\n\\njcorvinus: ah\\n\\n\\njcorvinus: hmm. Calibration systems of course. Validation systems\\n\\n\\nlauren0001: I\\'d want a security grade proof\\n\\n\\nlauren0001: not in terms of proving it to others, but what would you in particular want to know before you build such a thing\\n\\n\\n', 'lauren0001: https://www.youtube.com/channel/UCZfROMrGuOqrcj_XzHNCarw\\nLink: \"Lauren (has pinned channels)\" Description: \"I make music (rarely). I recommend a few channels, especially Saint Andrewism and Systems Innovation.\"\\n\\n\\ngoldsteel: The prof he bet against in the windmill thing was a real bellend as well.\\n\\n\\nlauren0001: he has \"taster content\" syndrome for sure\\n\\n\\ngoldsteel: His videos have been getting cringe and he keeps targeting dickheads to bait them into bets.\\n\\n\\nlauren0001: he\\'s great, he\\'s not great enough\\n\\n\\nlauren0001: yeah same\\n\\n\\ngoldsteel: I am tired of Derek. I unsubbed.\\n\\n\\ngoldsteel: Oh, yeah. Derek.\\n\\n\\nlauren0001: yeah, Veritasium caused quite a debate by making a video exploring a cool real thing and using lies to children\\n\\n\\ngoldsteel: I do know him, I watched his video recently on measuring the speed of voltage.\\n\\n\\ngoldsteel: Wait.\\n\\n\\ngoldsteel: I don\\'t think so.\\n\\n\\nlauren0001: he does an amazing job grounding things\\n\\n\\nlauren0001: you know about AlphaPhoenix?\\n\\n\\nlauren0001: mhm mhm okay I\\'m with you there then\\n\\n\\ngoldsteel: Contrast Applied Science.\\n\\n\\nlauren0001: i see, yeah, they do act as only fragments in a tutorial\\n\\n\\n', 'lauren0001: this looks like a fun api. I wish there was an open source thing this good https://www.assemblyai.com/features/audio-intelligence\\nLink: \"AssemblyAI Speech-to-Text API | Features\" Description: \"Features. Accurately convert speech to text with powerful AI models. Used by Fortune 500s, startups, and developers. Voted Best API of 2020, and funded by Y Combinator.\"\\n\\n\\n', 'lauren0001: https://youtu.be/xxXlD4e-wTE\\nLink: \"Forget Small ... What About Micro Nuclear Energy?\" Description: \"Revisiting Small Modular Reactors - The Future of Nuclear Energy? Get Surfshark VPN at  and enter promo code UNDECIDED for 83% off and 3 extra months for free! Although nuclear energy is reliable and a sustainable energy source, the nuclear energy debate rages on. It\\'s not considered the go-to solution in the ren...\"\\n\\n\\nyellow0720: :potatowo:\\n\\n\\nyellow0720: 👀\\n\\n\\n', 'goldsteel: Awful article that cites wikipedia about 10 times, but here\\'s the paper: https://www.nature.com/articles/s42005-022-00828-z#Sec5\\nLink: \"Experimental assessment of physical realism in a quantum-controlled...\" Description: \"Communications Physics - The concept of wave-particle duality is at the core of quantum mechanics principles but has seen an element of questioning with recent experiments were overlapping...\"\\n\\n\\n']]\n"
     ]
    }
   ],
   "source": [
    "train_data = []\n",
    "for pair in link_pairs:\n",
    "    for sem, summ in zip(pair['sembeds'], pair['summaries']):\n",
    "        train_data.append([summ, get_top7(sem, pair['link'])])\n",
    "\n",
    "print(len(train_data))\n",
    "print(train_data[0])\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# for each item in train data, shuffle the top7, and then replace each link string in the chat snippet with LINK-(1-8), and append the correct link ID to the end of the summary\n",
    "for i in range(len(train_data)):\n",
    "    # the last item in train_data[i][1] is the correct link context, so we need to track its position\n",
    "    correct_link = train_data[i][1][-1]\n",
    "    # shuffle the top7\n",
    "    random.shuffle(train_data[i][1])\n",
    "    correct_id = 0\n",
    "    for j in range(len(train_data[i][1])):\n",
    "        if train_data[i][1][j] == correct_link:\n",
    "            correct_id = j+1\n",
    "        # add [LINK-1] [LINK-2] ... [LINK-8] before the 'http' in the chat snippet\n",
    "        train_data[i][1][j] = train_data[i][1][j].replace('http', ' [LINK-' + str(j+1) + '] http')\n",
    "\n",
    "    # append the correct link ID to the end of the summary\n",
    "    train_data[i][0] = train_data[i][0] + ' LINK-' + str(correct_id)\n",
    "\n",
    "\n",
    "print(train_data[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "llama_tok.padding_side = 'left'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[  891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,   891,   891,   891,   891,   891,   891,   891,   891,   891,\n",
      "           891,     1,   530, 17986,  6559,  1139,   292, 10742, 29899,  1595,\n",
      "          2512,   868,  2877,   297, 12101,  7208,  1199, 29889, 21724, 29968,\n",
      "         29899, 29896]], device='mps:0')\n",
      "tensor([[891, 891, 891,  ...,  13,  13,  13],\n",
      "        [891, 891, 891,  ...,  13,  13,  13],\n",
      "        [891, 891, 891,  ...,  13,  13,  13],\n",
      "        ...,\n",
      "        [891, 891, 891,  ...,  13,  13,  13],\n",
      "        [891, 891, 891,  ...,  13,  13,  13],\n",
      "        [891, 891, 891,  ...,  13,  13,  13]], device='mps:0')\n"
     ]
    }
   ],
   "source": [
    "tokenized_train = []\n",
    "for i in range(len(train_data)):\n",
    "    item = []\n",
    "    item.append(llama_tok(train_data[i][0], padding='max_length', truncation=True, max_length=512, return_tensors='pt')['input_ids'])\n",
    "    item.append(llama_tok(train_data[i][1], padding='max_length', truncation=True, max_length=512, return_tensors='pt')['input_ids'])\n",
    "    tokenized_train.append(item)\n",
    "\n",
    "print(tokenized_train[0][0])\n",
    "print(tokenized_train[0][1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([7, 512, 4096])\n",
      "torch.Size([7, 512, 4096])\n",
      "torch.Size([7, 512, 4096])\n",
      "torch.Size([7, 512, 4096])\n",
      "torch.Size([7, 512, 4096])\n",
      "torch.Size([7, 512, 4096])\n",
      "torch.Size([7, 512, 4096])\n",
      "HColisolisHColisHCHCHCظظظظظظظظظظظظظظظظظظظظظظظظظظظظظظHCHCantinHCHCHCHCHCHCHCHCHColisHCHCHCHColisolisHCHCHCHCHCLSLSHCLSLSLSHCHCHCHCHCHCHCHCظظظظظظظظظظظظظظHCHCLSLSLSHCessaHCHCHColiseedHCHCeedHCHCHCHCHCHCHCHCHCHCHCHCHCLSLSLSHCLSLSLSLSLSLSLSLSLSLSLSLSLSLSLSLSLSظظظظظظظظظظظظظظHCHCHCHC Victoria Victoria Victoria Victoria VictoriaHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCظظظظظظظظظظظHCظظظظessaLSLSLSLSLSLSLSLSLSLSLStegerHCHCHCHCHCHCHCHCHCHCHCHCHCHCLSLSLSLSHCHCHCHCHCHCظHCHCHCHCHCHCHCershellershellHCLSHCHCHCHCHCLSLSLSHCHCLSLSLSLSLSLSLS Victoria VictoriaHCHC Victoria Victoria Victoria Victoria Victoria Victoria Victoria Victoria Victoria Victoria Victoria Victoria Victoria Victoria Victoria Victoria Victoria Victoria Victoria Victoria Victoria Victoria Victoria Victoria Victoria Victoria Victoria Victoria Victoria Victoria Victoria Victoria VictoriaLSLSLSLSLSLSLSLSLSLSLSLSLS Victoria Victoria Victoria VictoriaLSLSLSLSLSLSLSLSHCLSLSLS VictoriaolisolisolisolisolisolisolisolisolisolisolisolisolisHCHCLSLSershellHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCHCLSHCHCHCHCHCLSLSHCLSLSLSLSLSLS VictoriaHCHCظLSHColisolisolisolisolisolisiaeolisHCHCHCHCHCHCHCHCHCHCHCHColisolisershellHCershellolisolisLSershellHCHCHCHCLSLSHCHCHCHCHCywywHColisHCHCHColisolisolisظHCLSHCHCظظHColisHCLSHCHCantinolisцеHCظظLSLSظHCLS회HCLSЋolis\n",
      "12.033472061157227\n",
      "torch.Size([7, 512, 4096])\n",
      "torch.Size([7, 512, 4096])\n",
      "torch.Size([7, 512, 4096])\n",
      "torch.Size([7, 512, 4096])\n",
      "torch.Size([7, 512, 4096])\n",
      "torch.Size([7, 512, 4096])\n",
      "torch.Size([7, 512, 4096])\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800000; text-decoration-color: #800000\">╭─────────────────────────────── </span><span style=\"color: #800000; text-decoration-color: #800000; font-weight: bold\">Traceback </span><span style=\"color: #bf7f7f; text-decoration-color: #bf7f7f; font-weight: bold\">(most recent call last)</span><span style=\"color: #800000; text-decoration-color: #800000\"> ────────────────────────────────╮</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">&lt;module&gt;</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">42</span>                                                                                   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">39 │   # add a padding token to the beginning of the label</span>                                     <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">40 │   </span>label = torch.cat([torch.tensor([<span style=\"color: #0000ff; text-decoration-color: #0000ff\">0</span>]), label])                                           <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">41 │   </span>inp = torch.cat([tokenized_train[i][<span style=\"color: #0000ff; text-decoration-color: #0000ff\">0</span>], tokenized_train[i][<span style=\"color: #0000ff; text-decoration-color: #0000ff\">1</span>]])                         <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span>42 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   </span>output = model(inp)                                                                     <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">43 │   # argmax and decode</span>                                                                     <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">44 │   </span>stringy = torch.argmax(output, dim=-<span style=\"color: #0000ff; text-decoration-color: #0000ff\">1</span>)                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">45 │   </span>stringy = llama_tok.decode(stringy)                                                     <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">/opt/homebrew/Caskroom/miniconda/base/lib/python3.10/site-packages/torch/nn/modules/</span><span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">module.py</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">15</span> <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #0000ff; text-decoration-color: #0000ff\">18</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">_wrapped_call_impl</span>                                                                         <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1515 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">if</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._compiled_call_impl <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">is</span> <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">not</span> <span style=\"color: #0000ff; text-decoration-color: #0000ff\">None</span>:                                          <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1516 │   │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">return</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._compiled_call_impl(*args, **kwargs)  <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"># type: ignore[misc]</span>        <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1517 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">else</span>:                                                                             <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span>1518 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">return</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._call_impl(*args, **kwargs)                                       <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1519 │   </span>                                                                                      <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1520 │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">def</span> <span style=\"color: #00ff00; text-decoration-color: #00ff00\">_call_impl</span>(<span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>, *args, **kwargs):                                                <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1521 │   │   </span>forward_call = (<span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._slow_forward <span style=\"color: #0000ff; text-decoration-color: #0000ff\">if</span> torch._C._get_tracing_state() <span style=\"color: #0000ff; text-decoration-color: #0000ff\">else</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>.fo  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">/opt/homebrew/Caskroom/miniconda/base/lib/python3.10/site-packages/torch/nn/modules/</span><span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">module.py</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">15</span> <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #0000ff; text-decoration-color: #0000ff\">27</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">_call_impl</span>                                                                                 <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1524 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">if</span> <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">not</span> (<span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._backward_hooks <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._backward_pre_hooks <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._forward_hooks   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1525 │   │   │   │   </span><span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> _global_backward_pre_hooks <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> _global_backward_hooks                   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1526 │   │   │   │   </span><span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> _global_forward_hooks <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> _global_forward_pre_hooks):                   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span>1527 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">return</span> forward_call(*args, **kwargs)                                          <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1528 │   │   </span>                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1529 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">try</span>:                                                                              <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1530 │   │   │   </span>result = <span style=\"color: #0000ff; text-decoration-color: #0000ff\">None</span>                                                                 <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">forward</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">17</span>                                                                                    <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">14 │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">def</span> <span style=\"color: #00ff00; text-decoration-color: #00ff00\">forward</span>(<span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>, x):                                                                   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">15 │   │   </span>x = <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>.emb(x)                                                                     <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">16 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">for</span> i, block <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">in</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">enumerate</span>(<span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>.blocks):                                             <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span>17 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   │   │   </span>x = block(x)[<span style=\"color: #0000ff; text-decoration-color: #0000ff\">0</span>]                                                                 <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">18 │   │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">if</span> i % <span style=\"color: #0000ff; text-decoration-color: #0000ff\">4</span> == <span style=\"color: #0000ff; text-decoration-color: #0000ff\">0</span> <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">and</span> i != <span style=\"color: #0000ff; text-decoration-color: #0000ff\">0</span>:                                                       <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">19 │   │   │   │   # do self attention as normal, but then reshape the batch into [-1, rtr_</span>    <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">20 │   │   │   │   </span>tiled = x[<span style=\"color: #0000ff; text-decoration-color: #0000ff\">0</span>].unsqueeze(<span style=\"color: #0000ff; text-decoration-color: #0000ff\">0</span>).repeat(<span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>.rtr_num, <span style=\"color: #0000ff; text-decoration-color: #0000ff\">1</span>, <span style=\"color: #0000ff; text-decoration-color: #0000ff\">1</span>)                        <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">/opt/homebrew/Caskroom/miniconda/base/lib/python3.10/site-packages/torch/nn/modules/</span><span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">module.py</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">15</span> <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #0000ff; text-decoration-color: #0000ff\">18</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">_wrapped_call_impl</span>                                                                         <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1515 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">if</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._compiled_call_impl <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">is</span> <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">not</span> <span style=\"color: #0000ff; text-decoration-color: #0000ff\">None</span>:                                          <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1516 │   │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">return</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._compiled_call_impl(*args, **kwargs)  <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"># type: ignore[misc]</span>        <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1517 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">else</span>:                                                                             <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span>1518 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">return</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._call_impl(*args, **kwargs)                                       <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1519 │   </span>                                                                                      <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1520 │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">def</span> <span style=\"color: #00ff00; text-decoration-color: #00ff00\">_call_impl</span>(<span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>, *args, **kwargs):                                                <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1521 │   │   </span>forward_call = (<span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._slow_forward <span style=\"color: #0000ff; text-decoration-color: #0000ff\">if</span> torch._C._get_tracing_state() <span style=\"color: #0000ff; text-decoration-color: #0000ff\">else</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>.fo  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">/opt/homebrew/Caskroom/miniconda/base/lib/python3.10/site-packages/torch/nn/modules/</span><span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">module.py</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">15</span> <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #0000ff; text-decoration-color: #0000ff\">27</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">_call_impl</span>                                                                                 <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1524 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">if</span> <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">not</span> (<span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._backward_hooks <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._backward_pre_hooks <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._forward_hooks   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1525 │   │   │   │   </span><span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> _global_backward_pre_hooks <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> _global_backward_hooks                   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1526 │   │   │   │   </span><span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> _global_forward_hooks <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> _global_forward_pre_hooks):                   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span>1527 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">return</span> forward_call(*args, **kwargs)                                          <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1528 │   │   </span>                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1529 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">try</span>:                                                                              <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1530 │   │   │   </span>result = <span style=\"color: #0000ff; text-decoration-color: #0000ff\">None</span>                                                                 <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">/opt/homebrew/Caskroom/miniconda/base/lib/python3.10/site-packages/transformers/models/llama/</span><span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">mod</span> <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">eling_llama.py</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">292</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">forward</span>                                                                    <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">289 │   │   </span>hidden_states = <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>.input_layernorm(hidden_states)                                <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">290 │   │   </span>                                                                                   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">291 │   │   # Self Attention</span>                                                                   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span>292 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   │   </span>hidden_states, self_attn_weights, present_key_value = <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>.self_attn(              <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">293 │   │   │   </span>hidden_states=hidden_states,                                                   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">294 │   │   │   </span>attention_mask=attention_mask,                                                 <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">295 │   │   │   </span>position_ids=position_ids,                                                     <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">/opt/homebrew/Caskroom/miniconda/base/lib/python3.10/site-packages/torch/nn/modules/</span><span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">module.py</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">15</span> <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #0000ff; text-decoration-color: #0000ff\">18</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">_wrapped_call_impl</span>                                                                         <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1515 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">if</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._compiled_call_impl <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">is</span> <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">not</span> <span style=\"color: #0000ff; text-decoration-color: #0000ff\">None</span>:                                          <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1516 │   │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">return</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._compiled_call_impl(*args, **kwargs)  <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"># type: ignore[misc]</span>        <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1517 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">else</span>:                                                                             <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span>1518 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">return</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._call_impl(*args, **kwargs)                                       <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1519 │   </span>                                                                                      <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1520 │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">def</span> <span style=\"color: #00ff00; text-decoration-color: #00ff00\">_call_impl</span>(<span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>, *args, **kwargs):                                                <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1521 │   │   </span>forward_call = (<span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._slow_forward <span style=\"color: #0000ff; text-decoration-color: #0000ff\">if</span> torch._C._get_tracing_state() <span style=\"color: #0000ff; text-decoration-color: #0000ff\">else</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>.fo  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">/opt/homebrew/Caskroom/miniconda/base/lib/python3.10/site-packages/torch/nn/modules/</span><span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">module.py</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">15</span> <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #0000ff; text-decoration-color: #0000ff\">27</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">_call_impl</span>                                                                                 <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1524 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">if</span> <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">not</span> (<span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._backward_hooks <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._backward_pre_hooks <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>._forward_hooks   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1525 │   │   │   │   </span><span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> _global_backward_pre_hooks <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> _global_backward_hooks                   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1526 │   │   │   │   </span><span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> _global_forward_hooks <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> _global_forward_pre_hooks):                   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span>1527 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">return</span> forward_call(*args, **kwargs)                                          <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1528 │   │   </span>                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1529 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">try</span>:                                                                              <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1530 │   │   │   </span>result = <span style=\"color: #0000ff; text-decoration-color: #0000ff\">None</span>                                                                 <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">/opt/homebrew/Caskroom/miniconda/base/lib/python3.10/site-packages/transformers/models/llama/</span><span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">mod</span> <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">eling_llama.py</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">202</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">forward</span>                                                                    <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">199 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">if</span> past_key_value <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">is</span> <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">not</span> <span style=\"color: #0000ff; text-decoration-color: #0000ff\">None</span>:                                                     <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">200 │   │   │   </span>kv_seq_len += past_key_value[<span style=\"color: #0000ff; text-decoration-color: #0000ff\">0</span>].shape[-<span style=\"color: #0000ff; text-decoration-color: #0000ff\">2</span>]                                      <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">201 │   │   </span>cos, sin = <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>.rotary_emb(value_states, seq_len=kv_seq_len)                       <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span>202 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   │   </span>query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, s   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">203 │   │   # [bsz, nh, t, hd]</span>                                                                 <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">204 │   │   </span>                                                                                   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">205 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">if</span> past_key_value <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">is</span> <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">not</span> <span style=\"color: #0000ff; text-decoration-color: #0000ff\">None</span>:                                                     <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">/opt/homebrew/Caskroom/miniconda/base/lib/python3.10/site-packages/transformers/models/llama/</span><span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">mod</span> <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">eling_llama.py</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">136</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">apply_rotary_pos_emb</span>                                                       <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">133 │   </span>sin = sin.squeeze(<span style=\"color: #0000ff; text-decoration-color: #0000ff\">1</span>).squeeze(<span style=\"color: #0000ff; text-decoration-color: #0000ff\">0</span>)  <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"># [seq_len, dim]</span>                                      <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">134 │   </span>cos = cos[position_ids].unsqueeze(<span style=\"color: #0000ff; text-decoration-color: #0000ff\">1</span>)  <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"># [bs, 1, seq_len, dim]</span>                          <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">135 │   </span>sin = sin[position_ids].unsqueeze(<span style=\"color: #0000ff; text-decoration-color: #0000ff\">1</span>)  <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"># [bs, 1, seq_len, dim]</span>                          <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span>136 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   </span>q_embed = (q * cos) + (rotate_half(q) * sin)                                           <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">137 │   </span>k_embed = (k * cos) + (rotate_half(k) * sin)                                           <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">138 │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">return</span> q_embed, k_embed                                                                <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">139 </span>                                                                                           <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">/opt/homebrew/Caskroom/miniconda/base/lib/python3.10/site-packages/torch/utils/</span><span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">_device.py</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">76</span> in  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #00ff00; text-decoration-color: #00ff00\">__torch_function__</span>                                                                               <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">73 │   │   </span>kwargs = kwargs <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">or</span> {}                                                               <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">74 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">if</span> func <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">in</span> _device_constructors() <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">and</span> kwargs.get(<span style=\"color: #808000; text-decoration-color: #808000\">'device'</span>) <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">is</span> <span style=\"color: #0000ff; text-decoration-color: #0000ff\">None</span>:                 <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">75 │   │   │   </span>kwargs[<span style=\"color: #808000; text-decoration-color: #808000\">'device'</span>] = <span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>.device                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span>76 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">return</span> func(*args, **kwargs)                                                        <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">77 </span>                                                                                            <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">78 # NB: This is directly called from C++ in torch/csrc/Device.cpp</span>                             <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">79 </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">def</span> <span style=\"color: #00ff00; text-decoration-color: #00ff00\">device_decorator</span>(device, func):                                                         <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">╰──────────────────────────────────────────────────────────────────────────────────────────────────╯</span>\n",
       "<span style=\"color: #ff0000; text-decoration-color: #ff0000; font-weight: bold\">RuntimeError: </span>MPS backend out of memory <span style=\"font-weight: bold\">(</span>MPS allocated: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">114.43</span> GB, other allocations: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7.91</span> GB, max allowed: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">122.40</span> \n",
       "GB<span style=\"font-weight: bold\">)</span>. Tried to allocate <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">64.00</span> MB on private pool. Use <span style=\"color: #808000; text-decoration-color: #808000\">PYTORCH_MPS_HIGH_WATERMARK_RATIO</span>=<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0</span> to disable upper limit \n",
       "for memory allocations <span style=\"font-weight: bold\">(</span>may cause system failure<span style=\"font-weight: bold\">)</span>.\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[31m╭─\u001b[0m\u001b[31m──────────────────────────────\u001b[0m\u001b[31m \u001b[0m\u001b[1;31mTraceback \u001b[0m\u001b[1;2;31m(most recent call last)\u001b[0m\u001b[31m \u001b[0m\u001b[31m───────────────────────────────\u001b[0m\u001b[31m─╮\u001b[0m\n",
       "\u001b[31m│\u001b[0m in \u001b[92m<module>\u001b[0m:\u001b[94m42\u001b[0m                                                                                   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m39 \u001b[0m\u001b[2m│   \u001b[0m\u001b[2m# add a padding token to the beginning of the label\u001b[0m                                     \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m40 \u001b[0m\u001b[2m│   \u001b[0mlabel = torch.cat([torch.tensor([\u001b[94m0\u001b[0m]), label])                                           \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m41 \u001b[0m\u001b[2m│   \u001b[0minp = torch.cat([tokenized_train[i][\u001b[94m0\u001b[0m], tokenized_train[i][\u001b[94m1\u001b[0m]])                         \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m42 \u001b[2m│   \u001b[0moutput = model(inp)                                                                     \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m43 \u001b[0m\u001b[2m│   \u001b[0m\u001b[2m# argmax and decode\u001b[0m                                                                     \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m44 \u001b[0m\u001b[2m│   \u001b[0mstringy = torch.argmax(output, dim=-\u001b[94m1\u001b[0m)                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m45 \u001b[0m\u001b[2m│   \u001b[0mstringy = llama_tok.decode(stringy)                                                     \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[2;33m/opt/homebrew/Caskroom/miniconda/base/lib/python3.10/site-packages/torch/nn/modules/\u001b[0m\u001b[1;33mmodule.py\u001b[0m:\u001b[94m15\u001b[0m \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[94m18\u001b[0m in \u001b[92m_wrapped_call_impl\u001b[0m                                                                         \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1515 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mif\u001b[0m \u001b[96mself\u001b[0m._compiled_call_impl \u001b[95mis\u001b[0m \u001b[95mnot\u001b[0m \u001b[94mNone\u001b[0m:                                          \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1516 \u001b[0m\u001b[2m│   │   │   \u001b[0m\u001b[94mreturn\u001b[0m \u001b[96mself\u001b[0m._compiled_call_impl(*args, **kwargs)  \u001b[2m# type: ignore[misc]\u001b[0m        \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1517 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94melse\u001b[0m:                                                                             \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m1518 \u001b[2m│   │   │   \u001b[0m\u001b[94mreturn\u001b[0m \u001b[96mself\u001b[0m._call_impl(*args, **kwargs)                                       \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1519 \u001b[0m\u001b[2m│   \u001b[0m                                                                                      \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1520 \u001b[0m\u001b[2m│   \u001b[0m\u001b[94mdef\u001b[0m \u001b[92m_call_impl\u001b[0m(\u001b[96mself\u001b[0m, *args, **kwargs):                                                \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1521 \u001b[0m\u001b[2m│   │   \u001b[0mforward_call = (\u001b[96mself\u001b[0m._slow_forward \u001b[94mif\u001b[0m torch._C._get_tracing_state() \u001b[94melse\u001b[0m \u001b[96mself\u001b[0m.fo  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[2;33m/opt/homebrew/Caskroom/miniconda/base/lib/python3.10/site-packages/torch/nn/modules/\u001b[0m\u001b[1;33mmodule.py\u001b[0m:\u001b[94m15\u001b[0m \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[94m27\u001b[0m in \u001b[92m_call_impl\u001b[0m                                                                                 \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1524 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mif\u001b[0m \u001b[95mnot\u001b[0m (\u001b[96mself\u001b[0m._backward_hooks \u001b[95mor\u001b[0m \u001b[96mself\u001b[0m._backward_pre_hooks \u001b[95mor\u001b[0m \u001b[96mself\u001b[0m._forward_hooks   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1525 \u001b[0m\u001b[2m│   │   │   │   \u001b[0m\u001b[95mor\u001b[0m _global_backward_pre_hooks \u001b[95mor\u001b[0m _global_backward_hooks                   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1526 \u001b[0m\u001b[2m│   │   │   │   \u001b[0m\u001b[95mor\u001b[0m _global_forward_hooks \u001b[95mor\u001b[0m _global_forward_pre_hooks):                   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m1527 \u001b[2m│   │   │   \u001b[0m\u001b[94mreturn\u001b[0m forward_call(*args, **kwargs)                                          \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1528 \u001b[0m\u001b[2m│   │   \u001b[0m                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1529 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mtry\u001b[0m:                                                                              \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1530 \u001b[0m\u001b[2m│   │   │   \u001b[0mresult = \u001b[94mNone\u001b[0m                                                                 \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m in \u001b[92mforward\u001b[0m:\u001b[94m17\u001b[0m                                                                                    \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m14 \u001b[0m\u001b[2m│   \u001b[0m\u001b[94mdef\u001b[0m \u001b[92mforward\u001b[0m(\u001b[96mself\u001b[0m, x):                                                                   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m15 \u001b[0m\u001b[2m│   │   \u001b[0mx = \u001b[96mself\u001b[0m.emb(x)                                                                     \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m16 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mfor\u001b[0m i, block \u001b[95min\u001b[0m \u001b[96menumerate\u001b[0m(\u001b[96mself\u001b[0m.blocks):                                             \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m17 \u001b[2m│   │   │   \u001b[0mx = block(x)[\u001b[94m0\u001b[0m]                                                                 \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m18 \u001b[0m\u001b[2m│   │   │   \u001b[0m\u001b[94mif\u001b[0m i % \u001b[94m4\u001b[0m == \u001b[94m0\u001b[0m \u001b[95mand\u001b[0m i != \u001b[94m0\u001b[0m:                                                       \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m19 \u001b[0m\u001b[2m│   │   │   │   \u001b[0m\u001b[2m# do self attention as normal, but then reshape the batch into [-1, rtr_\u001b[0m    \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m20 \u001b[0m\u001b[2m│   │   │   │   \u001b[0mtiled = x[\u001b[94m0\u001b[0m].unsqueeze(\u001b[94m0\u001b[0m).repeat(\u001b[96mself\u001b[0m.rtr_num, \u001b[94m1\u001b[0m, \u001b[94m1\u001b[0m)                        \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[2;33m/opt/homebrew/Caskroom/miniconda/base/lib/python3.10/site-packages/torch/nn/modules/\u001b[0m\u001b[1;33mmodule.py\u001b[0m:\u001b[94m15\u001b[0m \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[94m18\u001b[0m in \u001b[92m_wrapped_call_impl\u001b[0m                                                                         \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1515 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mif\u001b[0m \u001b[96mself\u001b[0m._compiled_call_impl \u001b[95mis\u001b[0m \u001b[95mnot\u001b[0m \u001b[94mNone\u001b[0m:                                          \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1516 \u001b[0m\u001b[2m│   │   │   \u001b[0m\u001b[94mreturn\u001b[0m \u001b[96mself\u001b[0m._compiled_call_impl(*args, **kwargs)  \u001b[2m# type: ignore[misc]\u001b[0m        \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1517 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94melse\u001b[0m:                                                                             \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m1518 \u001b[2m│   │   │   \u001b[0m\u001b[94mreturn\u001b[0m \u001b[96mself\u001b[0m._call_impl(*args, **kwargs)                                       \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1519 \u001b[0m\u001b[2m│   \u001b[0m                                                                                      \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1520 \u001b[0m\u001b[2m│   \u001b[0m\u001b[94mdef\u001b[0m \u001b[92m_call_impl\u001b[0m(\u001b[96mself\u001b[0m, *args, **kwargs):                                                \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1521 \u001b[0m\u001b[2m│   │   \u001b[0mforward_call = (\u001b[96mself\u001b[0m._slow_forward \u001b[94mif\u001b[0m torch._C._get_tracing_state() \u001b[94melse\u001b[0m \u001b[96mself\u001b[0m.fo  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[2;33m/opt/homebrew/Caskroom/miniconda/base/lib/python3.10/site-packages/torch/nn/modules/\u001b[0m\u001b[1;33mmodule.py\u001b[0m:\u001b[94m15\u001b[0m \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[94m27\u001b[0m in \u001b[92m_call_impl\u001b[0m                                                                                 \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1524 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mif\u001b[0m \u001b[95mnot\u001b[0m (\u001b[96mself\u001b[0m._backward_hooks \u001b[95mor\u001b[0m \u001b[96mself\u001b[0m._backward_pre_hooks \u001b[95mor\u001b[0m \u001b[96mself\u001b[0m._forward_hooks   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1525 \u001b[0m\u001b[2m│   │   │   │   \u001b[0m\u001b[95mor\u001b[0m _global_backward_pre_hooks \u001b[95mor\u001b[0m _global_backward_hooks                   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1526 \u001b[0m\u001b[2m│   │   │   │   \u001b[0m\u001b[95mor\u001b[0m _global_forward_hooks \u001b[95mor\u001b[0m _global_forward_pre_hooks):                   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m1527 \u001b[2m│   │   │   \u001b[0m\u001b[94mreturn\u001b[0m forward_call(*args, **kwargs)                                          \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1528 \u001b[0m\u001b[2m│   │   \u001b[0m                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1529 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mtry\u001b[0m:                                                                              \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1530 \u001b[0m\u001b[2m│   │   │   \u001b[0mresult = \u001b[94mNone\u001b[0m                                                                 \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[2;33m/opt/homebrew/Caskroom/miniconda/base/lib/python3.10/site-packages/transformers/models/llama/\u001b[0m\u001b[1;33mmod\u001b[0m \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[1;33meling_llama.py\u001b[0m:\u001b[94m292\u001b[0m in \u001b[92mforward\u001b[0m                                                                    \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m289 \u001b[0m\u001b[2m│   │   \u001b[0mhidden_states = \u001b[96mself\u001b[0m.input_layernorm(hidden_states)                                \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m290 \u001b[0m\u001b[2m│   │   \u001b[0m                                                                                   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m291 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[2m# Self Attention\u001b[0m                                                                   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m292 \u001b[2m│   │   \u001b[0mhidden_states, self_attn_weights, present_key_value = \u001b[96mself\u001b[0m.self_attn(              \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m293 \u001b[0m\u001b[2m│   │   │   \u001b[0mhidden_states=hidden_states,                                                   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m294 \u001b[0m\u001b[2m│   │   │   \u001b[0mattention_mask=attention_mask,                                                 \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m295 \u001b[0m\u001b[2m│   │   │   \u001b[0mposition_ids=position_ids,                                                     \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[2;33m/opt/homebrew/Caskroom/miniconda/base/lib/python3.10/site-packages/torch/nn/modules/\u001b[0m\u001b[1;33mmodule.py\u001b[0m:\u001b[94m15\u001b[0m \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[94m18\u001b[0m in \u001b[92m_wrapped_call_impl\u001b[0m                                                                         \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1515 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mif\u001b[0m \u001b[96mself\u001b[0m._compiled_call_impl \u001b[95mis\u001b[0m \u001b[95mnot\u001b[0m \u001b[94mNone\u001b[0m:                                          \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1516 \u001b[0m\u001b[2m│   │   │   \u001b[0m\u001b[94mreturn\u001b[0m \u001b[96mself\u001b[0m._compiled_call_impl(*args, **kwargs)  \u001b[2m# type: ignore[misc]\u001b[0m        \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1517 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94melse\u001b[0m:                                                                             \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m1518 \u001b[2m│   │   │   \u001b[0m\u001b[94mreturn\u001b[0m \u001b[96mself\u001b[0m._call_impl(*args, **kwargs)                                       \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1519 \u001b[0m\u001b[2m│   \u001b[0m                                                                                      \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1520 \u001b[0m\u001b[2m│   \u001b[0m\u001b[94mdef\u001b[0m \u001b[92m_call_impl\u001b[0m(\u001b[96mself\u001b[0m, *args, **kwargs):                                                \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1521 \u001b[0m\u001b[2m│   │   \u001b[0mforward_call = (\u001b[96mself\u001b[0m._slow_forward \u001b[94mif\u001b[0m torch._C._get_tracing_state() \u001b[94melse\u001b[0m \u001b[96mself\u001b[0m.fo  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[2;33m/opt/homebrew/Caskroom/miniconda/base/lib/python3.10/site-packages/torch/nn/modules/\u001b[0m\u001b[1;33mmodule.py\u001b[0m:\u001b[94m15\u001b[0m \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[94m27\u001b[0m in \u001b[92m_call_impl\u001b[0m                                                                                 \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1524 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mif\u001b[0m \u001b[95mnot\u001b[0m (\u001b[96mself\u001b[0m._backward_hooks \u001b[95mor\u001b[0m \u001b[96mself\u001b[0m._backward_pre_hooks \u001b[95mor\u001b[0m \u001b[96mself\u001b[0m._forward_hooks   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1525 \u001b[0m\u001b[2m│   │   │   │   \u001b[0m\u001b[95mor\u001b[0m _global_backward_pre_hooks \u001b[95mor\u001b[0m _global_backward_hooks                   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1526 \u001b[0m\u001b[2m│   │   │   │   \u001b[0m\u001b[95mor\u001b[0m _global_forward_hooks \u001b[95mor\u001b[0m _global_forward_pre_hooks):                   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m1527 \u001b[2m│   │   │   \u001b[0m\u001b[94mreturn\u001b[0m forward_call(*args, **kwargs)                                          \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1528 \u001b[0m\u001b[2m│   │   \u001b[0m                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1529 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mtry\u001b[0m:                                                                              \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m1530 \u001b[0m\u001b[2m│   │   │   \u001b[0mresult = \u001b[94mNone\u001b[0m                                                                 \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[2;33m/opt/homebrew/Caskroom/miniconda/base/lib/python3.10/site-packages/transformers/models/llama/\u001b[0m\u001b[1;33mmod\u001b[0m \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[1;33meling_llama.py\u001b[0m:\u001b[94m202\u001b[0m in \u001b[92mforward\u001b[0m                                                                    \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m199 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mif\u001b[0m past_key_value \u001b[95mis\u001b[0m \u001b[95mnot\u001b[0m \u001b[94mNone\u001b[0m:                                                     \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m200 \u001b[0m\u001b[2m│   │   │   \u001b[0mkv_seq_len += past_key_value[\u001b[94m0\u001b[0m].shape[-\u001b[94m2\u001b[0m]                                      \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m201 \u001b[0m\u001b[2m│   │   \u001b[0mcos, sin = \u001b[96mself\u001b[0m.rotary_emb(value_states, seq_len=kv_seq_len)                       \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m202 \u001b[2m│   │   \u001b[0mquery_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, s   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m203 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[2m# [bsz, nh, t, hd]\u001b[0m                                                                 \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m204 \u001b[0m\u001b[2m│   │   \u001b[0m                                                                                   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m205 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mif\u001b[0m past_key_value \u001b[95mis\u001b[0m \u001b[95mnot\u001b[0m \u001b[94mNone\u001b[0m:                                                     \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[2;33m/opt/homebrew/Caskroom/miniconda/base/lib/python3.10/site-packages/transformers/models/llama/\u001b[0m\u001b[1;33mmod\u001b[0m \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[1;33meling_llama.py\u001b[0m:\u001b[94m136\u001b[0m in \u001b[92mapply_rotary_pos_emb\u001b[0m                                                       \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m133 \u001b[0m\u001b[2m│   \u001b[0msin = sin.squeeze(\u001b[94m1\u001b[0m).squeeze(\u001b[94m0\u001b[0m)  \u001b[2m# [seq_len, dim]\u001b[0m                                      \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m134 \u001b[0m\u001b[2m│   \u001b[0mcos = cos[position_ids].unsqueeze(\u001b[94m1\u001b[0m)  \u001b[2m# [bs, 1, seq_len, dim]\u001b[0m                          \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m135 \u001b[0m\u001b[2m│   \u001b[0msin = sin[position_ids].unsqueeze(\u001b[94m1\u001b[0m)  \u001b[2m# [bs, 1, seq_len, dim]\u001b[0m                          \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m136 \u001b[2m│   \u001b[0mq_embed = (q * cos) + (rotate_half(q) * sin)                                           \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m137 \u001b[0m\u001b[2m│   \u001b[0mk_embed = (k * cos) + (rotate_half(k) * sin)                                           \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m138 \u001b[0m\u001b[2m│   \u001b[0m\u001b[94mreturn\u001b[0m q_embed, k_embed                                                                \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m139 \u001b[0m                                                                                           \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[2;33m/opt/homebrew/Caskroom/miniconda/base/lib/python3.10/site-packages/torch/utils/\u001b[0m\u001b[1;33m_device.py\u001b[0m:\u001b[94m76\u001b[0m in  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[92m__torch_function__\u001b[0m                                                                               \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m73 \u001b[0m\u001b[2m│   │   \u001b[0mkwargs = kwargs \u001b[95mor\u001b[0m {}                                                               \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m74 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mif\u001b[0m func \u001b[95min\u001b[0m _device_constructors() \u001b[95mand\u001b[0m kwargs.get(\u001b[33m'\u001b[0m\u001b[33mdevice\u001b[0m\u001b[33m'\u001b[0m) \u001b[95mis\u001b[0m \u001b[94mNone\u001b[0m:                 \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m75 \u001b[0m\u001b[2m│   │   │   \u001b[0mkwargs[\u001b[33m'\u001b[0m\u001b[33mdevice\u001b[0m\u001b[33m'\u001b[0m] = \u001b[96mself\u001b[0m.device                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m76 \u001b[2m│   │   \u001b[0m\u001b[94mreturn\u001b[0m func(*args, **kwargs)                                                        \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m77 \u001b[0m                                                                                            \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m78 \u001b[0m\u001b[2m# NB: This is directly called from C++ in torch/csrc/Device.cpp\u001b[0m                             \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m79 \u001b[0m\u001b[94mdef\u001b[0m \u001b[92mdevice_decorator\u001b[0m(device, func):                                                         \u001b[31m│\u001b[0m\n",
       "\u001b[31m╰──────────────────────────────────────────────────────────────────────────────────────────────────╯\u001b[0m\n",
       "\u001b[1;91mRuntimeError: \u001b[0mMPS backend out of memory \u001b[1m(\u001b[0mMPS allocated: \u001b[1;36m114.43\u001b[0m GB, other allocations: \u001b[1;36m7.91\u001b[0m GB, max allowed: \u001b[1;36m122.40\u001b[0m \n",
       "GB\u001b[1m)\u001b[0m. Tried to allocate \u001b[1;36m64.00\u001b[0m MB on private pool. Use \u001b[33mPYTORCH_MPS_HIGH_WATERMARK_RATIO\u001b[0m=\u001b[1;36m0\u001b[0m\u001b[1;36m.0\u001b[0m to disable upper limit \n",
       "for memory allocations \u001b[1m(\u001b[0mmay cause system failure\u001b[1m)\u001b[0m.\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "class LlamaRetrofit(torch.nn.Module):\n",
    "    def __init__(self, llama, rtr_num=7):\n",
    "        super().__init__()\n",
    "\n",
    "        self.rtr_num = rtr_num\n",
    "\n",
    "        self.emb = llama.model.embed_tokens\n",
    "        self.blocks = llama.model.layers\n",
    "        self.norm = llama.model.norm\n",
    "        self.head = llama.lm_head\n",
    "\n",
    "        self.cross_attn = torch.nn.ModuleList([torch.nn.MultiheadAttention(4096, 32, batch_first=True) for _ in range(len(self.blocks)//4)])\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.emb(x)\n",
    "        for i, block in enumerate(self.blocks):\n",
    "            x = block(x)[0]\n",
    "            if i % 4 == 0 and i != 0:\n",
    "                # do self attention as normal, but then reshape the batch into [-1, rtr_num+1, seq_len, hidden_size], and use the q, k, and v from the current layer to perform cross attention from item 0 to 1-rtr_num, and vis versa, to move information between the retrieved items and the query.\n",
    "                tiled = x[0].unsqueeze(0).repeat(self.rtr_num, 1, 1)\n",
    "                print(tiled.shape)\n",
    "                new_query = self.cross_attn[i//4](tiled, x[1:], x[1:])[0]\n",
    "                # average across batch\n",
    "                new_query = new_query.mean(dim=0)\n",
    "                x = torch.cat([new_query.unsqueeze(0), x[1:]], dim=0)\n",
    "        x = self.norm(x)\n",
    "        x = self.head(x)\n",
    "        return x[0]\n",
    "\n",
    "# create test model and pass a sample batch of 16\n",
    "model = LlamaRetrofit(llama_mod)\n",
    "# move to cuda\n",
    "model = model.cuda()\n",
    "\n",
    "optimizer = torch.optim.Adam(model.cross_attn.parameters(), lr=1e-4)\n",
    "optimizer.zero_grad()\n",
    "\n",
    "loss_fn = torch.nn.CrossEntropyLoss()\n",
    "for i in range(len(tokenized_train)):\n",
    "    label = tokenized_train[i][0][0][1:]\n",
    "    # add a padding token to the beginning of the label\n",
    "    label = torch.cat([torch.tensor([0]), label])\n",
    "    inp = torch.cat([tokenized_train[i][0], tokenized_train[i][1]])\n",
    "    output = model(inp)\n",
    "    # argmax and decode\n",
    "    stringy = torch.argmax(output, dim=-1)\n",
    "    stringy = llama_tok.decode(stringy)\n",
    "    print(stringy)\n",
    "    loss = loss_fn(output, label)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    optimizer.zero_grad()\n",
    "    print(loss.item())\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([ 3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130, 16020, 16020, 16020, 16020, 16020, 16020, 16020, 16020,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130, 16020, 16020, 16020, 16020,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130, 27823,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130, 16020,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130, 16020,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130, 17856, 17856, 17856,  3130,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130, 16020,\n",
      "         3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,  3130,\n",
      "         3130,  3130], device='mps:0')\n",
      "ilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilon Aqu Aqu Aqu Aqu Aqu Aqu Aqu Aquilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilon Aqu Aqu Aqu AquilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonNililonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilon Aquilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilonilon Aquiloniloniloniloniloniloniloniloniloniloniloniloniloniloniloniloniloniloninéinéinéilonilonilonilonilonilonilonilonilonilon Aquilonilonilonilonilonilonilonilonilonilonilonilon\n"
     ]
    }
   ],
   "source": [
    "print(stringy)\n",
    "print(llama_tok.decode(stringy))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Token indices sequence length is longer than the specified maximum sequence length for this model (28985 > 512). Running this sequence through the model will result in indexing errors\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 28985])\n",
      "torch.Size([1, 30485])\n"
     ]
    }
   ],
   "source": [
    "# test_adapter = LlamaT5Adapter(t5_mod, llama_mod, ctx_num=16)\n",
    "\n",
    "# optimizer = torch.optim.Adam(test_adapter.parameters(), lr=1e-4)\n",
    "# sample_out = test_adapter(torch.tensor([[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]]), torch.tensor([[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]]))\n",
    "# print(sample_out.shape)\n",
    "\n",
    "\n",
    "# sample_out.sum().backward()\n",
    "# optimizer.step()\n",
    "\n",
    "import sqlite3\n",
    "\n",
    "conn = sqlite3.connect('data.db')\n",
    "\n",
    "# create a cursor object\n",
    "cur = conn.cursor()\n",
    "\n",
    "cur.execute('SELECT content, author FROM discord ORDER BY timestamp')\n",
    "data = cur.fetchall()\n",
    "\n",
    "# concat author name to all messages\n",
    "data = [f'{author}: {content}' for content, author in data]\n",
    "\n",
    "# combine all messages into one string\n",
    "data = '\\n\\n'.join(data)\n",
    "\n",
    "data = data[:100000]\n",
    "\n",
    "t5_tokenized_data = t5_tok.encode(data, return_tensors='pt')\n",
    "llama_tokenized_data = llama_tok.encode(data, return_tensors='pt')\n",
    "print(t5_tokenized_data.shape)\n",
    "print(llama_tokenized_data.shape)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0, batch 10240\n",
      "Loss: 11.575541496276855\n",
      "Epoch 0, batch 11264\n",
      "Loss: 9.961394309997559\n",
      "Epoch 0, batch 12288\n"
     ]
    }
   ],
   "source": [
    "\n",
    "adapter = LlamaT5Adapter(t5_mod, llama_mod, ctx_num=8)\n",
    "\n",
    "\n",
    "params_to_optimize = list(adapter.llama_queries.parameters()) + list(adapter.t5_keys.parameters()) + list(adapter.t5_values.parameters())\n",
    "optimizer = torch.optim.Adam(params_to_optimize, lr=1e-4)\n",
    "# optimizer = torch.optim.Adam(adapter.parameters(), lr=1e-4)\n",
    "\n",
    "total_llama_tokens = llama_tokenized_data.shape[1]\n",
    "\n",
    "# for the training loop, for each chunk passed to the llama decoder, the previous 16 chunks of 512 t5-tokens are passed to the t5 encoder\n",
    "epochs = 4\n",
    "llama_batch_size = 1024\n",
    "t5_batch_size = 512\n",
    "for epoch in range(epochs):\n",
    "    for i in range(t5_batch_size*20, total_llama_tokens, llama_batch_size):\n",
    "        print(f'Epoch {epoch}, batch {i}')\n",
    "        # get the next batch of data\n",
    "        llama_in = llama_tokenized_data[:, i:i+llama_batch_size]\n",
    "        # to get approximately the same spot in the chat for the t5 tokenization, use the current index/total_llama_tokens, and multiply by the total t5 tokens\n",
    "        t5_index = int(i/total_llama_tokens*t5_tokenized_data.shape[1])\n",
    "        t5_in = t5_tokenized_data[:, t5_index-t5_batch_size*8:t5_index]\n",
    "        # reshape the t5 input to be 16 chunks of 512 tokens\n",
    "        t5_in = t5_in.reshape(-1, t5_batch_size)\n",
    "        # pass the data through the model\n",
    "        out = adapter(llama_in, t5_in)\n",
    "        # calculate loss by shifting the target data by 1\n",
    "        loss = torch.nn.functional.cross_entropy(out[:, :-1].reshape(-1, out.shape[-1]), llama_in[:, 1:].reshape(-1))\n",
    "        # backpropagate\n",
    "        loss.backward()\n",
    "        # clip gradients\n",
    "        torch.nn.utils.clip_grad_norm_(params_to_optimize, 1.0)\n",
    "        # update parameters\n",
    "        optimizer.step()\n",
    "        # zero gradients\n",
    "        optimizer.zero_grad()\n",
    "        print(f'Loss: {loss.item()}')\n",
    "        del llama_in, t5_in, out, loss\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.9"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}