PlayingWithMemo-pytorch-esr.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "bd95042e-0d56-401d-ae0c-0fbe1eee44e5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GPU: NVIDIA RTX A6000 is available.\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import math\n",
    "import os\n",
    "\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    print(f\"GPU: {torch.cuda.get_device_name(0)} is available.\")\n",
    "else:\n",
    "    print(\"No GPU available. Training will run on CPU.\")\n",
    "\n",
    "verbose = False"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ac14a52-6180-4c0d-997e-b24745c624f9",
   "metadata": {},
   "source": [
    "## MeMo Tokenizer and input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "7956914c-ec7b-4fa7-a6a2-dbff28f8cec2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from MeMoPyTorch.modelling_memo_tokenizer import MeMoTokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8673f471-46ce-4875-b74d-e8a0378a01ad",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The tokenizer class you load from this checkpoint is not the same type as the class this function is called from. It may result in unexpected tokenization. \n",
      "The tokenizer class you load from this checkpoint is 'GPTNeoXTokenizer'. \n",
      "The class this function is called from is 'MeMoTokenizer'.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Setting pad token and pad token id = <|endoftext|>, 0\n"
     ]
    }
   ],
   "source": [
    "max_length = 12 \n",
    "tokenizer = MeMoTokenizer.from_pretrained(\"EleutherAI/gpt-neox-20b\", \n",
    "                                          truncation_side = 'left',\n",
    "                                          padding_side='left', max_length=max_length, head_number=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c159d5f5-2bf5-4d04-bde2-c897be8f494f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(tensor([[18886,   256, 36144,  4164,  1809,    80,  1448,   295,   532,  1584,\n",
      "            13, 50190]]), tensor([[  256, 36144,  4164,  1809,    80,  1448,   295,   532,  1584,    13,\n",
      "         50190,    15]]))\n"
     ]
    }
   ],
   "source": [
    "with open(\"testo_di_prova.txt\") as my_first_text_f:\n",
    "    my_first_text = my_first_text_f.read()\n",
    "\n",
    "token_ids = tokenizer.encode(my_first_text)#, return_tensors='pt')\n",
    "print(token_ids) # return max len + 1 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "32c458fb-4eeb-4ffd-91d6-0f22d4d6566f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(dict_keys(['input_ids', 'labels']), torch.Size([52, 12]))"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "memo_input = tokenizer.get_text_batch_encoding([my_first_text, my_first_text[0:10]])\n",
    "memo_input.keys(), memo_input['input_ids'].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "427781b3-ae1c-430a-b961-6500d44c4100",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|>Cosimo di\n",
      "<|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|>Cosimo di Giovanni\n",
      "\n",
      " de' Medici detto il Vecchio o Pater\n",
      "' Medici detto il Vecchio o Pater patri\n",
      "\n",
      "æ (Firenze, 27 settembre 1389\n",
      " (Firenze, 27 settembre 1389 –\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for i in range(3):\n",
    "    print(tokenizer.decode(memo_input['input_ids'][i]))\n",
    "    print(tokenizer.decode(memo_input['labels'][i]))\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57bb7a04-5b90-48f5-be13-96ca53ce73a7",
   "metadata": {},
   "source": [
    "## MeMo Embedding layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e3e6eeb9-4f10-4fb6-b6c4-53dcfedca561",
   "metadata": {},
   "outputs": [],
   "source": [
    "from MeMoPyTorch.modelling_memo_embedding import MeMoEmbedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "94661fdd-cef9-4a41-b197-9d5f40645e25",
   "metadata": {},
   "outputs": [],
   "source": [
    "d,h,l = 1024, 4, 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "2612d880-6777-4052-9b05-f32fb1f6dec0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MeMo embedding initilialization\n"
     ]
    }
   ],
   "source": [
    "embedding = MeMoEmbedding(\n",
    "    num_embeddings=tokenizer.vocab_size,\n",
    "    embedding_dim=d,\n",
    "    padding_idx=tokenizer.pad_token_id, #0\n",
    "    _freeze=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "3835b9f2-c1a7-4673-b9a4-ba5710fc0df6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[   0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,\n",
      "         5089],\n",
      "        [   0,    0,    0,    0,    0,    0,    0,    0,    0, 2447, 6945,  287,\n",
      "         6004]])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[[ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],\n",
       "         [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],\n",
       "         [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],\n",
       "         ...,\n",
       "         [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],\n",
       "         [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],\n",
       "         [-0.0023,  0.0027, -0.0391,  ..., -0.0083, -0.0081, -0.0272]],\n",
       "\n",
       "        [[ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],\n",
       "         [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],\n",
       "         [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],\n",
       "         ...,\n",
       "         [-0.0065, -0.0010,  0.0132,  ...,  0.0424, -0.0111,  0.0262],\n",
       "         [ 0.0285, -0.0263,  0.0528,  ..., -0.0215,  0.0131,  0.0006],\n",
       "         [ 0.0702,  0.0040,  0.0162,  ...,  0.0326, -0.0240, -0.0420]]])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "input_tokens_ids = tokenizer(['Test', 'Un altro Test'])['input_ids']\n",
    "print(input_tokens_ids)\n",
    "\n",
    "input_embeddings = embedding.forward(input_tokens_ids)\n",
    "input_embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "cf18be91-e462-43c4-adb8-fceb129469e2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([52, 12])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "memo_input = tokenizer.get_text_batch_encoding([my_first_text, my_first_text[10:30]])\n",
    "\n",
    "memo_input['input_ids'].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "cba11bf1-12c3-4b6e-96e0-7009ffb0d60f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[    0,     0,     0,     0,     0,     0,     0,     0,     0, 38577,\n",
       "         17622,  1073],\n",
       "        [  372,     8,  9718,    74,   843,   936,  4164, 43876, 41380,   258,\n",
       "           367,   727],\n",
       "        [ 5507,   313, 15723,   445,  2721,    13,  3435,  3414,   358,  3381,\n",
       "         15410,    26],\n",
       "        [ 9776,  1266,    74,    13,   337, 11703,   639, 39337,  1638,  1540,\n",
       "            10, 12187],\n",
       "        [  440,  2314,  4173,   299,  8913,  2942,   250,   352,  6770,    80,\n",
       "            13,  2248],\n",
       "        [  861,   410,   372, 32924,  1073, 33813,   445,  2721,   299,  2248,\n",
       "            80,  1484],\n",
       "        [ 1073,   659,  4611,  1073,   391,   300,   466,  5711, 14804,  1431,\n",
       "           304, 19702],\n",
       "        [   74,    15, 14929,  1327,  1323, 10081, 24843, 15438,   412, 16406,\n",
       "         38055,  9821],\n",
       "        [ 3737,  1073,   391,   300,   466,  5711, 39814,   260,   770,  5991,\n",
       "           313,  1962],\n",
       "        [18006, 22217, 42722, 10863,   262,  7958,  1593, 12704,  5940,  2719,\n",
       "           538, 21287]])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "memo_input['input_ids'][0:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "59a26e8e-22f3-499a-bba2-f3f714debdb1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([52, 12, 1024]), torch.Size([52, 12, 1024]))"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "input_embeddings = embedding.encode(memo_input['input_ids'])\n",
    "output_symbols = embedding.encode(memo_input['labels'])\n",
    "\n",
    "input_embeddings.shape, output_symbols.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "3adbd482-ce71-4d86-a049-a9ef8737ee31",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([52, 12])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[    0,     0,     0,     0,     0,     0,     0,     0,     0, 38577,\n",
       "         17622,  1073],\n",
       "        [  372,     8,  9718,    74,   843,   936,  4164, 43876, 41380,   258,\n",
       "           367,   727],\n",
       "        [ 5507,   313, 15723,   445,  2721,    13,  3435,  3414,   358,  3381,\n",
       "         15410,    26],\n",
       "        [ 9776,  1266,    74,    13,   337, 11703,   639, 39337,  1638,  1540,\n",
       "            10, 12187],\n",
       "        [  440,  2314,  4173,   299,  8913,  2942,   250,   352,  6770,    80,\n",
       "            13,  2248],\n",
       "        [  861,   410,   372, 32924,  1073, 33813,   445,  2721,   299,  2248,\n",
       "            80,  1484],\n",
       "        [ 1073,   659,  4611,  1073,   391,   300,   466,  5711, 14804,  1431,\n",
       "           304, 19702],\n",
       "        [   74,    15, 14929,  1327,  1323, 10081, 24843, 15438,   412, 16406,\n",
       "         38055,  9821],\n",
       "        [ 3737,  1073,   391,   300,   466,  5711, 39814,   260,   770,  5991,\n",
       "           313,  1962],\n",
       "        [18006, 22217, 42722, 10863,   262,  7958,  1593, 12704,  5940,  2719,\n",
       "           538, 21287]])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "decoded, _ = embedding.decode(input_embeddings)\n",
    "print(decoded.shape)\n",
    "\n",
    "decoded[0:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "616e69f2-501c-41fc-8f7e-21b6b3ae2f9c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[True, True, True, True, True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True, True, True],\n",
       "        [True, True, True, True, True, True, True, True, True, True, True, True]])"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "decoded[0:10] == memo_input['input_ids'][:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "776eb6ee-c5ec-4b24-afb2-127832d4fc04",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],\n",
       "        [ 0.0000,  0.9983, -0.0050,  ..., -0.0302, -0.0083, -0.0071],\n",
       "        [ 0.0000, -0.0050,  0.9900,  ...,  0.0180,  0.0399,  0.0036],\n",
       "        ...,\n",
       "        [ 0.0000, -0.0302,  0.0180,  ...,  0.9862,  0.0286, -0.0203],\n",
       "        [ 0.0000, -0.0083,  0.0399,  ...,  0.0286,  1.0189, -0.0050],\n",
       "        [ 0.0000, -0.0071,  0.0036,  ..., -0.0203, -0.0050,  1.0564]])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(50248.9219)\n",
      "tensor(-1119.1719)\n"
     ]
    }
   ],
   "source": [
    "sims = embedding.weight @ embedding.weight.T\n",
    "display(sims)\n",
    "diag_sum = torch.sum(sims[1:, 1: ].diag()) # almost 1 in each entry\n",
    "print(diag_sum) # obs vs expected\n",
    "print(torch.sum(sims[1:, 1:]) - diag_sum) #almost 0... more or less"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68109c97-0d07-4737-a2e1-1506927e7897",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "d65c6b2f-7994-4035-bf57-5e389e6c0db7",
   "metadata": {},
   "source": [
    "## Test layer and MeMo CMM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "851420e5-b642-4196-9d39-81ec08bbffa4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from MeMoPyTorch.modelling_memo_layer import MeMoLayer, ProjectionSequence, ProjectionTokens, CorrelationMatrixMemory"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca806ba1-8230-4e22-8171-a96695a0059c",
   "metadata": {},
   "source": [
    "### Check initialization of each matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "f6ef81f9-02b6-4236-9542-7ec40f80a333",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([4096, 1024]) (trasposed wrt saved one) in_features=4096, out_features=1024\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(tensor([1.0001, 0.9734, 1.0113,  ..., 1.0054, 1.0096, 1.0599],\n",
       "        grad_fn=<DiagonalBackward0_copy>),\n",
       " tensor([0.2441, 0.2397, 0.2492,  ..., 0.2716, 0.2557, 0.2687],\n",
       "        grad_fn=<DiagonalBackward0_copy>))"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "d,h,l = 1024, 4, 3\n",
    "proj = ProjectionSequence(d, d*h)\n",
    "print(proj.weight.shape, proj.extra_repr())\n",
    "(proj.weight.T @ proj.weight).diag(), (proj.weight @ proj.weight.T).diag()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "b310ff39-9e6c-420e-9062-d83abfa8ef81",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([1.0071, 1.0193, 0.9939,  ..., 1.0096, 0.9539, 0.9837]),\n",
       " tensor([0.2454, 0.2464, 0.2486,  ..., 0.2683, 0.2605, 0.2546]))"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Prj = torch.normal(0, 1/math.sqrt(d*h), size=(d,d*h))\n",
    "Prj = torch.transpose(Prj, 0, 1)\n",
    "(Prj.T @ Prj).diag(), (Prj @ Prj.T).diag()"
   ]
  },
  {
   "cell_type": "raw",
   "id": "bfe9a656-6558-47a2-a49a-1fcf73865c97",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "37f1066d-be0d-4772-b06a-4b2cc02fe98f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1024\n",
      "4\n",
      "torch.Size([256, 1024]) in_features=1024, out_features=256\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(tensor([0.9130, 1.1256, 0.8835,  ..., 1.0044, 0.9630, 1.0328],\n",
       "        grad_fn=<DiagonalBackward0_copy>),\n",
       " tensor([3.9768, 4.1737, 4.2684, 4.1132, 4.0868, 4.2425, 4.2230, 3.8434, 4.0362,\n",
       "         4.0547, 4.0841, 4.1058, 3.8918, 4.0659, 3.9973, 3.9419, 3.9983, 4.0927,\n",
       "         4.3244, 3.9605, 3.6646, 4.1196, 3.8822, 3.8914, 3.8621, 3.8363, 4.0533,\n",
       "         3.8372, 3.9975, 4.1140, 3.9401, 4.1465, 3.9762, 4.3454, 3.9064, 4.2744,\n",
       "         3.7955, 3.8557, 3.9269, 4.0063, 3.9984, 4.1154, 4.1589, 3.9601, 3.6307,\n",
       "         4.0627, 3.7175, 4.1304, 3.7287, 4.1450, 4.0616, 4.3163, 4.1937, 3.9228,\n",
       "         4.3139, 3.6658, 4.0031, 3.9427, 4.0584, 3.8599, 4.1947, 4.1495, 3.8413,\n",
       "         4.2497, 3.9633, 3.5159, 3.7868, 4.0483, 4.1109, 4.2331, 4.2501, 3.8947,\n",
       "         3.8700, 4.1005, 4.0520, 3.9067, 3.8875, 3.9716, 4.1152, 4.0778, 3.7914,\n",
       "         4.0581, 4.0259, 3.6324, 4.0859, 3.7057, 4.0759, 4.3153, 3.7562, 4.2928,\n",
       "         3.8849, 4.1671, 4.0208, 4.0959, 4.4941, 4.1431, 4.0250, 3.8631, 4.1924,\n",
       "         4.0749, 3.9606, 4.3228, 4.3304, 3.8730, 3.7526, 3.5913, 3.9791, 3.9778,\n",
       "         4.0275, 4.1521, 3.9836, 4.0522, 3.8879, 4.0047, 4.0402, 4.2642, 3.9911,\n",
       "         3.8229, 3.9110, 4.3203, 4.1355, 4.1162, 3.6573, 4.2348, 3.9384, 4.1847,\n",
       "         4.0710, 3.8432, 4.0793, 4.2618, 4.1772, 3.7868, 4.1433, 3.9188, 4.3340,\n",
       "         4.1209, 4.2902, 4.0526, 4.4291, 3.4608, 3.9765, 4.0040, 3.9520, 4.2562,\n",
       "         4.0864, 3.9089, 4.2130, 3.9969, 3.9996, 4.1191, 4.2089, 4.2102, 4.1928,\n",
       "         4.2302, 4.0049, 3.9763, 3.8184, 4.1753, 4.2091, 3.8947, 3.8695, 3.8703,\n",
       "         4.4131, 3.7003, 3.9282, 4.0047, 4.0285, 4.1500, 3.7685, 3.9728, 3.8149,\n",
       "         4.2218, 3.7093, 3.7745, 3.8366, 4.0434, 4.0631, 3.9403, 3.9097, 4.1317,\n",
       "         4.2786, 3.8261, 3.9639, 3.8418, 4.1069, 3.9730, 3.9540, 3.8077, 3.8001,\n",
       "         4.2947, 4.0457, 3.8041, 4.0149, 3.7611, 4.0451, 3.9350, 4.0928, 4.0308,\n",
       "         3.7998, 4.2200, 4.0626, 4.0589, 4.2061, 3.7565, 4.0606, 4.0845, 3.7497,\n",
       "         3.8644, 4.0558, 3.8699, 3.7022, 3.8193, 4.1033, 4.2067, 4.1566, 4.1678,\n",
       "         3.6865, 3.7897, 4.0604, 3.8512, 3.8398, 4.0965, 3.6062, 4.0845, 4.2356,\n",
       "         3.9543, 4.0724, 4.2175, 3.8336, 3.8921, 3.8776, 4.1143, 4.1395, 3.9572,\n",
       "         3.9234, 4.2603, 4.0341, 4.0701, 4.1503, 3.7909, 4.3133, 3.8848, 3.6730,\n",
       "         4.0135, 3.8419, 4.0423, 4.1062, 3.6979, 4.0586, 3.9010, 4.0482, 3.7967,\n",
       "         3.5832, 4.5566, 3.9330, 4.0356], grad_fn=<DiagonalBackward0_copy>))"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(d)\n",
    "print(h)\n",
    "\n",
    "d_k = d // h\n",
    "\n",
    "W_v = ProjectionTokens(d, d_k)\n",
    "\n",
    "print(W_v.weight.shape, W_v.extra_repr())\n",
    "\n",
    "\n",
    "### always used transposed! so check with .T\n",
    "(W_v.weight.T @ W_v.weight).diag(), (W_v.weight @ W_v.weight.T).diag()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "b6d736ff-f9d1-4fbf-9f28-5e939fd16192",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([3.8353, 3.9763, 4.3228, 3.8681, 3.8375, 4.1265, 3.8689, 3.8987, 4.2349,\n",
       "         3.8796, 3.9763, 3.6021, 3.7336, 3.6950, 3.9160, 4.1376, 4.0095, 4.1417,\n",
       "         3.9849, 4.2199, 3.8680, 3.6905, 3.9728, 4.2266, 4.3437, 4.0601, 3.8447,\n",
       "         3.8382, 4.0107, 3.7765, 3.9026, 3.9907, 3.7311, 4.1689, 4.1451, 3.8639,\n",
       "         3.8962, 4.3349, 3.9581, 3.9291, 3.9250, 3.9108, 3.7227, 3.8750, 4.0860,\n",
       "         3.9869, 4.1752, 4.0651, 3.8667, 4.1384, 4.0423, 4.1828, 3.7566, 4.0609,\n",
       "         3.9843, 4.1172, 4.0428, 4.3638, 4.0453, 4.1944, 3.9861, 3.9967, 3.5535,\n",
       "         3.8496, 3.6559, 4.0801, 4.0982, 4.0738, 4.3134, 3.9767, 4.0875, 4.1507,\n",
       "         3.6378, 4.0853, 3.9793, 4.2400, 4.0486, 4.2980, 3.8818, 4.0489, 4.1768,\n",
       "         3.7635, 4.0278, 3.7948, 3.8342, 4.1616, 4.3265, 3.9240, 3.7758, 4.0395,\n",
       "         3.8655, 3.8947, 4.0042, 3.8414, 3.9757, 3.8218, 3.8125, 4.1606, 3.9218,\n",
       "         4.0527, 3.9117, 4.0999, 3.9146, 3.8490, 4.1715, 3.9676, 4.0928, 3.6956,\n",
       "         3.9751, 4.1215, 4.2063, 4.0846, 3.7404, 4.1148, 3.9799, 4.2358, 3.9574,\n",
       "         3.9888, 3.7914, 4.0309, 4.2669, 3.9133, 3.6545, 3.9479, 4.0787, 3.9721,\n",
       "         3.9840, 4.1040, 3.9432, 4.1900, 3.6693, 3.8747, 4.3377, 3.9578, 3.6932,\n",
       "         3.7772, 4.2170, 3.9374, 4.1474, 3.9566, 3.8984, 4.0653, 3.8625, 4.1531,\n",
       "         4.0030, 4.1917, 4.0490, 4.1464, 3.6995, 4.1988, 3.9048, 4.1751, 3.8209,\n",
       "         3.7138, 4.1377, 4.0773, 4.1703, 4.0533, 3.9918, 3.9047, 4.1181, 3.9627,\n",
       "         4.2953, 3.6410, 3.8896, 4.1702, 4.0069, 3.7293, 4.2583, 3.6601, 3.8508,\n",
       "         4.1348, 3.7614, 3.9681, 4.0771, 3.9354, 3.7495, 4.2745, 4.2349, 4.0858,\n",
       "         4.0110, 4.2233, 3.8109, 3.9742, 3.8753, 3.9954, 4.0722, 3.9852, 4.1820,\n",
       "         4.1669, 4.0183, 3.9887, 3.9640, 4.0965, 3.9221, 4.0630, 3.9899, 4.1322,\n",
       "         3.5362, 3.7825, 4.2704, 3.8348, 3.9129, 4.0327, 3.9213, 4.2357, 4.0316,\n",
       "         4.0401, 4.0129, 3.8732, 3.5990, 3.9698, 3.9882, 3.6947, 4.0470, 4.0765,\n",
       "         3.7347, 3.9067, 3.9219, 4.0172, 4.1780, 4.0201, 4.3322, 3.8193, 4.0039,\n",
       "         4.3300, 4.0767, 4.1337, 3.6960, 4.4786, 3.9141, 3.7216, 3.8459, 3.9324,\n",
       "         3.8382, 4.1555, 3.9672, 4.0571, 4.1875, 3.9470, 4.0674, 4.0975, 3.8496,\n",
       "         3.7938, 4.1240, 4.0696, 4.3337, 3.9160, 4.2105, 4.1791, 4.2121, 3.7637,\n",
       "         3.9771, 3.9175, 4.0455, 4.1539]),\n",
       " tensor([1.1023, 0.9831, 1.0222,  ..., 0.9989, 0.9508, 1.1465]))"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "W_v_single_head = torch.normal(0, 1/math.sqrt(d_k), size=(d,d_k))\n",
    "\n",
    "(W_v_single_head.T @ W_v_single_head).diag(), (W_v_single_head @ W_v_single_head.T).diag()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c545596f-fbb2-4e0b-b3f1-1e4f37ebc903",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3fd69212-c26e-4404-95a6-7442d6d814ea",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "3ff823bd-ce40-49fd-b67b-5c12dc739614",
   "metadata": {},
   "source": [
    "### Check memorization on single layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "10a33880-0afe-4de6-a594-682be1581cee",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "MeMoLayer(\n",
       "  (W_v_single_head): ProjectionTokens(in_features=1024, out_features=256)\n",
       "  (Prj): ProjectionSequence((trasposed wrt saved one) in_features=4096, out_features=1024)\n",
       "  (CMM): CorrelationMatrixMemory(in_features=1024, out_features=1024)\n",
       ")"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "d,h,l = 1024, 4, 3\n",
    "\n",
    "layer = MeMoLayer(d, h)\n",
    "layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "6fdbe2b1-b421-4f7e-8f61-18a82c9b0eb1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(52, 12, 1024)"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batch_size, current_length, d = input_embeddings.shape\n",
    "batch_size, current_length, d "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "7664712d-7ce1-4d64-b1f4-54abb019ffee",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([52, 12, 1024])"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output_symbols.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "f8c91851-c546-48bf-90c6-bb4c7ed5484c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(4096) torch.Size([4, 1024])\n"
     ]
    }
   ],
   "source": [
    "current_length = int(input_embeddings.shape[1]/ h)\n",
    "\n",
    "input_sequence = input_embeddings.reshape((batch_size, current_length, h, d))\n",
    "\n",
    "current_output_symbols = output_symbols[:, [(x+1)*h-1 for x in range(0,current_length)]]\n",
    "j = 2 \n",
    "print(sum(sum(input_sequence[0][j] == input_embeddings[0][4*j:4*(j+1)])), input_sequence[0][j].shape)\n",
    "\n",
    "(batch_size, blocks,h,d) = input_sequence.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "99480537-8c54-4dae-85dd-838ce5bfabb9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([52, 3, 4, 1024]), torch.Size([52, 3, 1024]))"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "input_sequence.shape, current_output_symbols.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "2a2ba406-c0a0-4008-9aa0-836372382cca",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "MeMoLayer(\n",
       "  (W_v_single_head): ProjectionTokens(in_features=1024, out_features=256)\n",
       "  (Prj): ProjectionSequence((trasposed wrt saved one) in_features=4096, out_features=1024)\n",
       "  (CMM): CorrelationMatrixMemory(in_features=1024, out_features=1024)\n",
       ")"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "layer = MeMoLayer(d, h)\n",
    "display(layer)\n",
    "## update the input sequence for the next layer\n",
    "_, seq_encoding_for_the_last_layer = layer.memorize(input_sequence, current_output_symbols, is_last=False)\n",
    "layer.directly_memorize(seq_encoding_for_the_last_layer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "4c46ccae-66bb-4e72-abaa-b7abdfd0839c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([52, 3, 4, 1024]), torch.Size([3, 4, 1024]))"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "input_sequence.shape, input_sequence[3].shape # batch (52 elements of chunks 4*4*1024)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "8467cfe3-4e46-45ea-b3fc-bac785a84ebb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([52, 1024])\n"
     ]
    }
   ],
   "source": [
    "_, seq_encoding_for_the_last_layer = layer.retrieve(input_sequence)\n",
    "\n",
    "print(seq_encoding_for_the_last_layer.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "4a3cd023-37a9-426a-9df2-b7d25a459ce8",
   "metadata": {},
   "outputs": [],
   "source": [
    "logits = layer.directly_retrieve(seq_encoding_for_the_last_layer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "32a917d5-4a51-4163-9bac-21f960de5290",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([48505, 20110,  1108, 48019,    80, 19216,  9718,  1113,  4927,    66,\n",
      "        19216,  1448,  2122, 41530,   187,  4172,   246,   659, 10986, 30975,\n",
      "           80, 12931,   352, 14134,  2721,   258,  8830,    87,   826,    15,\n",
      "        17532,   729, 26798, 41070,  6575,   299,   266,  3737, 20889,   287,\n",
      "          512,   354,   250,   247,    70,   275, 16128,  2680,    74, 13679,\n",
      "           15,   209]) tensor([0.7027, 1.0707, 0.9809, 1.0174, 0.8847, 1.1506, 1.0516, 0.9499, 1.0354,\n",
      "        1.7528, 1.2026, 1.0802, 0.9133, 1.0406, 0.9373, 0.9744, 0.9165, 0.9665,\n",
      "        1.0506, 1.0039, 0.9117, 0.9218, 0.8659, 0.9423, 1.2641, 0.9583, 1.0097,\n",
      "        1.1661, 1.1091, 0.9234, 0.8966, 1.0579, 1.1814, 0.9961, 0.9950, 0.9661,\n",
      "        1.0283, 1.1281, 0.8819, 0.9572, 0.9551, 1.0385, 0.9571, 1.0440, 1.0158,\n",
      "        1.0891, 1.0594, 0.9931, 1.0877, 0.9315, 0.9413, 0.9924],\n",
      "       grad_fn=<MaxBackward0>)\n"
     ]
    }
   ],
   "source": [
    "retreived_output_symbol_vector, m = embedding.decode(logits)\n",
    "print(retreived_output_symbol_vector, m)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "1acd6b66-770e-42d1-b01a-cce2d9bd2523",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([48505, 20110,  1108, 48019,    80, 19216,  9718,  1113,  4927,    66,\n",
       "        19216,  1448,  2122, 41530,   187,  4172,   246,   659, 10986, 30975,\n",
       "           80, 12931,   352, 14134,  2721,   258,  8830,    87,   826,    15,\n",
       "        17532,   729, 26798, 41070,  6575,   299,   266,  3737, 20889,   287,\n",
       "          512,   354,   250,   247,    70,   275, 16128,  2680,    74, 13679,\n",
       "           15,   209])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(52) over torch.Size([52])\n"
     ]
    }
   ],
   "source": [
    "o = embedding.decode(current_output_symbols[:, -1])[0]\n",
    "display(o)\n",
    "\n",
    "print(sum(o == retreived_output_symbol_vector), 'over', retreived_output_symbol_vector.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "4fe7568a-9319-4beb-b56c-89c8fd366f56",
   "metadata": {},
   "outputs": [],
   "source": [
    "#### test single block"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "370fac45-4707-43d5-88f8-df78ee1c931c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([3, 4, 1024]) torch.Size([3, 1024])\n",
      "torch.Size([4, 1024]) torch.Size([1024])\n"
     ]
    }
   ],
   "source": [
    "print(input_sequence[3].shape, current_output_symbols[3].shape)\n",
    "print(input_sequence[3][0].shape, current_output_symbols[3][0].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30d28e7b-29c7-4368-9111-360fd75ab9b8",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "01cd5f78-d3a8-4be2-952a-2db537b10110",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "153/156\n"
     ]
    }
   ],
   "source": [
    "total = 0\n",
    "correct = 0 \n",
    "\n",
    "for batch_index in range(len(memo_input['input_ids'])):\n",
    "    #print(\"input ids\", memo_input[0]['input_ids'][batch_index])\n",
    "    #print()\n",
    "    \n",
    "    for i in range(len(current_output_symbols[batch_index])):\n",
    "        #display(embedding.decode(input_sequence[batch_index][i]), embedding.decode(current_output_symbols[batch_index][i]))\n",
    "        true = embedding.decode(current_output_symbols[batch_index][i])[0].item()\n",
    "        \n",
    "        _, seq_encoding_for_the_last_layer  = layer.retrieve(input_sequence[batch_index][i].unsqueeze(0).unsqueeze(0))\n",
    "        \n",
    "        retreived_output_symbol_vector, m = embedding.decode(layer.directly_retrieve(seq_encoding_for_the_last_layer))\n",
    "        pred = retreived_output_symbol_vector.item()\n",
    "\n",
    "        total += 1\n",
    "        correct += pred == true\n",
    "\n",
    "print(f\"{correct}/{total}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "592772e2-6b27-4071-bd5f-a441876381fd",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "7b22599b-d21d-43bc-b728-dbadd0e4263d",
   "metadata": {},
   "source": [
    "### Check with batch size of 1 and output probs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "30176822-1cf4-4b5e-8fbc-7e4375fdecab",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[2520,  310,  247, 1071,  323,  247, 1077, 2159, 2159, 3425,  273, 1249]]) tensor([[  310,   247,  1071,   323,   247,  1077,  2159,  2159,  3425,   273,\n",
      "          1249, 21761]])\n",
      "(tensor([[  323,  2159, 21761]]), tensor([[0.9773, 1.0502, 1.0602]]))\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(torch.Size([1, 3, 4, 1024]), torch.Size([1, 3, 1024]))"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "memo_input = tokenizer.get_text_batch_encoding(['this is a test for a very short short sequence of 12 tokens'])\n",
    "input_ids, labels = memo_input['input_ids'], memo_input['labels']\n",
    "print(input_ids, labels)\n",
    "\n",
    "input_embeddings = embedding.encode(input_ids)\n",
    "#print(input_embeddings.shape)\n",
    "\n",
    "output_embeddings = embedding.encode(labels)\n",
    "#print(output_embeddings.shape)\n",
    "\n",
    "\n",
    "current_length = max_length\n",
    "\n",
    "current_length = int(current_length/h)\n",
    "input_sequence = input_embeddings.reshape((1, current_length, h, d))\n",
    "\n",
    "output_symbols = output_embeddings[:, [(x+1)*h-1 for x in range(0,current_length)]] ## the output symbol is always the same tokem?\n",
    "print(embedding.decode(output_symbols))\n",
    "input_sequence.shape, output_symbols.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "c4968bbc-e2df-44ca-93ac-9610a1ab071c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[[2520,  310,  247, 1071],\n",
       "          [ 323,  247, 1077, 2159],\n",
       "          [2159, 3425,  273, 1249]]]),\n",
       " tensor([[  323,  2159, 21761]]),\n",
       " tensor([[2520,  310,  247, 1071,  323,  247, 1077, 2159, 2159, 3425,  273, 1249]]))"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "embedding.decode(input_sequence)[0], embedding.decode(output_symbols)[0], input_ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "8cfaf35a-6ff8-4979-9727-cedcb78d8c9e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 3, 4, 1024])"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "input_sequence.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "ccc5bed4-5902-4f96-bab2-f1eefe5207c6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 3, 1024])"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output_symbols.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "813fecba-8575-4e2b-8271-e94af40c2600",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "MeMoLayer(\n",
       "  (W_v_single_head): ProjectionTokens(in_features=1024, out_features=256)\n",
       "  (Prj): ProjectionSequence((trasposed wrt saved one) in_features=4096, out_features=1024)\n",
       "  (CMM): CorrelationMatrixMemory(in_features=1024, out_features=1024)\n",
       ")"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 1024])\n",
      "tensor([323]) tensor([0.9535], grad_fn=<MaxBackward0>)\n",
      "torch.Size([1, 1024])\n",
      "tensor([2159]) tensor([1.1330], grad_fn=<MaxBackward0>)\n",
      "torch.Size([1, 1024])\n",
      "tensor([21761]) tensor([1.0683], grad_fn=<MaxBackward0>)\n"
     ]
    }
   ],
   "source": [
    "layer = MeMoLayer(d, h)\n",
    "display(layer)\n",
    "\n",
    "## update the input sequence for the next layer\n",
    "_, seq_encoding_for_the_last_layer = layer.memorize(input_sequence, output_symbols, is_last=False)\n",
    "layer.directly_memorize(seq_encoding_for_the_last_layer)\n",
    "\n",
    "for i in range(0,3):\n",
    "    _, seq_encoding_for_the_last_layer = layer.retrieve(input_sequence[0][i].unsqueeze(0).unsqueeze(0))\n",
    "    print(seq_encoding_for_the_last_layer.shape)\n",
    "                                                        \n",
    "    retreived_output_symbol_vector, m = embedding.decode(layer.directly_retrieve(seq_encoding_for_the_last_layer))\n",
    "    print(retreived_output_symbol_vector, m)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "364c89cd-cb85-4742-a214-c1fe1812e0de",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "80c3e709-edf1-4366-9766-ed6403ecedb9",
   "metadata": {},
   "source": [
    "## Test the entire MeMo model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "6a73e8f4-a4d6-4dcf-9698-3c9853976b8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from MeMoPyTorch.modelling_memo import MeMo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "9113cb02-449d-4ff0-b0d4-0cf121004e4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from MeMoPyTorch.modelling_memo_tokenizer import MeMoTokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "0a3e1627-b177-455c-9a81-7ab248fb1fc0",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"testo_di_prova.txt\") as my_first_text_f:\n",
    "    my_first_text = my_first_text_f.read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "c58b7e91-677e-44f4-a738-93b3d3bb9ce6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "384 4\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The tokenizer class you load from this checkpoint is not the same type as the class this function is called from. It may result in unexpected tokenization. \n",
      "The tokenizer class you load from this checkpoint is 'GPTNeoXTokenizer'. \n",
      "The class this function is called from is 'MeMoTokenizer'.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Setting pad token and pad token id = <|endoftext|>, 0\n"
     ]
    }
   ],
   "source": [
    "max_length = 384\n",
    "print(max_length, h)\n",
    "tokenizer = MeMoTokenizer.from_pretrained(\"EleutherAI/gpt-neox-20b\", \n",
    "                                          padding_side='left', truncation_side='left', \n",
    "                                          max_length=max_length, head_number=h)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "6cd21d5d-2e4e-4b8d-aebc-b8f78421a038",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 384])"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "memo_input = tokenizer.get_text_batch_encoding(my_first_text)\n",
    "memo_input['labels'].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "bacfc15c-56ea-49e3-9628-cec62cc9019e",
   "metadata": {},
   "outputs": [],
   "source": [
    "device='cuda:0' #'cpu'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "2e7021d0-26d9-448b-ba86-23c38c87499c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1024, 4, 3)"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "d,h,l"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "8c2b3fb1-8a1c-4678-8831-da82aeece4fe",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MeMo embedding initilialization\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "MeMo(\n",
       "  (encoder): MeMoEmbedding(50254, 1024, padding_idx=0)\n",
       "  (layers): ModuleList(\n",
       "    (0-2): 3 x MeMoLayer(\n",
       "      (W_v_single_head): ProjectionTokens(in_features=1024, out_features=256)\n",
       "      (Prj): ProjectionSequence((trasposed wrt saved one) in_features=4096, out_features=1024)\n",
       "      (CMM): CorrelationMatrixMemory(in_features=1024, out_features=1024)\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = MeMo(inner_dim=d, \n",
    "             num_of_heads=h, \n",
    "             num_of_layers=3, \n",
    "             chunk_length=max_length, \n",
    "             num_embeddings=tokenizer.vocab_size, \n",
    "             padding_idx=tokenizer.pad_token_id, \n",
    "             device=device) #MeMoModel\n",
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "850ad4cc-61c2-45f3-bcf3-2114468331bd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([2, 384, 1024]), torch.Size([2, 384, 1024]), 384)"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "input_sequence =  model.encoder.encode(memo_input['input_ids'])\n",
    "output_symbols = model.encoder.encode(memo_input['labels'])\n",
    "\n",
    "(batch_size, current_length, d) = input_sequence.shape\n",
    "last_layer = model.layers[model.l-1]\n",
    "\n",
    "current_length = model.chunk_length\n",
    "\n",
    "input_sequence.shape, output_symbols.shape, current_length"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "6c1208c5-8186-4740-8d3c-65b57a792cbf",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[0, 1, 2, 3], [1, 2, 3, 4], [2, 3, 4, 5], [3, 4, 5, 6], [4, 5, 6, 7]]"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "layer_level = 0\n",
    "input_index = [[j for j in range(i - model.h ** (layer_level + 1), i, model.h ** ((layer_level + 1) - 1))] \n",
    "               for i in range(model.h ** (layer_level + 1), current_length + 1)]\n",
    "input_index[0:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "f8d5d83c-8633-49c3-ac64-83088339792b",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([2, 384, 1024]), torch.Size([2, 384, 1024]))"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "original_input_seq = torch.clone(input_sequence)\n",
    "original_out = torch.clone(output_symbols)\n",
    "\n",
    "original_input_seq.shape, original_out.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e689a9e-ab03-41eb-9fc6-832c6f95080f",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "81da7dce-28c9-4dcf-a6ba-3f7ea6a0ae3d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n",
      "torch.Size([2, 381, 1024])\n",
      "torch.Size([2, 381, 4, 1024])\n"
     ]
    }
   ],
   "source": [
    "#for layer_level in range(model.l):\n",
    "layer_level = 0 #1,2\n",
    "          \n",
    "print(model.h ** (layer_level + 1) < current_length + 1)\n",
    "## update the input sequence for the next layer\n",
    "layer_output_idxs = [i - model.h ** ((layer_level + 1) - 1) for i in range(model.h ** (layer_level + 1), current_length + 1)]\n",
    "output_symbols = output_symbols[:, layer_output_idxs]\n",
    "print(output_symbols.shape)\n",
    "\n",
    "input_index = [[j for j in range(i - model.h ** (layer_level + 1), i, model.h ** ((layer_level + 1) - 1))] \n",
    "               for i in range(model.h ** (layer_level + 1), current_length + 1)]\n",
    "\n",
    "\n",
    "input_sequence = input_sequence[:, input_index]\n",
    "print(input_sequence.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "68743a34-c4ff-46ed-82e1-e319cd3391b5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "        True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "        True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "        True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "        True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "        True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "        True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "        True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "        True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "        True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "        True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "        True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "        True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "        True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "        True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "        True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "        True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "        True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "        True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "        True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "        True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "        True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "        True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "        True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "        True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "        True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "        True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "        True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "        True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "        True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "        True, True, True, True, True, True, True, True, True, True, True, True,\n",
       "        True, True, True, True, True, True, True, True, True], device='cuda:0')"
      ]
     },
     "execution_count": 53,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.encoder.decode(output_symbols[0])[0] == model.encoder.decode(torch.stack([original_out[0][i] for i in layer_output_idxs]))[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19285dec-a99a-419b-85c8-22985b57a009",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c237394a-cb47-413a-bac2-5a52333b243d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c2a7f82-162f-4089-add0-bdf07295d30a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "b74e5e8f-2f0e-4838-87ca-3a55ae3b8081",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tqdm\n",
    "\n",
    "\n",
    "class Evaluation:\n",
    "    def check_memorization(self, model, tokenizer, text, # device='cpu',\n",
    "                           starting_point=None):\n",
    "        if starting_point == None:\n",
    "            basic_block = model.h ** model.l\n",
    "        else:\n",
    "            basic_block = starting_point\n",
    "        \n",
    "        \n",
    "        input_ = tokenizer(my_first_text, padding='longest', truncation='do_not_truncate', max_length=None)\n",
    "        input_ = tokenizer.pad(input_, pad_to_multiple_of=basic_block)\n",
    "        input_ids = input_['input_ids']\n",
    "                \n",
    "        count = 0\n",
    "        correct = 0\n",
    "        max_length = tokenizer.max_length\n",
    "        (batch_size, number_of_tokens) = input_ids.shape\n",
    "\n",
    "        #print(f\"(batch_size, number_of_tokens) = {(batch_size, number_of_tokens)}\")\n",
    "        \n",
    "        for i in tqdm.tqdm(range(basic_block,  number_of_tokens - 1)):\n",
    "            text_tokens = input_ids[:, i - basic_block:i]\n",
    "            \n",
    "            (batch_size, number_of_tokens) = text_tokens.shape\n",
    "            \n",
    "            text_tokens = torch.concat((torch.zeros((batch_size, max_length-1-number_of_tokens), \n",
    "                                                    dtype=torch.int), \n",
    "                                        text_tokens), axis=1\n",
    "                                      )\n",
    "            \n",
    "            #print(i - basic_block, i)\n",
    "            out, max_value = model.retrieve(text_tokens)\n",
    "            #print(out, input_ids[:, i])\n",
    "            #print(out[0].item())\n",
    "            \n",
    "            count += batch_size\n",
    "            correct += torch.sum(out.to('cpu') == input_ids[:, i])\n",
    "        \n",
    "                           \n",
    "        return correct / count\n",
    "\n",
    "    def check_pretokenized(self, model, tokenizer, input_ids,# device='cpu',\n",
    "                           starting_point=None):\n",
    "        if starting_point == None:\n",
    "            basic_block = model.h ** model.l\n",
    "        else:\n",
    "            basic_block = starting_point\n",
    "                \n",
    "        count = 0\n",
    "        correct = 0\n",
    "        max_length = tokenizer.max_length\n",
    "        (batch_size, number_of_tokens) = input_ids.shape\n",
    "\n",
    "        #print(f\"(batch_size, number_of_tokens) = {(batch_size, number_of_tokens)}\")\n",
    "        \n",
    "        for i in tqdm.tqdm(range(basic_block,  number_of_tokens - 1)):\n",
    "            text_tokens = input_ids[:, i - basic_block:i]\n",
    "            \n",
    "            (batch_size, number_of_tokens) = text_tokens.shape\n",
    "            \n",
    "            text_tokens = torch.concat((torch.zeros((batch_size, max_length-1-number_of_tokens), \n",
    "                                                    dtype=torch.int), \n",
    "                                        text_tokens), axis=1\n",
    "                                      )\n",
    "            \n",
    "            #print(i - basic_block, i)\n",
    "            out, max_value = model.retrieve(text_tokens)\n",
    "            #print(out, input_ids[:, i])\n",
    "            #print(out[0].item())\n",
    "            \n",
    "            count += batch_size\n",
    "            correct += torch.sum(out.to('cpu') == input_ids[:, i])\n",
    "        \n",
    "                           \n",
    "        return correct / count\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "f8a2bd82-a24f-4274-a2bf-26170f8abfb4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MeMo embedding initilialization\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "torch.Size([16, 384])"
      ]
     },
     "execution_count": 55,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = MeMo(inner_dim=d, \n",
    "             num_of_heads=h, \n",
    "             num_of_layers=l, \n",
    "             chunk_length=max_length, \n",
    "             num_embeddings=tokenizer.vocab_size, \n",
    "             padding_idx=tokenizer.pad_token_id, \n",
    "             device=device)\n",
    "\n",
    "memo_input = tokenizer.get_text_batch_encoding([my_first_text]*8)\n",
    "\n",
    "\n",
    "memo_input['input_ids'].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2d0a856-4f26-40da-b600-fc8b2bdbe7b2",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "bca7a003-745f-418d-96c0-9f02bc9cce3a",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "model.memorize_text(memo_input)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "c12ee0c6-1d8e-402c-9078-17ae941c9ffe",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 319/319 [00:00<00:00, 737.06it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Degree of memorization: %f  tensor(0.8072)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "e = Evaluation()\n",
    "out = e.check_pretokenized(model, tokenizer, memo_input['input_ids'])\n",
    "print(\"Degree of memorization: %f \", out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "f431f7b1-27ac-4f7f-a476-d8f61102de19",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.forget_text(memo_input)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "186d2b13-ace1-4bee-9591-b6f5530e0a5f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 319/319 [00:00<00:00, 764.43it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Degree of memorization: %f  tensor(0.0815)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "out = e.check_pretokenized(model, tokenizer, memo_input['input_ids'])\n",
    "print(\"Degree of memorization: %f \", out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "c72f7371-61f8-4f80-b332-ec7d8703d855",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MeMo embedding initilialization\n",
      "CMM pre learning\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Parameter containing:\n",
       "tensor([[0., 0., 0.,  ..., 0., 0., 0.],\n",
       "        [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "        [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "        ...,\n",
       "        [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "        [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "        [0., 0., 0.,  ..., 0., 0., 0.]], device='cuda:0', requires_grad=True)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([2, 384])\n",
      "torch.Size([2, 384])\n",
      "torch.Size([2, 384])\n",
      "torch.Size([2, 384])\n",
      "torch.Size([2, 384])\n",
      "torch.Size([2, 384])\n",
      "torch.Size([2, 384])\n",
      "torch.Size([2, 384])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[ 1.0101e+00,  1.2280e-02,  1.7227e-02,  ..., -1.2366e-02,\n",
       "          1.4448e-02, -1.8096e-02],\n",
       "        [ 1.2280e-02,  9.7455e-01,  1.1633e-02,  ...,  5.7634e-03,\n",
       "          3.4262e-02,  2.4340e-02],\n",
       "        [ 1.7227e-02,  1.1633e-02,  9.8092e-01,  ...,  2.9026e-02,\n",
       "          2.8689e-03, -4.5811e-03],\n",
       "        ...,\n",
       "        [-1.2366e-02,  5.7634e-03,  2.9026e-02,  ...,  9.8428e-01,\n",
       "          9.1292e-04,  1.4481e-02],\n",
       "        [ 1.4448e-02,  3.4262e-02,  2.8689e-03,  ...,  9.1292e-04,\n",
       "          9.8084e-01,  5.4475e-03],\n",
       "        [-1.8096e-02,  2.4340e-02, -4.5811e-03,  ...,  1.4481e-02,\n",
       "          5.4475e-03,  1.0217e+00]])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "tensor([[ 1.7443e-02, -1.2456e-02, -5.1137e-03,  ...,  1.0043e-02,\n",
       "          1.6821e-02, -1.0690e-05],\n",
       "        [ 2.0084e-02,  2.4817e-02,  2.6959e-02,  ...,  9.1169e-03,\n",
       "         -2.1050e-02, -8.9389e-03],\n",
       "        [-4.9506e-03,  6.9840e-03, -2.0625e-03,  ..., -8.0438e-03,\n",
       "          3.0667e-03,  8.5841e-03],\n",
       "        ...,\n",
       "        [ 4.6615e-03,  7.4334e-03,  1.6063e-02,  ..., -2.8373e-02,\n",
       "         -7.0026e-03,  1.4033e-02],\n",
       "        [-7.4981e-03,  1.0286e-02,  1.4566e-02,  ...,  2.8840e-03,\n",
       "         -7.8137e-03,  2.1140e-02],\n",
       "        [-5.8409e-03,  2.4287e-03,  1.1041e-02,  ...,  2.5658e-02,\n",
       "         -6.4533e-03,  3.8758e-03]])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 319/319 [00:00<00:00, 796.28it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Degree of memorization: %f  tensor(0.9514)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 319/319 [00:00<00:00, 743.09it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Degree of memorization: %f  tensor(0.9514)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "model = MeMo(inner_dim=d, \n",
    "             num_of_heads=h, \n",
    "             num_of_layers=l, \n",
    "             chunk_length=max_length, \n",
    "             num_embeddings=tokenizer.vocab_size, \n",
    "             padding_idx=tokenizer.pad_token_id, \n",
    "             device=device)\n",
    "print(\"CMM pre learning\")\n",
    "display(model.layers[0].CMM.weight)\n",
    "\n",
    "\n",
    "bs = 8\n",
    "for b in range(bs):\n",
    "    memo_input = tokenizer.get_text_batch_encoding(my_first_text)\n",
    "    print(memo_input['input_ids'].shape)\n",
    "\n",
    "    model.memorize_text(memo_input)\n",
    "\n",
    "Prj = model.layers[0].Prj.weight.detach().cpu()\n",
    "CMM = model.layers[0].CMM.weight.detach().cpu()\n",
    "\n",
    "display(Prj.T @ Prj)\n",
    "display(CMM)\n",
    "\n",
    "e = Evaluation()\n",
    "out = e.check_pretokenized(model, tokenizer, memo_input['input_ids'])\n",
    "print(\"Degree of memorization: %f \", out)\n",
    "\n",
    "\n",
    "model.forget_text(memo_input)\n",
    "\n",
    "out = e.check_pretokenized(model, tokenizer, memo_input['input_ids'])\n",
    "print(\"Degree of memorization: %f \", out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "id": "67c1f7ed-02cd-4f13-aa01-e8821b11f302",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MeMo embedding initilialization\n",
      "CMM pre learning\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Parameter containing:\n",
       "tensor([[0., 0., 0.,  ..., 0., 0., 0.],\n",
       "        [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "        [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "        ...,\n",
       "        [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "        [0., 0., 0.,  ..., 0., 0., 0.],\n",
       "        [0., 0., 0.,  ..., 0., 0., 0.]], device='cuda:0', requires_grad=True)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "memorizing the same text iteration = 0\n",
      "memorizing the same text iteration = 1\n",
      "memorizing the same text iteration = 2\n",
      "memorizing the same text iteration = 3\n",
      "memorizing the same text iteration = 4\n",
      "memorizing the same text iteration = 5\n",
      "memorizing the same text iteration = 6\n",
      "memorizing the same text iteration = 7\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[ 1.0114e+00, -7.7352e-03,  3.1193e-02,  ...,  9.8616e-03,\n",
       "         -5.1539e-03,  8.5160e-03],\n",
       "        [-7.7352e-03,  1.0026e+00, -6.3145e-03,  ...,  1.2942e-02,\n",
       "          6.2130e-03, -3.2290e-03],\n",
       "        [ 3.1193e-02, -6.3145e-03,  9.9515e-01,  ...,  3.7192e-03,\n",
       "          8.0864e-04, -3.9914e-02],\n",
       "        ...,\n",
       "        [ 9.8616e-03,  1.2942e-02,  3.7192e-03,  ...,  9.9843e-01,\n",
       "         -1.8007e-02,  1.0587e-02],\n",
       "        [-5.1539e-03,  6.2130e-03,  8.0864e-04,  ..., -1.8007e-02,\n",
       "          9.6098e-01,  9.6622e-03],\n",
       "        [ 8.5160e-03, -3.2290e-03, -3.9914e-02,  ...,  1.0587e-02,\n",
       "          9.6622e-03,  9.9557e-01]])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "tensor([[-0.0053, -0.0290, -0.0211,  ..., -0.0158,  0.0136,  0.0058],\n",
       "        [-0.0051,  0.0091,  0.0132,  ...,  0.0086, -0.0140, -0.0399],\n",
       "        [-0.0129, -0.0180, -0.0300,  ...,  0.0116,  0.0044,  0.0201],\n",
       "        ...,\n",
       "        [ 0.0466,  0.0165, -0.0146,  ..., -0.0134,  0.0170,  0.0188],\n",
       "        [ 0.0155,  0.0005, -0.0193,  ..., -0.0363,  0.0022,  0.0193],\n",
       "        [ 0.0015,  0.0067, -0.0050,  ..., -0.0334, -0.0207,  0.0009]])"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 319/319 [00:00<00:00, 622.56it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Degree of memorization after memorizing: %f  tensor(0.9420)\n",
      "forgetting the same text iteration = 0\n",
      "forgetting the same text iteration = 1\n",
      "forgetting the same text iteration = 2\n",
      "forgetting the same text iteration = 3\n",
      "forgetting the same text iteration = 4\n",
      "forgetting the same text iteration = 5\n",
      "forgetting the same text iteration = 6\n",
      "forgetting the same text iteration = 7\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 319/319 [00:00<00:00, 494.24it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Degree of memorization after forgetting: %f  tensor(0.0815)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "model = MeMo(inner_dim=d, \n",
    "             num_of_heads=h, \n",
    "             num_of_layers=l, \n",
    "             chunk_length=max_length, \n",
    "             num_embeddings=tokenizer.vocab_size, \n",
    "             padding_idx=tokenizer.pad_token_id, \n",
    "             device=device)\n",
    "print(\"CMM pre learning\")\n",
    "display(model.layers[0].CMM.weight)\n",
    "\n",
    "\n",
    "bs = 8\n",
    "for b in range(bs):\n",
    "    print(f\"memorizing the same text iteration = {b}\")\n",
    "    memo_input = tokenizer.get_text_batch_encoding(my_first_text)\n",
    "    model.memorize_text(memo_input)\n",
    "\n",
    "Prj = model.layers[0].Prj.weight.detach().cpu()\n",
    "CMM = model.layers[0].CMM.weight.detach().cpu()\n",
    "\n",
    "display(Prj.T @ Prj)\n",
    "display(CMM)\n",
    "\n",
    "e = Evaluation()\n",
    "out = e.check_pretokenized(model, tokenizer, memo_input['input_ids'])\n",
    "print(\"Degree of memorization after memorizing: %f \", out)\n",
    "\n",
    "for b in range(bs):\n",
    "    print(f\"forgetting the same text iteration = {b}\")\n",
    "    memo_input = tokenizer.get_text_batch_encoding(my_first_text)\n",
    "    model.forget_text(memo_input)\n",
    "\n",
    "out = e.check_pretokenized(model, tokenizer, memo_input['input_ids'])\n",
    "print(\"Degree of memorization after forgetting: %f \", out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41935a13-0b0f-4b93-b129-0a1e88241fcd",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d237416-54ea-4b21-8d19-84599ff9e34a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ad99bb2-af43-4e42-a447-6871f4ee47f1",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "346e794d-b1c9-438e-b66f-14fd17c1ba6e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c4d5e5c-a117-4378-970a-d42b76a0ed12",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43b9e371-c48a-479a-8cec-fc87c9aaacdb",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}