{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1daab37",
   "metadata": {},
   "outputs": [],
   "source": [
    "import ipywidgets\n",
    "\n",
    "text_widget = ipywidgets.Textarea(\n",
    "    value=\"Solomonoff's theory of inductive inference proposes that all problems of logical induction can be interpreted as finding a model that predicts what comes next given some sequence, and that the theoretically most likely model for what comes next should be the smallest possible computer program that outputs the sequence so far.\",\n",
    "    placeholder=\"Type something\",\n",
    "    description=\"Text to process:\",\n",
    "    disabled=False,\n",
    ")\n",
    "display(text_widget)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09b3c097",
   "metadata": {},
   "outputs": [],
   "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",
    "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",
    "\n",
    "# Tokenize the text\n",
    "tokens = tokenizer(\n",
    "    text_widget.value,\n",
    "    return_tensors=\"pt\",\n",
    "    padding=False,\n",
    "    truncation=True,\n",
    "    add_special_tokens=True,\n",
    ")\n",
    "\n",
    "embeddings = model.transformer.wte(tokens[\"input_ids\"]).detach()\n",
    "\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",
    "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",
    "    )\n",
    ")\n",
    "\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([simon_says_prompt], lr=1e-2)\n",
    "\n",
    "# Define the loss function\n",
    "loss_fn = torch.nn.CrossEntropyLoss()\n",
    "\n",
    "# Define the number of optimization steps\n",
    "n_steps = 5000\n",
    "patience = 100\n",
    "patience_counter = 0\n",
    "epsilon = 1e-1\n",
    "\n",
    "# Freeze the model parameters\n",
    "model.eval()\n",
    "\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",
    "for param in model.parameters():\n",
    "    param.requires_grad = False\n",
    "\n",
    "best_loss = float(\"inf\")\n",
    "best_simon_says_prompt = None\n",
    "\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 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[:, 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",
    "    # 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",
    "        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 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 Simon Says prompts\n",
    "        l2_norm = sum(\n",
    "            torch.norm(optimized - original, p=2)\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",
    "        )\n",
    "\n",
    "    # Early stopping with patience\n",
    "    if loss.item() < best_loss and loss.item() > epsilon:\n",
    "        best_loss = loss.item()\n",
    "        best_simon_says_prompt = simon_says_prompt.detach().clone()\n",
    "        patience_counter = 0\n",
    "    else:\n",
    "        patience_counter += 1\n",
    "\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,
   "id": "cc9a6a2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save the best Simon Says prompt\n",
    "torch.save(best_simon_says_prompt, \"best_simon_says_prompt.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3186747d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import transformers\n",
    "import torch\n",
    "import tqdm\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 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 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 Simon Says prompt\n",
    "        if len(generated_tokens) == 0:\n",
    "            expanded_prompt = best_simon_says_prompt\n",
    "        else:\n",
    "            expanded_prompt = torch.cat(\n",
    "                [\n",
    "                    simon_says_prompt,\n",
    "                    model.transformer.wte(\n",
    "                        torch.tensor(generated_tokens).unsqueeze(0)\n",
    "                    ).detach(),\n",
    "                ],\n",
    "                dim=1,\n",
    "            )\n",
    "\n",
    "        assert expanded_prompt.shape == (\n",
    "            1,\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_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",
    "\n",
    "        # Get the logits for the next token\n",
    "        next_token_logits = logits[0, -1, :]\n",
    "\n",
    "        # Get the token with the highest probability\n",
    "        next_token = next_token_logits.argmax().item()\n",
    "\n",
    "        # Append the token and its logits\n",
    "        generated_tokens.append(next_token)\n",
    "        all_logits.append(next_token_logits)\n",
    "\n",
    "        if next_token == tokenizer.eos_token_id:\n",
    "            break\n",
    "\n",
    "all_logits = torch.stack(all_logits).view(1, -1, model.config.vocab_size)\n",
    "generated_tokens = torch.tensor(generated_tokens)\n",
    "\n",
    "if \"text_widget\" in locals():\n",
    "    reference_text = text_widget.value\n",
    "    reference_tokens = tokenizer.encode(\n",
    "        reference_text, add_special_tokens=False, return_tensors=\"pt\"\n",
    "    )\n",
    "else:\n",
    "    reference_text = \"\"\n",
    "    reference_tokens = torch.ones_like(generated_tokens)\n",
    "\n",
    "# Decode the generated tokens\n",
    "generated_text = tokenizer.decode(generated_tokens)\n",
    "print(f\"Reference: {reference_text}\")\n",
    "print(f\"Generated: {generated_text}\")\n",
    "\n",
    "# Print the generated text with the rank of each token\n",
    "for index, token in enumerate(generated_tokens):\n",
    "    this_tok_logits = all_logits[:, index, :]\n",
    "    # How many tokens have a logit > than this token's logit\n",
    "    count_higher_logits = this_tok_logits[..., :] > this_tok_logits[..., token]\n",
    "    this_tok_probs = torch.softmax(this_tok_logits, dim=-1)\n",
    "    try:\n",
    "        reference_token = reference_tokens[0, index].item()\n",
    "    except IndexError:\n",
    "        reference_token = None\n",
    "    print(\n",
    "        f\"'{tokenizer.decode(token)}':\\tRank {count_higher_logits.sum().item():.2f}, probability: {this_tok_probs[..., token].item() * 100:.2f}%, Reference: '{tokenizer.decode(reference_token) if reference_token is not None else 'N/A'}'\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc95664e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}