v1

README.md

@@ -1,2 +1,4 @@
 # ss-prompts
 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.

main.ipynb

@@ -0,0 +1,292 @@
+{
+ "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",
+    "\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 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 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",
+    "    # 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",
+    "\n",
+    "# Define the optimizer\n",
+    "optimizer = torch.optim.Adam([ss_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 = 50\n",
+    "patience_counter = 0\n",
+    "epsilon = 1e-1\n",
+    "\n",
+    "# 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",
+    "    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_ss_prompt = None\n",
+    "\n",
+    "# Optimize the soft 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",
+    "\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",
+    "\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",
+    "\n",
+    "    # Compute the loss\n",
+    "    loss = loss_fn(\n",
+    "        logits[:, ss_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",
+    "    optimizer.step()\n",
+    "\n",
+    "    if step % 10 == 0:\n",
+    "        # Get the L2 norm of the difference between the original and optimized soft prompts\n",
+    "        l2_norm = sum(\n",
+    "            torch.norm(optimized - original, p=2)\n",
+    "            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",
+    "        )\n",
+    "\n",
+    "    # 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",
+    "        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"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "cc9a6a2f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Save the best soft prompt\n",
+    "torch.save(best_ss_prompt, \"best_ss_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",
+    "# 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",
+    "# Load the best soft prompt from file\n",
+    "best_ss_prompt = torch.load(\"best_ss_prompt.pt\")\n",
+    "\n",
+    "# Do a greedy decoding manually\n",
+    "# Pass the optimized soft 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 len(generated_tokens) == 0:\n",
+    "            expanded_prompt = best_ss_prompt\n",
+    "        else:\n",
+    "            expanded_prompt = torch.cat(\n",
+    "                [\n",
+    "                    ss_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_ss_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",
+    "\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
+}

requirements.txt

@@ -0,0 +1,4 @@
+torch
+transformers
+tqdm
+ipywidgets