test_babi.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '1'\n",
    "\n",
    "import json\n",
    "import logging\n",
    "import os\n",
    "import math\n",
    "import shutil\n",
    "from pathlib import Path\n",
    "from itertools import chain\n",
    "\n",
    "# from dotenv import load_dotenv\n",
    "import torch\n",
    "import numpy as np\n",
    "import datasets\n",
    "import transformers\n",
    "from torch.utils.data import DataLoader\n",
    "from huggingface_hub import hf_hub_download\n",
    "\n",
    "\n",
    "from transformers import AutoConfig, AutoTokenizer, HfArgumentParser, AutoModelForCausalLM\n",
    "\n",
    "from modeling_rmt.language_modeling import MemoryCell, RecurrentWrapper\n",
    "from modeling_amt.language_modeling import AssociativeMemoryCell, AssociativeRecurrentWrapper\n",
    "from modeling_rmt.lm_parallel_mem import MemoryCell as PMemoryCell, RecurrentWrapper as PRecurrentWrapper\n",
    "\n",
    "from torch.nn.utils.rnn import pad_sequence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "nums_segments = [1024]\n",
    "\n",
    "segment_size = 512\n",
    "block_size = segment_size\n",
    "\n",
    "test_size = 1000\n",
    "batch_size = 1\n",
    "\n",
    "model_name = 'armt'\n",
    "num_mem_tokens = 10\n",
    "model_path = './models'\n",
    "d_mem = 64\n",
    "\n",
    "task_dataset = \"qa1_single-supporting-fact\"\n",
    "max_n_facts = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "if model_name == 'armt':\n",
    "    mem_cell_cls = AssociativeMemoryCell\n",
    "    rec_wrap_cls = AssociativeRecurrentWrapper\n",
    "elif model_name == 'rmt':\n",
    "    mem_cell_cls = MemoryCell\n",
    "    rec_wrap_cls = RecurrentWrapper\n",
    "elif model_name == 'prmt':\n",
    "    mem_cell_cls = PMemoryCell\n",
    "    rec_wrap_cls = PRecurrentWrapper\n",
    "if model_name != 'armt':\n",
    "    d_mem = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AssociativeRecurrentWrapper(\n",
       "  (memory_cell): AssociativeMemoryCell(\n",
       "    (model): GPT2LMHeadModel(\n",
       "      (transformer): GPT2Model(\n",
       "        (wte): Embedding(50257, 768)\n",
       "        (wpe): Embedding(1024, 768)\n",
       "        (drop): Dropout(p=0.1, inplace=False)\n",
       "        (h): ModuleList(\n",
       "          (0-11): 12 x AssociativeLayerWrapper(\n",
       "            (W_mq): Linear(in_features=768, out_features=64, bias=False)\n",
       "            (W_mk): Linear(in_features=768, out_features=64, bias=False)\n",
       "            (W_mv): Linear(in_features=768, out_features=768, bias=False)\n",
       "            (W_mb): Linear(in_features=768, out_features=1, bias=True)\n",
       "            (layer): GPT2Block(\n",
       "              (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "              (attn): GPT2Attention(\n",
       "                (c_attn): Conv1D()\n",
       "                (c_proj): Conv1D()\n",
       "                (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "                (resid_dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "              (mlp): GPT2MLP(\n",
       "                (c_fc): Conv1D()\n",
       "                (c_proj): Conv1D()\n",
       "                (act): NewGELUActivation()\n",
       "                (dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "        (ln_f): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      )\n",
       "      (lm_head): Linear(in_features=768, out_features=50257, bias=False)\n",
       "    )\n",
       "    (W_mq): ModuleList()\n",
       "    (layers): ModuleList(\n",
       "      (0-11): 12 x AssociativeLayerWrapper(\n",
       "        (W_mq): Linear(in_features=768, out_features=64, bias=False)\n",
       "        (W_mk): Linear(in_features=768, out_features=64, bias=False)\n",
       "        (W_mv): Linear(in_features=768, out_features=768, bias=False)\n",
       "        (W_mb): Linear(in_features=768, out_features=1, bias=True)\n",
       "        (layer): GPT2Block(\n",
       "          (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "          (attn): GPT2Attention(\n",
       "            (c_attn): Conv1D()\n",
       "            (c_proj): Conv1D()\n",
       "            (attn_dropout): Dropout(p=0.1, inplace=False)\n",
       "            (resid_dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "          (mlp): GPT2MLP(\n",
       "            (c_fc): Conv1D()\n",
       "            (c_proj): Conv1D()\n",
       "            (act): NewGELUActivation()\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(\"gpt2\")\n",
    "mem_cell_args = dict(\n",
    "        base_model=model,\n",
    "        num_mem_tokens=num_mem_tokens,\n",
    ")\n",
    "if d_mem is not None:\n",
    "    mem_cell_args['d_mem'] = d_mem\n",
    "\n",
    "\n",
    "cell = mem_cell_cls(**mem_cell_args, wrap_pos=False)\n",
    "model = rec_wrap_cls(cell, segment_size=block_size, k2=-1)\n",
    "\n",
    "model_cpt = os.path.join(model_path, \"model_best/pytorch_model.bin\")\n",
    "cpt = torch.load(model_cpt, map_location='cuda')\n",
    "model.load_state_dict(cpt, strict=False)\n",
    "model.eval().to('cuda')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from babilong_utils import TaskDataset, NoiseInjectionDataset, SentenceSampler\n",
    "\n",
    "babi_path = \"/home/rodkin/lab/t5-experiments/data/tasks_1-20_v1-2/en-10k\"\n",
    "\n",
    "noise_dataset = datasets.load_dataset(\"pg19\")\n",
    "noise_dataset_test = noise_dataset['test']\n",
    "test_path = os.path.join(babi_path, f\"{task_dataset}_test.txt\")\n",
    "task_dataset_test = TaskDataset(test_path, max_n_facts=max_n_facts)\n",
    "qa_margin = 20\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"gpt2\")\n",
    "for i, num_segments in enumerate(nums_segments):\n",
    "    sample_size_ = segment_size * num_segments\n",
    "    sample_size = sample_size_ - qa_margin\n",
    "    noise_sampler_test = SentenceSampler(noise_dataset_test, tokenizer=tokenizer, max_sentence_len=None, shuffle=True, random_seed=42)\n",
    "test_dataset = NoiseInjectionDataset(task_dataset=task_dataset_test,\n",
    "                                     noise_sampler=noise_sampler_test,\n",
    "                                     tokenizer=tokenizer,\n",
    "                                     sample_size=sample_size,\n",
    "                                     mixed_length_ratio=0,\n",
    "                                     task_start_pct=None,\n",
    "                                     task_end_pct=None\n",
    "                                    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data.distributed import DistributedSampler\n",
    "\n",
    "id_pad_value = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id\n",
    "gen_token = tokenizer.encode('GEN')[0]\n",
    "eos_token = tokenizer.eos_token_id\n",
    "\n",
    "def collate_fn(batch):\n",
    "        targets = [torch.tensor(b['target_tokens']) for b in batch]\n",
    "        input_ids = [torch.tensor(b['input_tokens'] + b['question_tokens'] + [gen_token] + b['target_tokens'] + [eos_token]) for b in batch]\n",
    "        gen_inputs = [torch.tensor(b['input_tokens'] + b['question_tokens'] + [gen_token]) for b in batch]\n",
    "\n",
    "        attention_mask = [torch.ones_like(b, dtype=int) for b in input_ids]\n",
    "        labels_mask = [torch.zeros_like(b, dtype=bool) for b in input_ids]\n",
    "        for m, t in zip(labels_mask, targets):\n",
    "            m[-len(t) - 2:] = True\n",
    "\n",
    "        input_ids = pad_sequence(input_ids, padding_value=id_pad_value, batch_first=True)\n",
    "        gen_inputs = pad_sequence(gen_inputs, padding_value=id_pad_value, batch_first=True)\n",
    "        attention_mask = pad_sequence(attention_mask, padding_value=0, batch_first=True)\n",
    "        labels_mask = pad_sequence(labels_mask, padding_value=0, batch_first=True)\n",
    "\n",
    "        collated = {}\n",
    "        collated['input_ids'] = collated['labels'] = input_ids\n",
    "        collated['input_ids_generate'] = gen_inputs\n",
    "        collated['labels_mask'] = labels_mask\n",
    "        collated['attention_mask'] = attention_mask.bool()\n",
    "        collated['attention_mask_generate'] = (gen_inputs != id_pad_value).bool()\n",
    "        collated['target_text'] = [b['answer'] for b in batch]\n",
    "        return collated\n",
    "\n",
    "test_dataloader = DataLoader(batch_size=batch_size, dataset=test_dataset, collate_fn=collate_fn, pin_memory=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0dbe2d36a2b0474992e647ecc4906552",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/999 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 524274])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 524274])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 524274])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n",
      "Token indices sequence length is longer than the specified maximum sequence length for this model (2108 > 1024). Running this sequence through the model will result in indexing errors\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 524274])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 524274])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 524274])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 524274])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 524274])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 524274])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 524274])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 524274])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 524274])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 524274])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 524274])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 524274])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 524274])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 524274])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 524274])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 524274])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 524274])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 524274])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 524274])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 524274])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 524274])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 524274])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 524274])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 524274])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 524274])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 524274])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 524274])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 524274])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 524274])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 524274])\n"
     ]
    }
   ],
   "source": [
    "from tqdm.auto import tqdm\n",
    "device = 'cuda'\n",
    "metrics = []\n",
    "bar = tqdm(enumerate(test_dataloader), total=len(test_dataloader))\n",
    "with torch.no_grad():\n",
    "  for j, batch in bar:\n",
    "    inp_gen = batch['input_ids_generate'].to(device)\n",
    "    print(inp_gen.shape)\n",
    "    gen_mask = batch['attention_mask_generate'].to(device)\n",
    "    p = model.generate(input_ids=inp_gen, attention_mask=gen_mask, max_new_tokens=10)\n",
    "    generation_outputs = tokenizer.batch_decode([d for d in p], add_special_tokens=False)\n",
    "    for i, o in enumerate(generation_outputs):\n",
    "        if '<|endoftext|>' in o:\n",
    "            generation_outputs[i] = o.split('<|endoftext|>')[1].strip()\n",
    "    y = batch['target_text']\n",
    "    exact_match = np.mean([text == pred for text, pred in zip (y, generation_outputs)])\n",
    "    metrics.append(exact_match)\n",
    "    bar.set_description(f\"current exact match: {np.mean(metrics)}\")\n",
    "\n",
    "metric = np.mean(metrics)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def last_log_forward(model, input_ids, attention_mask):\n",
    "    model.memory_cell.zero_mem()\n",
    "    segmented = model.segment(input_ids=input_ids, attention_mask=attention_mask)\n",
    "\n",
    "    for seg_num, segment in enumerate(segmented[:-1]):\n",
    "        _ = model.memory_cell(**segment, output_hidden_states=True, zero_mem=False)\n",
    "\n",
    "    final_segment = segmented[-1]\n",
    "    out = model.memory_cell(**final_segment, zero_mem=False)\n",
    "    model.memory_cell.zero_mem()\n",
    "    return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "981ff1432bcc4a6f85b564a6f2bfd2e9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/63 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "ename": "OutOfMemoryError",
     "evalue": "CUDA out of memory. Tried to allocate 1.56 GiB (GPU 0; 10.91 GiB total capacity; 8.58 GiB already allocated; 1016.06 MiB free; 9.76 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mOutOfMemoryError\u001b[0m                          Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[9], line 10\u001b[0m\n\u001b[1;32m      8\u001b[0m inp \u001b[38;5;241m=\u001b[39m batch[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124minput_ids\u001b[39m\u001b[38;5;124m'\u001b[39m]\u001b[38;5;241m.\u001b[39mto(device)\n\u001b[1;32m      9\u001b[0m mask \u001b[38;5;241m=\u001b[39m batch[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mattention_mask\u001b[39m\u001b[38;5;124m'\u001b[39m]\u001b[38;5;241m.\u001b[39mto(device)\n\u001b[0;32m---> 10\u001b[0m output \u001b[38;5;241m=\u001b[39m \u001b[43mlast_log_forward\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minp\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmask\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     11\u001b[0m predictions \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39margmax(output[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mlogits\u001b[39m\u001b[38;5;124m'\u001b[39m], dim\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m)\n\u001b[1;32m     12\u001b[0m labels_mask \u001b[38;5;241m=\u001b[39m model\u001b[38;5;241m.\u001b[39msegment(labels_mask\u001b[38;5;241m=\u001b[39mbatch[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mlabels_mask\u001b[39m\u001b[38;5;124m'\u001b[39m])[\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m][\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mlabels_mask\u001b[39m\u001b[38;5;124m'\u001b[39m]\n",
      "Cell \u001b[0;32mIn[8], line 6\u001b[0m, in \u001b[0;36mlast_log_forward\u001b[0;34m(model, input_ids, attention_mask)\u001b[0m\n\u001b[1;32m      3\u001b[0m segmented \u001b[38;5;241m=\u001b[39m model\u001b[38;5;241m.\u001b[39msegment(input_ids\u001b[38;5;241m=\u001b[39minput_ids, attention_mask\u001b[38;5;241m=\u001b[39mattention_mask)\n\u001b[1;32m      5\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m seg_num, segment \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28menumerate\u001b[39m(segmented[:\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m]):\n\u001b[0;32m----> 6\u001b[0m     _ \u001b[38;5;241m=\u001b[39m \u001b[43mmodel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmemory_cell\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43msegment\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43moutput_hidden_states\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mzero_mem\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mFalse\u001b[39;49;00m\u001b[43m)\u001b[49m\n\u001b[1;32m      8\u001b[0m final_segment \u001b[38;5;241m=\u001b[39m segmented[\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m]\n\u001b[1;32m      9\u001b[0m out \u001b[38;5;241m=\u001b[39m model\u001b[38;5;241m.\u001b[39mmemory_cell(\u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mfinal_segment, zero_mem\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mFalse\u001b[39;00m)\n",
      "File \u001b[0;32m~/lab/env/lib/python3.8/site-packages/torch/nn/modules/module.py:1501\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1496\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1497\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1498\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1499\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1500\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1501\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1502\u001b[0m \u001b[38;5;66;03m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1503\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[38;5;241m=\u001b[39m [], []\n",
      "File \u001b[0;32m/cephfs/home/rodkin/lab/t5-experiments/modeling_amt/language_modeling.py:189\u001b[0m, in \u001b[0;36mAssociativeMemoryCell.forward\u001b[0;34m(self, input_ids, labels, labels_mask, zero_mem, **kwargs)\u001b[0m\n\u001b[1;32m    184\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mzero_mem()\n\u001b[1;32m    187\u001b[0m seg_kwargs \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mprocess_input(input_ids, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[0;32m--> 189\u001b[0m out \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmodel\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mseg_kwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    191\u001b[0m out \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mprocess_output(out, labels, labels_mask, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[1;32m    193\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m out\n",
      "File \u001b[0;32m~/lab/env/lib/python3.8/site-packages/torch/nn/modules/module.py:1501\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1496\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1497\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1498\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1499\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1500\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1501\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1502\u001b[0m \u001b[38;5;66;03m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1503\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[38;5;241m=\u001b[39m [], []\n",
      "File \u001b[0;32m~/lab/env/lib/python3.8/site-packages/transformers/models/gpt2/modeling_gpt2.py:1098\u001b[0m, in \u001b[0;36mGPT2LMHeadModel.forward\u001b[0;34m(self, input_ids, past_key_values, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, encoder_hidden_states, encoder_attention_mask, labels, use_cache, output_attentions, output_hidden_states, return_dict)\u001b[0m\n\u001b[1;32m   1095\u001b[0m     torch\u001b[38;5;241m.\u001b[39mcuda\u001b[38;5;241m.\u001b[39mset_device(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtransformer\u001b[38;5;241m.\u001b[39mfirst_device)\n\u001b[1;32m   1096\u001b[0m     hidden_states \u001b[38;5;241m=\u001b[39m hidden_states\u001b[38;5;241m.\u001b[39mto(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlm_head\u001b[38;5;241m.\u001b[39mweight\u001b[38;5;241m.\u001b[39mdevice)\n\u001b[0;32m-> 1098\u001b[0m lm_logits \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlm_head\u001b[49m\u001b[43m(\u001b[49m\u001b[43mhidden_states\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1100\u001b[0m loss \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m   1101\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m labels \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m   1102\u001b[0m     \u001b[38;5;66;03m# move labels to correct device to enable model parallelism\u001b[39;00m\n",
      "File \u001b[0;32m~/lab/env/lib/python3.8/site-packages/torch/nn/modules/module.py:1501\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1496\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1497\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1498\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1499\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1500\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1501\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1502\u001b[0m \u001b[38;5;66;03m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1503\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[38;5;241m=\u001b[39m [], []\n",
      "File \u001b[0;32m~/lab/env/lib/python3.8/site-packages/torch/nn/modules/linear.py:114\u001b[0m, in \u001b[0;36mLinear.forward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    113\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;28minput\u001b[39m: Tensor) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Tensor:\n\u001b[0;32m--> 114\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mF\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlinear\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mweight\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbias\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[0;31mOutOfMemoryError\u001b[0m: CUDA out of memory. Tried to allocate 1.56 GiB (GPU 0; 10.91 GiB total capacity; 8.58 GiB already allocated; 1016.06 MiB free; 9.76 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF"
     ]
    }
   ],
   "source": [
    "from tqdm.auto import tqdm\n",
    "\n",
    "metrics = []\n",
    "bar = tqdm(enumerate(test_dataloader), total=len(test_dataloader))\n",
    "device = 'cuda'\n",
    "with torch.no_grad():\n",
    "  for j, batch in bar:\n",
    "    inp = batch['input_ids'].to(device)\n",
    "    mask = batch['attention_mask'].to(device)\n",
    "    output = last_log_forward(model, inp, mask)\n",
    "    predictions = torch.argmax(output['logits'], dim=-1)\n",
    "    labels_mask = model.segment(labels_mask=batch['labels_mask'])[-1]['labels_mask']\n",
    "\n",
    "    predicted_labels = [p[m] for p, m in zip(predictions, labels_mask)]\n",
    "    predicted_labels = tokenizer.batch_decode(predicted_labels, add_special_tokens=False)\n",
    "    for i, l in enumerate(predicted_labels):\n",
    "        if '<|endoftext|>' in l:\n",
    "            eos_ind = predicted_labels[i].index('<|endoftext|>')\n",
    "            predicted_labels[i] = predicted_labels[i][:eos_ind]\n",
    "    y = batch['target_text']\n",
    "    exact_match = np.mean([text == pred for text, pred in zip (y, predicted_labels)])\n",
    "    metrics.append(exact_match)\n",
    "    bar.set_description(f\"current exact match: {np.mean(metrics)}\")\n",
    "\n",
    "metric = np.mean(metrics)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Failed to detect the name of this notebook, you can set it manually with the WANDB_NOTEBOOK_NAME environment variable to enable code saving.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mirodkin\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "wandb version 0.16.2 is available!  To upgrade, please run:\n",
       " $ pip install wandb --upgrade"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.15.10"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>/cephfs/home/rodkin/lab/wandb/run-20240204_110834-frfyc1qx</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href='https://wandb.ai/irodkin/associative_retrieval/runs/frfyc1qx' target=\"_blank\">prmt</a></strong> to <a href='https://wandb.ai/irodkin/associative_retrieval' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View project at <a href='https://wandb.ai/irodkin/associative_retrieval' target=\"_blank\">https://wandb.ai/irodkin/associative_retrieval</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run at <a href='https://wandb.ai/irodkin/associative_retrieval/runs/frfyc1qx' target=\"_blank\">https://wandb.ai/irodkin/associative_retrieval/runs/frfyc1qx</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import numpy as np\n",
    "import wandb\n",
    "run = wandb.init(project='associative_retrieval', entity='irodkin', name=f'{model_name}')\n",
    "table = wandb.Table(data=np.vstack([nums_pairs, metrics]).T, columns=['evaluated_on', 'test/exact_match'])\n",
    "mem_cap = wandb.Table(data=np.vstack([nums_pairs, pairs_in_mem]).T, columns=['evaluated_on', 'test/pairs_in_mem'])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "Waiting for W&B process to finish... <strong style=\"color:green\">(success).</strong>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run <strong style=\"color:#cdcd00\">prmt</strong> at: <a href='https://wandb.ai/irodkin/associative_retrieval/runs/frfyc1qx' target=\"_blank\">https://wandb.ai/irodkin/associative_retrieval/runs/frfyc1qx</a><br/>Synced 4 W&B file(s), 2 media file(s), 2 artifact file(s) and 0 other file(s)"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Find logs at: <code>./wandb/run-20240204_110834-frfyc1qx/logs</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "line = run.plot_table(\"wandb/line/v0\", table, {\"x\":'evaluated_on', \"y\":'test/exact_match'})\n",
    "line_cap = run.plot_table(\"wandb/line/v0\", mem_cap, {\"x\":'evaluated_on', \"y\":'test/pairs_in_mem'})\n",
    "wandb.log({f'exact_match_on_pairs': line})\n",
    "wandb.log({f'capacity_on_pairs': line_cap})\n",
    "wandb.finish()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}