Gradio demo

README.md

@@ -1,4 +1,14 @@
-# ss-prompts
-Simon Says Prompts
+---
+title: "Simon Says Prompts"
+emoji: "🗣️"
+colorFrom: "blue"
+colorTo: "pink"
+sdk: "gradio"
+sdk_version: "5.26.0"
+app_file: "demo.py"
+pinned: false
+---
+
+# Simon Says Prompts
 
 This repo is a demo of how to generate an embedding that gets fed into a given LLM and is optimized to make the LLM output a specific text under a greedy decoding scheme.

demo.py

@@ -0,0 +1,318 @@
+import gradio as gr
+import torch
+import transformers
+import time
+from safetensors.torch import save_file, load_file
+import tempfile
+from io import BytesIO
+import logging
+
+# Load the tokenizer and model
+tokenizer = transformers.AutoTokenizer.from_pretrained("openai-community/gpt2")
+model = transformers.AutoModelForCausalLM.from_pretrained("openai-community/gpt2")
+
+
+# Update the optimization function in demo.py to align with the notebook
+
+
+def optimize_simon_says_prompt(
+    input_text: str,
+    number_of_simon_says_tokens: int,
+    n_steps: int,
+    lr: float,
+    progress=gr.Progress(track_tqdm=False),  # Gradio progress tracking
+) -> tuple[str, torch.Tensor]:
+    """
+    Optimize a Simon Says prompt based on the input text and display the optimization process.
+
+    Parameters:
+        input_text (str): The input text provided by the user.
+        number_of_simon_says_tokens (int): Number of Simon Says tokens to optimize.
+        n_steps (int): Number of optimization steps.
+        lr (float): Learning rate for the optimization process.
+        progress (gr.Progress): Gradio progress tracking.
+
+    Returns:
+        The optimized Simon Says prompt
+    """
+    # Tokenize the input text
+    tokens = tokenizer(
+        input_text,
+        return_tensors="pt",
+        padding=False,
+        truncation=True,
+        add_special_tokens=True,
+    )
+    embeddings = model.transformer.wte(tokens["input_ids"]).detach()
+
+    # Initialize a random Simon Says prompt
+    simon_says_prompt = torch.randn(
+        1, number_of_simon_says_tokens, model.config.n_embd, requires_grad=True
+    )
+    optimizer = torch.optim.Adam([simon_says_prompt], lr=lr)
+    loss_fn = torch.nn.CrossEntropyLoss()
+
+    best_loss: float = float("inf")
+    best_simon_says_prompt: torch.Tensor = None
+
+    progress(0, desc="Starting optimization...")
+    time.sleep(1)
+
+    for step in range(n_steps):
+        optimizer.zero_grad()
+        expanded_prompt = torch.cat([simon_says_prompt, embeddings], dim=1)
+        logits = model(inputs_embeds=expanded_prompt).logits
+        probs = torch.softmax(logits[:, simon_says_prompt.size(-2) - 1 : -1], dim=-1)
+        ranks = (
+            torch.sum(
+                probs > probs.gather(2, tokens["input_ids"].unsqueeze(-1)), dim=-1
+            )
+            + 1
+        )
+        loss = loss_fn(
+            logits[:, simon_says_prompt.size(-2) - 1 : -1].reshape(-1, logits.size(-1)),
+            tokens["input_ids"].reshape(-1),
+        )
+        loss.backward()
+        optimizer.step()
+
+        avg_rank = ranks.float().mean().item()
+        progress(
+            step / n_steps,
+            desc=f"Step {step}, Loss: {loss.item():.4f}, Avg Rank: {avg_rank:.2f}, Max Rank: {ranks.max().item()}",
+        )
+
+        logging.info(
+            f"Step {step}, Loss: {loss.item():.4f}, Avg Rank: {avg_rank:.2f}, Max Rank: {ranks.max().item()}"
+        )
+
+        if loss.item() < best_loss:
+            best_loss = loss.item()
+            best_simon_says_prompt = simon_says_prompt.detach().clone()
+
+        # If all ranks are 1, stop the optimization (perfect prediction)
+        if torch.all(ranks == 1):
+            break
+
+    return best_simon_says_prompt
+
+
+# Modify the download_tensor function to save the tensor as a safetensors file
+
+
+def download_tensor(tensor):
+    """
+    Save a tensor to a safetensors file for download.
+
+    Parameters:
+        tensor (torch.Tensor): The tensor to be saved.
+
+    Returns:
+        str: The file path of the saved tensor.
+    """
+    file_path = "optimized_tensor.safetensors"
+    save_file({"optimized_tensor": tensor}, file_path)
+    return file_path
+
+
+def upload_tensor(file):
+    """
+    Load a tensor from an uploaded safetensors file.
+
+    Parameters:
+        file (bytes): The uploaded file containing the safetensors data.
+
+    Returns:
+        torch.Tensor: The loaded tensor.
+
+    Raises:
+        ValueError: If the safetensors file is invalid or the header is too large.
+    """
+    if isinstance(file, bytes):
+        file = BytesIO(file)  # Wrap bytes in a BytesIO object
+
+    with tempfile.NamedTemporaryFile(delete=True) as temp_file:
+        temp_file.write(
+            file.read()
+        )  # Directly write the BytesIO content to the temporary file
+        temp_file.flush()
+
+        try:
+            tensor_data = load_file(temp_file.name)
+        except Exception as e:
+            raise ValueError(f"Failed to load safetensors file: {e}")
+
+        if "optimized_tensor" not in tensor_data:
+            raise ValueError(
+                "The safetensors file does not contain the expected 'optimized_tensor' key."
+            )
+
+        return tensor_data["optimized_tensor"]
+
+
+def greedy_decode_with_ss_prompt(
+    ss_prompt: torch.Tensor, progress=gr.Progress()
+) -> str:
+    """
+    Perform greedy decoding using an uploaded optimized tensor and input text.
+
+    Parameters:
+        ss_prompt (torch.Tensor): The uploaded optimized tensor.
+        progress (gr.Progress): Gradio progress tracking.
+
+    Returns:
+        str: The generated text.
+    """
+    generated_tokens = []
+    all_logits = []
+
+    progress(0, desc="Starting greedy decoding...")
+
+    with torch.no_grad():
+        for i in progress.tqdm(range(150), desc="Decoding..."):
+            if len(generated_tokens) == 0:
+                expanded_prompt = ss_prompt
+            else:
+                expanded_prompt = torch.cat(
+                    [
+                        ss_prompt,
+                        model.transformer.wte(
+                            torch.tensor(generated_tokens).unsqueeze(0)
+                        ).detach(),
+                    ],
+                    dim=1,
+                )
+
+            logits = model(inputs_embeds=expanded_prompt).logits
+            next_token_logits = logits[0, -1, :]
+            next_token = next_token_logits.argmax().item()
+
+            logging.info(
+                f"Step {i}, Next Token: {next_token}, Logit: {next_token_logits[next_token].item()}"
+            )
+
+            generated_tokens.append(next_token)
+            all_logits.append(next_token_logits)
+
+            if next_token == tokenizer.eos_token_id:
+                break
+
+    generated_tokens = torch.tensor(generated_tokens)
+    generated_text = tokenizer.decode(generated_tokens, skip_special_tokens=True)
+
+    return generated_text
+
+
+def process_and_generate(
+    input_text: str,
+    number_of_simon_says_tokens: int,
+    n_steps: int,
+    lr: float,
+) -> tuple[str, str]:
+    """
+    Optimize the Simon Says prompt, display the optimization process, and generate text based on the input text.
+
+    Parameters:
+        input_text (str): The input text provided by the user.
+        number_of_simon_says_tokens (int): Number of Simon Says tokens to optimize.
+        n_steps (int): Number of optimization steps.
+        lr (float): Learning rate for the optimization process.
+
+    Returns:
+        tuple: The optimized Simon Says prompt and the greedy-decoded text.
+    """
+    optimized_prompt = optimize_simon_says_prompt(
+        input_text=input_text,
+        number_of_simon_says_tokens=number_of_simon_says_tokens,
+        n_steps=n_steps,
+        lr=lr,
+    )
+
+    # Generate text using the optimized prompt
+    generated_text: str = greedy_decode_with_ss_prompt(optimized_prompt)
+
+    return (
+        generated_text,
+        download_tensor(optimized_prompt),
+    )  # Return the optimized tensor for download
+
+
+def process_with_uploaded_tensor(
+    input_text: str, uploaded_tensor: torch.Tensor
+) -> tuple[str, str]:
+    """
+    Process the uploaded tensor and generate text based on the input text.
+
+    Parameters:
+        input_text (str): The input text provided by the user.
+        uploaded_tensor (torch.Tensor): The uploaded optimized tensor.
+
+    Returns:
+        tuple: The generated text and the file path of the uploaded tensor.
+    """
+    generated_text = greedy_decode_with_ss_prompt(uploaded_tensor)
+    return generated_text, None
+
+
+theme = gr.themes.Soft(
+    primary_hue="fuchsia",
+    secondary_hue="cyan",
+    neutral_hue="gray",
+    radius_size="none",
+    font=[
+        gr.themes.GoogleFont("IBM Plex Sans"),
+        "ui-sans-serif",
+        "system-ui",
+        "sans-serif",
+    ],
+    font_mono=[
+        gr.themes.GoogleFont("IBM Plex Mono"),
+        "ui-monospace",
+        "Consolas",
+        "monospace",
+    ],
+)
+
+# Update the Gradio interface to include configurable parameters
+demo = gr.Interface(
+    theme=theme,
+    title="Simon Says Prompt Optimization and Text Generation",
+    fn=lambda input_text, number_of_simon_says_tokens, n_steps, lr, uploaded_file: (
+        process_with_uploaded_tensor(input_text, upload_tensor(uploaded_file))
+        if uploaded_file
+        else process_and_generate(
+            input_text, number_of_simon_says_tokens, n_steps, lr
+        )
+    ),
+    inputs=[
+        gr.Textbox(
+            lines=5,
+            placeholder="Enter your text here...",
+            label="Input Text",
+            value="Hello world! I'm Aldan, happy to be here.",
+        ),
+        gr.Slider(
+            minimum=1, maximum=10, step=1, value=4, label="Number of Simon Says Tokens"
+        ),
+        gr.Slider(
+            minimum=100,
+            maximum=10000,
+            step=100,
+            value=5000,
+            label="Number of Optimization Steps",
+        ),
+        gr.Slider(
+            minimum=1e-5, maximum=1e-1, step=1e-5, value=1e-2, label="Learning Rate"
+        ),
+        gr.File(label="Upload Optimized Tensor (Optional)", type="binary"),
+    ],
+    outputs=[
+        gr.Textbox(label="Generated Text"),
+        gr.File(label="Download Optimized Tensor", type="filepath"),
+    ],
+    description="This demo optimizes a Simon Says prompt based on your input text, displays the optimization process, and generates text using the optimized prompt. Optionally, you can upload a pre-optimized tensor for inference.",
+)
+
+# Ensure the Gradio interface is correctly launched
+if __name__ == "__main__":
+    demo.launch(debug=True, show_error=True)

main.ipynb

@@ -2,25 +2,10 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "id": "d1daab37",
    "metadata": {},
-   "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"
-    }
-   ],
+   "outputs": [],
    "source": [
     "import ipywidgets\n",
     "\n",
@@ -35,683 +20,10 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "id": "09b3c097",
    "metadata": {},
-   "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"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "import transformers\n",
     "import torch\n",
@@ -738,30 +50,28 @@
     "\n",
     "embeddings = model.transformer.wte(tokens[\"input_ids\"]).detach()\n",
     "\n",
-    "# We'll use a soft prompt - a special token and its hidden states for the first token that we'll use to condition the model. We'll optimize the hidden states of this token to maximize the likelihood of the text that follows it.\n",
-    "# Generate a soft prompt by creating random hidden states for the first token\n",
+    "# We'll use a Simon Says prompt - a special token and its hidden states for the first token that we'll use to condition the model. We'll optimize the hidden states of this token to maximize the likelihood of the text that follows it.\n",
+    "# Generate a Simon Says prompt by creating random hidden states for the first token\n",
     "# We'll optimize these hidden states to maximize the likelihood of the text that follows it\n",
     "# past_key_values (Tuple[Tuple[torch.Tensor]] of length config.n_layers) — Contains precomputed hidden-states (key and values in the attention blocks) as computed by the model (see past_key_values output below). Can be used to speed up sequential decoding. The input_ids which have their past given to this model should not be passed as input_ids as they have already been computed.\n",
     "# Shape: past_key_values (Tuple[Tuple[torch.Tensor]] of length config.n_layers)\n",
     "# with each tuple having 2 tensors of shape (batch_size, num_heads, sequence_length, embed_size_per_head))\n",
     "# The precision of the hidden states should be fp16\n",
-    "ss_prompt = (\n",
+    "simon_says_prompt = (\n",
     "    # One tensor of shape (1, 1, embed_size_per_head)\n",
     "    torch.randn(\n",
     "        1,\n",
     "        4,\n",
     "        model.config.n_embd,\n",
     "        requires_grad=True,\n",
-    "        # dtype=torch.float16,\n",
     "    )\n",
-    "    # for _ in range(model.config.n_layer)\n",
     ")\n",
     "\n",
-    "# Copy the soft prompts since we'll optimize them and they'll change\n",
-    "original_ss_prompt = tuple(key.detach().clone() for key in ss_prompt)\n",
+    "# Copy the Simon Says prompts since we'll optimize them and they'll change\n",
+    "original_simon_says_prompt = tuple(key.detach().clone() for key in simon_says_prompt)\n",
     "\n",
     "# Define the optimizer\n",
-    "optimizer = torch.optim.Adam([ss_prompt], lr=1e-2)\n",
+    "optimizer = torch.optim.Adam([simon_says_prompt], lr=1e-2)\n",
     "\n",
     "# Define the loss function\n",
     "loss_fn = torch.nn.CrossEntropyLoss()\n",
@@ -775,8 +85,8 @@
     "# Freeze the model parameters\n",
     "model.eval()\n",
     "\n",
-    "# Check that the soft prompt is optimizable (requires_grad=True)\n",
-    "for key in ss_prompt:\n",
+    "# Check that the Simon Says prompt is optimizable (requires_grad=True)\n",
+    "for key in simon_says_prompt:\n",
     "    assert key.requires_grad\n",
     "\n",
     "# Disable gradient computation for the model\n",
@@ -784,20 +94,20 @@
     "    param.requires_grad = False\n",
     "\n",
     "best_loss = float(\"inf\")\n",
-    "best_ss_prompt = None\n",
+    "best_simon_says_prompt = None\n",
     "\n",
-    "# Optimize the soft prompt\n",
+    "# Optimize the Simon Says prompt\n",
     "for step in tqdm.tqdm(range(n_steps)):\n",
     "    # Zero the gradients\n",
     "    optimizer.zero_grad()\n",
     "\n",
-    "    # Add the optimizable SS prompt to the embeddings\n",
-    "    expanded_prompt = torch.cat([ss_prompt, embeddings], dim=1)\n",
+    "    # Add the optimizable Simon Says prompt to the embeddings\n",
+    "    expanded_prompt = torch.cat([simon_says_prompt, embeddings], dim=1)\n",
     "\n",
     "    # Generate the logits for the text\n",
     "    logits: torch.Tensor = model(inputs_embeds=expanded_prompt).logits\n",
     "\n",
-    "    probs = torch.softmax(logits[:, ss_prompt.size(-2) - 1 :-1], dim=-1)\n",
+    "    probs = torch.softmax(logits[:, simon_says_prompt.size(-2) - 1 :-1], dim=-1)\n",
     "\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",
@@ -806,21 +116,21 @@
     "\n",
     "    # Compute the loss\n",
     "    loss = loss_fn(\n",
-    "        logits[:, ss_prompt.size(-2) - 1 :-1].reshape(-1, logits.size(-1)),\n",
+    "        logits[:, simon_says_prompt.size(-2) - 1 :-1].reshape(-1, logits.size(-1)),\n",
     "        tokens[\"input_ids\"].reshape(-1),\n",
     "    )\n",
     "\n",
     "    # Backpropagate the gradients\n",
     "    loss.backward()\n",
     "\n",
-    "    # Optimize the soft prompt\n",
+    "    # Optimize the Simon Says prompt\n",
     "    optimizer.step()\n",
     "\n",
     "    if step % 10 == 0:\n",
-    "        # Get the L2 norm of the difference between the original and optimized soft prompts\n",
+    "        # Get the L2 norm of the difference between the original and optimized Simon Says prompts\n",
     "        l2_norm = sum(\n",
     "            torch.norm(optimized - original, p=2)\n",
-    "            for optimized, original in zip(ss_prompt, original_ss_prompt)\n",
+    "            for optimized, original in zip(simon_says_prompt, original_simon_says_prompt)\n",
     "        )\n",
     "        print(\n",
     "            f\"Step {step}, Loss: {loss.item()}, L2 norm: {l2_norm.item()}, avg rank: {ranks.float().mean().item()}\"\n",
@@ -829,7 +139,7 @@
     "    # Early stopping with patience\n",
     "    if loss.item() < best_loss and loss.item() > epsilon:\n",
     "        best_loss = loss.item()\n",
-    "        best_ss_prompt = ss_prompt.detach().clone()\n",
+    "        best_simon_says_prompt = simon_says_prompt.detach().clone()\n",
     "        patience_counter = 0\n",
     "    else:\n",
     "        patience_counter += 1\n",
@@ -846,187 +156,21 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": null,
    "id": "cc9a6a2f",
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Save the best soft prompt\n",
-    "torch.save(best_ss_prompt, \"best_ss_prompt.pt\")"
+    "# Save the best Simon Says prompt\n",
+    "torch.save(best_simon_says_prompt, \"best_simon_says_prompt.pt\")"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": null,
    "id": "3186747d",
    "metadata": {},
-   "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"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "import os\n",
     "import transformers\n",
@@ -1042,24 +186,24 @@
     "        \"openai-community/gpt2\"\n",
     "    )\n",
     "\n",
-    "# Load the best soft prompt from file\n",
-    "best_ss_prompt = torch.load(\"best_ss_prompt.pt\")\n",
+    "# Load the best Simon Says prompt from file\n",
+    "best_simon_says_prompt = torch.load(\"best_simon_says_prompt.pt\")\n",
     "\n",
     "# Do a greedy decoding manually\n",
-    "# Pass the optimized soft prompt to the model\n",
+    "# Pass the optimized Simon Says prompt to the model\n",
     "# We can't use .generate() since we need to pass the inputs_embeds\n",
     "all_logits = []\n",
     "generated_tokens = []\n",
     "with torch.no_grad():\n",
     "    for i in tqdm.tqdm(range(150)):\n",
     "        # Generate the logits for the next token using what we've generated so far\n",
-    "        # If there are no generated tokens yet, just take the soft prompt\n",
+    "        # If there are no generated tokens yet, just take the Simon Says prompt\n",
     "        if len(generated_tokens) == 0:\n",
-    "            expanded_prompt = best_ss_prompt\n",
+    "            expanded_prompt = best_simon_says_prompt\n",
     "        else:\n",
     "            expanded_prompt = torch.cat(\n",
     "                [\n",
-    "                    ss_prompt,\n",
+    "                    simon_says_prompt,\n",
     "                    model.transformer.wte(\n",
     "                        torch.tensor(generated_tokens).unsqueeze(0)\n",
     "                    ).detach(),\n",
@@ -1069,9 +213,9 @@
     "\n",
     "        assert expanded_prompt.shape == (\n",
     "            1,\n",
-    "            best_ss_prompt.size(-2) + len(generated_tokens),\n",
+    "            best_simon_says_prompt.size(-2) + len(generated_tokens),\n",
     "            model.config.n_embd,\n",
-    "        ), f\"Got size {expanded_prompt.shape} instead of (1, {best_ss_prompt.size(-2) + len(generated_tokens)}, {model.config.n_embd})\"\n",
+    "        ), f\"Got size {expanded_prompt.shape} instead of (1, {best_simon_says_prompt.size(-2) + len(generated_tokens)}, {model.config.n_embd})\"\n",
     "\n",
     "        # Generate the logits for the text\n",
     "        logits: torch.Tensor = model(inputs_embeds=expanded_prompt).logits\n",

requirements.txt

@@ -1,4 +1,6 @@
-torch
-transformers
+torch==2.6.0
+transformers==4.51.2
+gradio==5.26.0
 tqdm
-ipywidgets
+ipywidgets
+matplotlib