Update code

main.ipynb

@@ -2,10 +2,25 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "id": "d1daab37",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "6be175a38e1843ab83831cb3a9a06296",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Textarea(value=\"Solomonoff's theory of inductive inference proposes that all problems of logical induction can…"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
    "source": [
     "import ipywidgets\n",
     "\n",
@@ -20,14 +35,688 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 6,
    "id": "09b3c097",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Optimizing: Hey this is a test text of me, Aldan, writing some random text in this text box. Will this work? Perhaps!\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  0%|          | 2/5000 [00:00<13:45,  6.05it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 0, Loss: 5.8886637687683105, L2 norm: 0.5542519688606262, avg rank: 655.2222290039062\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  0%|          | 12/5000 [00:01<10:57,  7.58it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 10, Loss: 5.159086227416992, L2 norm: 3.46441388130188, avg rank: 439.6666564941406\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  0%|          | 22/5000 [00:03<12:31,  6.62it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 20, Loss: 4.893388748168945, L2 norm: 5.287822246551514, avg rank: 380.629638671875\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  1%|          | 32/5000 [00:04<10:37,  7.79it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 30, Loss: 4.576204299926758, L2 norm: 7.122143745422363, avg rank: 313.0740661621094\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  1%|          | 42/5000 [00:05<10:42,  7.71it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 40, Loss: 4.330563545227051, L2 norm: 8.640031814575195, avg rank: 189.8518524169922\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  1%|          | 52/5000 [00:07<13:39,  6.04it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 50, Loss: 4.107393264770508, L2 norm: 10.037625312805176, avg rank: 130.6666717529297\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  1%|          | 62/5000 [00:08<11:33,  7.12it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 60, Loss: 3.8990049362182617, L2 norm: 11.52373218536377, avg rank: 105.9259262084961\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  1%|▏         | 72/5000 [00:10<19:46,  4.15it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 70, Loss: 3.689441442489624, L2 norm: 12.885846138000488, avg rank: 96.77777862548828\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  2%|▏         | 82/5000 [00:13<14:27,  5.67it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 80, Loss: 3.664461374282837, L2 norm: 14.095556259155273, avg rank: 82.48148345947266\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  2%|▏         | 92/5000 [00:14<12:06,  6.76it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 90, Loss: 3.359441041946411, L2 norm: 15.264603614807129, avg rank: 76.62963104248047\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  2%|▏         | 102/5000 [00:15<11:25,  7.14it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 100, Loss: 3.2641189098358154, L2 norm: 16.456912994384766, avg rank: 59.55555725097656\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  2%|▏         | 111/5000 [00:17<10:21,  7.87it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 110, Loss: 3.4184069633483887, L2 norm: 17.494909286499023, avg rank: 45.22222137451172\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  2%|▏         | 122/5000 [00:19<14:21,  5.66it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 120, Loss: 3.057454824447632, L2 norm: 18.361045837402344, avg rank: 47.592594146728516\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  3%|▎         | 132/5000 [00:20<11:23,  7.12it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 130, Loss: 2.752321243286133, L2 norm: 19.076860427856445, avg rank: 37.185184478759766\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  3%|▎         | 142/5000 [00:21<10:12,  7.93it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 140, Loss: 2.492600679397583, L2 norm: 19.873262405395508, avg rank: 29.0\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  3%|▎         | 152/5000 [00:23<12:10,  6.64it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 150, Loss: 2.2693941593170166, L2 norm: 20.689470291137695, avg rank: 25.22222137451172\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  3%|▎         | 162/5000 [00:24<11:26,  7.05it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 160, Loss: 2.0094830989837646, L2 norm: 21.427452087402344, avg rank: 20.185184478759766\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  3%|▎         | 172/5000 [00:26<11:22,  7.08it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 170, Loss: 1.8882372379302979, L2 norm: 22.13208770751953, avg rank: 23.074073791503906\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  4%|▎         | 182/5000 [00:27<09:56,  8.08it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 180, Loss: 1.8304041624069214, L2 norm: 22.583101272583008, avg rank: 28.33333396911621\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  4%|▍         | 192/5000 [00:28<10:58,  7.30it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 190, Loss: 1.5684080123901367, L2 norm: 23.05557632446289, avg rank: 22.037036895751953\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  4%|▍         | 202/5000 [00:30<09:58,  8.02it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 200, Loss: 1.3705590963363647, L2 norm: 23.487506866455078, avg rank: 18.925926208496094\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  4%|▍         | 212/5000 [00:31<09:58,  8.00it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 210, Loss: 1.2061578035354614, L2 norm: 23.888507843017578, avg rank: 15.481481552124023\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  4%|▍         | 222/5000 [00:32<10:36,  7.50it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 220, Loss: 1.0748673677444458, L2 norm: 24.25291633605957, avg rank: 13.333333015441895\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  5%|▍         | 232/5000 [00:34<11:45,  6.76it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 230, Loss: 0.9695132374763489, L2 norm: 24.588470458984375, avg rank: 12.592592239379883\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  5%|▍         | 242/5000 [00:35<11:47,  6.73it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 240, Loss: 0.8693955540657043, L2 norm: 24.909320831298828, avg rank: 11.518518447875977\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  5%|▌         | 252/5000 [00:36<09:53,  7.99it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 250, Loss: 0.7695979475975037, L2 norm: 25.205760955810547, avg rank: 10.037036895751953\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  5%|▌         | 262/5000 [00:38<10:06,  7.81it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 260, Loss: 0.6803638935089111, L2 norm: 25.493518829345703, avg rank: 8.666666984558105\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  5%|▌         | 272/5000 [00:39<11:42,  6.73it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 270, Loss: 0.6074316501617432, L2 norm: 25.7545108795166, avg rank: 7.407407283782959\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  6%|▌         | 282/5000 [00:40<10:20,  7.60it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 280, Loss: 0.5503782629966736, L2 norm: 25.987356185913086, avg rank: 6.666666507720947\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  6%|▌         | 292/5000 [00:42<10:42,  7.33it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 290, Loss: 0.5049920678138733, L2 norm: 26.197214126586914, avg rank: 6.111111164093018\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  6%|▌         | 302/5000 [00:43<11:42,  6.68it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 300, Loss: 0.46782854199409485, L2 norm: 26.38979721069336, avg rank: 5.407407283782959\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  6%|▌         | 312/5000 [00:45<10:34,  7.39it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 310, Loss: 0.4365224838256836, L2 norm: 26.568010330200195, avg rank: 4.666666507720947\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  6%|▋         | 322/5000 [00:46<10:26,  7.46it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 320, Loss: 0.4172615706920624, L2 norm: 26.73436737060547, avg rank: 4.296296119689941\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  7%|▋         | 332/5000 [00:48<11:34,  6.73it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 330, Loss: 0.44890207052230835, L2 norm: 26.890453338623047, avg rank: 4.777777671813965\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  7%|▋         | 342/5000 [00:49<09:44,  7.97it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 340, Loss: 1.6345257759094238, L2 norm: 27.1015625, avg rank: 10.0\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  7%|▋         | 352/5000 [00:50<10:12,  7.59it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 350, Loss: 1.079402208328247, L2 norm: 27.411144256591797, avg rank: 6.592592716217041\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  7%|▋         | 362/5000 [00:52<12:10,  6.35it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 360, Loss: 0.8217418193817139, L2 norm: 27.712562561035156, avg rank: 4.703703880310059\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  7%|▋         | 372/5000 [00:53<11:04,  6.97it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 370, Loss: 0.6470327973365784, L2 norm: 28.024415969848633, avg rank: 1.8888888359069824\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  8%|▊         | 382/5000 [00:55<10:21,  7.43it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 380, Loss: 1.0997235774993896, L2 norm: 28.2260684967041, avg rank: 7.185184955596924\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  8%|▊         | 392/5000 [00:56<10:58,  7.00it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 390, Loss: 0.64601731300354, L2 norm: 28.461565017700195, avg rank: 2.9259259700775146\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  8%|▊         | 402/5000 [00:58<09:58,  7.68it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 400, Loss: 0.482523649930954, L2 norm: 28.71732521057129, avg rank: 1.5185185670852661\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  8%|▊         | 412/5000 [00:59<09:54,  7.71it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 410, Loss: 0.38174617290496826, L2 norm: 28.943510055541992, avg rank: 1.2592592239379883\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  8%|▊         | 422/5000 [01:01<12:41,  6.01it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 420, Loss: 0.29940587282180786, L2 norm: 29.132537841796875, avg rank: 1.1111111640930176\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  9%|▊         | 432/5000 [01:02<10:03,  7.57it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 430, Loss: 0.23488281667232513, L2 norm: 29.293710708618164, avg rank: 1.0740740299224854\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  9%|▉         | 442/5000 [01:03<09:35,  7.91it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 440, Loss: 0.19830304384231567, L2 norm: 29.436723709106445, avg rank: 1.0370370149612427\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  9%|▉         | 452/5000 [01:05<10:41,  7.09it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 450, Loss: 0.17161428928375244, L2 norm: 29.54309844970703, avg rank: 1.0370370149612427\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  9%|▉         | 454/5000 [01:05<10:54,  6.94it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Perfect ranks achieved at step 454, stopping optimization.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    }
+   ],
    "source": [
     "import transformers\n",
     "import torch\n",
     "import tqdm\n",
+    "torch.manual_seed(42)  # Set a fixed seed for reproducibility\n",
     "\n",
     "print(f\"Optimizing: {text_widget.value}\")\n",
     "\n",
@@ -79,7 +768,7 @@
     "\n",
     "# Define the number of optimization steps\n",
     "n_steps = 5000\n",
-    "patience = 50\n",
+    "patience = 100\n",
     "patience_counter = 0\n",
     "epsilon = 1e-1\n",
     "\n",
@@ -112,8 +801,8 @@
     "\n",
     "    # Compute the ranks of the input IDs, i.e. how many tokens would have been more likely than the correct one (the label, the input IDs)\n",
     "    \n",
-    "    count_higher_logits = (logits[:, ss_prompt.size(-2) - 1 :-1] > logits[:, -1:]).float()\n",
-    "    ranks = count_higher_logits.sum(dim=-1)\n",
+    "    # Calculate the ranks by summing the probabilities of tokens with higher logits than the correct token\n",
+    "    ranks = torch.sum(probs > probs.gather(2, tokens[\"input_ids\"].unsqueeze(-1)), dim=-1) + 1\n",
     "\n",
     "    # Compute the loss\n",
     "    loss = loss_fn(\n",
@@ -134,7 +823,7 @@
     "            for optimized, original in zip(ss_prompt, original_ss_prompt)\n",
     "        )\n",
     "        print(\n",
-    "            f\"Step {step}, Loss: {loss.item()}, L2 norm: {l2_norm.item()}, avg rank: {ranks.mean().item()}\"\n",
+    "            f\"Step {step}, Loss: {loss.item()}, L2 norm: {l2_norm.item()}, avg rank: {ranks.float().mean().item()}\"\n",
     "        )\n",
     "\n",
     "    # Early stopping with patience\n",
@@ -147,12 +836,17 @@
     "\n",
     "    if patience_counter >= patience:\n",
     "        print(f\"Early stopping at step {step} with best loss {best_loss}\")\n",
+    "        break\n",
+    "\n",
+    "    # If the ranks are perfect (all 1), stop\n",
+    "    if torch.all(ranks == 1):\n",
+    "        print(f\"Perfect ranks achieved at step {step}, stopping optimization.\")\n",
     "        break"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 7,
    "id": "cc9a6a2f",
    "metadata": {},
    "outputs": [],
@@ -163,22 +857,190 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 8,
    "id": "3186747d",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 150/150 [00:10<00:00, 14.83it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Reference: Hey this is a test text of me, Aldan, writing some random text in this text box. Will this work? Perhaps!\n",
+      "Generated: Hey this is a test text of me, Aldan, writing some random text in this text box. Will this work? Perhaps! I'll have to test this out. Maybe. I'll have to test this out some more. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe. Maybe\n",
+      "'Hey':\tRank 0.00, probability: 71.02%, Reference: 'Hey'\n",
+      "' this':\tRank 0.00, probability: 87.04%, Reference: ' this'\n",
+      "' is':\tRank 0.00, probability: 99.65%, Reference: ' is'\n",
+      "' a':\tRank 0.00, probability: 96.23%, Reference: ' a'\n",
+      "' test':\tRank 0.00, probability: 90.31%, Reference: ' test'\n",
+      "' text':\tRank 0.00, probability: 95.40%, Reference: ' text'\n",
+      "' of':\tRank 0.00, probability: 62.63%, Reference: ' of'\n",
+      "' me':\tRank 0.00, probability: 93.07%, Reference: ' me'\n",
+      "',':\tRank 0.00, probability: 95.56%, Reference: ','\n",
+      "' Ald':\tRank 0.00, probability: 71.15%, Reference: ' Ald'\n",
+      "'an':\tRank 0.00, probability: 84.45%, Reference: 'an'\n",
+      "',':\tRank 0.00, probability: 97.80%, Reference: ','\n",
+      "' writing':\tRank 0.00, probability: 95.00%, Reference: ' writing'\n",
+      "' some':\tRank 0.00, probability: 88.69%, Reference: ' some'\n",
+      "' random':\tRank 0.00, probability: 97.07%, Reference: ' random'\n",
+      "' text':\tRank 0.00, probability: 89.14%, Reference: ' text'\n",
+      "' in':\tRank 0.00, probability: 90.19%, Reference: ' in'\n",
+      "' this':\tRank 0.00, probability: 98.12%, Reference: ' this'\n",
+      "' text':\tRank 0.00, probability: 96.47%, Reference: ' text'\n",
+      "' box':\tRank 0.00, probability: 95.55%, Reference: ' box'\n",
+      "'.':\tRank 0.00, probability: 97.38%, Reference: '.'\n",
+      "' Will':\tRank 0.00, probability: 90.65%, Reference: ' Will'\n",
+      "' this':\tRank 0.00, probability: 93.50%, Reference: ' this'\n",
+      "' work':\tRank 0.00, probability: 49.78%, Reference: ' work'\n",
+      "'?':\tRank 0.00, probability: 96.33%, Reference: '?'\n",
+      "' Perhaps':\tRank 0.00, probability: 76.97%, Reference: ' Perhaps'\n",
+      "'!':\tRank 0.00, probability: 42.45%, Reference: '!'\n",
+      "' I':\tRank 0.00, probability: 11.67%, Reference: 'N/A'\n",
+      "''ll':\tRank 0.00, probability: 10.19%, Reference: 'N/A'\n",
+      "' have':\tRank 0.00, probability: 53.08%, Reference: 'N/A'\n",
+      "' to':\tRank 0.00, probability: 77.45%, Reference: 'N/A'\n",
+      "' test':\tRank 0.00, probability: 8.06%, Reference: 'N/A'\n",
+      "' this':\tRank 0.00, probability: 40.56%, Reference: 'N/A'\n",
+      "' out':\tRank 0.00, probability: 51.98%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 26.31%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 27.11%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 22.16%, Reference: 'N/A'\n",
+      "' I':\tRank 0.00, probability: 20.00%, Reference: 'N/A'\n",
+      "''ll':\tRank 0.00, probability: 21.35%, Reference: 'N/A'\n",
+      "' have':\tRank 0.00, probability: 59.45%, Reference: 'N/A'\n",
+      "' to':\tRank 0.00, probability: 74.45%, Reference: 'N/A'\n",
+      "' test':\tRank 0.00, probability: 17.91%, Reference: 'N/A'\n",
+      "' this':\tRank 0.00, probability: 70.72%, Reference: 'N/A'\n",
+      "' out':\tRank 0.00, probability: 87.33%, Reference: 'N/A'\n",
+      "' some':\tRank 0.00, probability: 28.23%, Reference: 'N/A'\n",
+      "' more':\tRank 0.00, probability: 26.20%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 22.73%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 36.98%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 54.91%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 64.39%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 65.85%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 65.80%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 51.26%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 60.87%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 35.56%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 48.26%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 28.23%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 43.46%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 23.53%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 42.47%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 21.56%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 44.99%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 22.61%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 48.24%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 19.47%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 51.23%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 19.13%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 54.27%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 20.28%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 56.09%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 23.08%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 60.23%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 25.86%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 63.49%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 28.66%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 65.52%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 29.81%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 67.58%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 32.28%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 69.40%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 33.50%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 70.15%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 34.01%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 71.20%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 36.83%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 71.29%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 38.54%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 72.09%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 40.10%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 71.66%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 41.78%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 72.04%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 42.55%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 71.18%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 44.33%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 70.62%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 44.84%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 70.70%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 44.89%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 71.80%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 46.26%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 71.88%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 46.11%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 72.24%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 44.98%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 73.18%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 45.08%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 75.74%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 45.33%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 76.68%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 46.90%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 78.19%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 46.31%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 78.19%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 48.19%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 79.38%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 48.56%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 80.29%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 49.49%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 79.51%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 51.13%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 81.70%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 52.29%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 80.87%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 52.84%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 81.61%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 54.15%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 81.67%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 55.17%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 82.23%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 55.47%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 82.99%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 57.04%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 83.32%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 58.06%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 84.47%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 59.85%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 83.82%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 58.62%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 85.14%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 60.26%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 85.67%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 62.56%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 86.05%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 64.28%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 86.93%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 67.84%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 87.52%, Reference: 'N/A'\n",
+      "'.':\tRank 0.00, probability: 73.01%, Reference: 'N/A'\n",
+      "' Maybe':\tRank 0.00, probability: 88.17%, Reference: 'N/A'\n"
+     ]
+    }
+   ],
    "source": [
     "import os\n",
     "import transformers\n",
     "import torch\n",
     "import tqdm\n",
     "\n",
-    "# tokenizer: transformers.PreTrainedTokenizer = (\n",
-    "#     transformers.AutoTokenizer.from_pretrained(\"openai-community/gpt2\")\n",
-    "# )\n",
-    "# model: transformers.PreTrainedModel = transformers.AutoModelForCausalLM.from_pretrained(\n",
-    "#     \"openai-community/gpt2\"\n",
-    "# )\n",
+    "if \"tokenizer\" not in locals():\n",
+    "    tokenizer: transformers.PreTrainedTokenizer = (\n",
+    "        transformers.AutoTokenizer.from_pretrained(\"openai-community/gpt2\")\n",
+    "    )\n",
+    "if \"model\" not in locals():\n",
+    "    model: transformers.PreTrainedModel = transformers.AutoModelForCausalLM.from_pretrained(\n",
+    "        \"openai-community/gpt2\"\n",
+    "    )\n",
     "\n",
     "# Load the best soft prompt from file\n",
     "best_ss_prompt = torch.load(\"best_ss_prompt.pt\")\n",