embed.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModelForSeq2SeqLM\n",
    "import torch\n",
    "torch.set_default_device('mps')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from FlagEmbedding import FlagModel\n",
    "model = FlagModel('BAAI/bge-large-en', query_instruction_for_retrieval=\"Represent this sentence for searching relevant passages:\")\n",
    "\n",
    "# q_embeddings = model.encode_queries(queries)\n",
    "# p_embeddings = model.encode(passages)\n",
    "# scores = q_embeddings @ p_embeddings.T\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# load link_pairs.json\n",
    "import json\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('link_pairs.json', 'r') as f:\n",
    "    link_pairs = json.load(f)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "links = [pair['link'] for pair in link_pairs]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Batches: 100%|██████████| 13/13 [26:36<00:00, 122.79s/it]\n"
     ]
    }
   ],
   "source": [
    "context_emb = model.encode(links)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "summaries = []\n",
    "for pair in link_pairs:\n",
    "    for summary in pair['summaries']:\n",
    "        summaries.append(summary)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Batches: 100%|██████████| 50/50 [04:17<00:00,  5.16s/it]\n"
     ]
    }
   ],
   "source": [
    "summary_emb = model.encode_queries(summaries)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4\n"
     ]
    }
   ],
   "source": [
    "print(len(link_pairs[0]['summaries']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "# for each pair, add a third item, the embedding vector\n",
    "for i, pair in enumerate(link_pairs):\n",
    "    pair['lembed'] = context_emb[i].tolist()\n",
    "    summary_embeds = []\n",
    "    for j, summary in enumerate(pair['summaries']):\n",
    "        summary_embeds.append(summary_emb[i*4+j].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": null,
   "metadata": {},
   "outputs": [],
   "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
}