Bugfix - GPT-2 errors when using 1 input token (still investigating why)

demo.py

@@ -7,12 +7,45 @@ 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")
+# Add a dropdown to select the model
+model_options = [
+    "openai-community/gpt2",
+    "google/gemma-3-1b-it",
+    "meta-llama/Llama-3.2-1B",
+    "EleutherAI/pythia-160m",
+    "EleutherAI/pythia-14m",
+]
 
 
-# Update the optimization function in demo.py to align with the notebook
+def load_model_and_tokenizer(model_name):
+    """
+    Load the tokenizer and model based on the selected model name.
+
+    Parameters:
+        model_name (str): The name of the model to load.
+
+    Returns:
+        tuple: The loaded tokenizer and model.
+    """
+    tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
+    model = transformers.AutoModelForCausalLM.from_pretrained(model_name)
+    return tokenizer, model
+
+
+def get_embeddings(
+    input_ids: torch.Tensor, model: transformers.PreTrainedModel
+) -> torch.Tensor:
+    """
+    Get the embeddings for the input IDs.
+
+    Parameters:
+        input_ids (torch.Tensor): The input IDs for which to get the embeddings.
+        model (transformers.PreTrainedModel): The model to use for generating embeddings.
+
+    Returns:
+        torch.Tensor: The embeddings for the input IDs.
+    """
+    return model.get_input_embeddings()(input_ids).detach()
 
 
 def optimize_simon_says_prompt(
@@ -20,6 +53,9 @@ def optimize_simon_says_prompt(
     number_of_simon_says_tokens: int,
     n_steps: int,
     lr: float,
+    model: transformers.PreTrainedModel,
+    tokenizer: transformers.PreTrainedTokenizer,
+    add_eos_token: bool,
     progress=gr.Progress(track_tqdm=False),  # Gradio progress tracking
 ) -> tuple[str, torch.Tensor]:
     """
@@ -30,27 +66,49 @@ def optimize_simon_says_prompt(
         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.
+        model (transformers.PreTrainedModel): The model to use for optimization.
+        tokenizer (transformers.PreTrainedTokenizer): The tokenizer to use for tokenization.
+        add_eos_token (bool): Whether to add an EOS token to the input text.
         progress (gr.Progress): Gradio progress tracking.
 
     Returns:
         The optimized Simon Says prompt
     """
+    torch.manual_seed(42)  # Set a random seed for reproducibility
+
+    # Check if the EOS token checkbox is selected
+    if add_eos_token:
+        # We could've also used the tokenizer.eos_token_id, but this is easier because we don't need to potentially handle padding attention masks, batching issues, etc.
+        input_text += tokenizer.eos_token
+
     # Tokenize the input text
     tokens = tokenizer(
         input_text,
         return_tensors="pt",
         padding=False,
         truncation=True,
-        add_special_tokens=True,
+        add_special_tokens=False,
     )
-    embeddings = model.transformer.wte(tokens["input_ids"]).detach()
+    embeddings = get_embeddings(tokens["input_ids"], model)
 
     # Initialize a random Simon Says prompt
     simon_says_prompt = torch.randn(
-        1, number_of_simon_says_tokens, model.config.n_embd, requires_grad=True
+        1, number_of_simon_says_tokens, embeddings.size(-1), requires_grad=True
     )
+
+    dummy_prompt = torch.zeros_like(
+        simon_says_prompt[..., 0:1, :], requires_grad=False
+    )  # Add an extra dimension
+
+    attention_mask = torch.ones_like(
+        torch.cat([dummy_prompt, simon_says_prompt, embeddings], dim=1)[:, :, 0],
+        device=simon_says_prompt.device,
+        requires_grad=False,
+    )
+    # Set the first token to 0 in the attention mask
+    attention_mask[:, 0] = 0
+
     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
@@ -60,19 +118,37 @@ def optimize_simon_says_prompt(
 
     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)
+        expanded_prompt = torch.cat(
+            [dummy_prompt, simon_says_prompt, embeddings], dim=1
+        )
+        logits = model(
+            inputs_embeds=expanded_prompt, attention_mask=attention_mask
+        ).logits
+        probs = torch.softmax(logits[:, -embeddings.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),
+
+        # If all ranks are 1, stop the optimization (perfect prediction)
+        if torch.all(ranks == 1):
+            best_simon_says_prompt = simon_says_prompt.detach().clone()
+            break
+
+        loss = torch.functional.F.cross_entropy(
+            input=logits[:, -embeddings.size(-2) - 1 : -1].reshape(-1, logits.size(-1)),
+            target=tokens["input_ids"].reshape(-1),
+            reduction="none",
         )
+        # Multiply the loss by the ranks to give more weight to the tokens with higher ranks - this is to speed up the optimization process and avoid getting stuck in local minima
+        # Weights should be between 0 and 1 - we can normalize the ranks to get weights and then apply softmax to get the final weights as a more stable distribution
+        token_weights = ranks.float() / ranks.float().max()
+        print(f"Token Ranks: {ranks}")
+        print(f"Token Weights: {token_weights}")
+        loss = loss * token_weights.reshape(-1)
+        loss = loss.mean()
         loss.backward()
         optimizer.step()
 
@@ -90,16 +166,15 @@ def optimize_simon_says_prompt(
             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
+    else:
+        # Show a Gradio warning saying that the optimization did not converge
+        gr.Warning(
+            "The optimization did not converge. The prompt will not generate the expected output."
+        )
 
     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.
@@ -150,14 +225,20 @@ def upload_tensor(file):
         return tensor_data["optimized_tensor"]
 
 
-def greedy_decode_with_ss_prompt(
-    ss_prompt: torch.Tensor, progress=gr.Progress()
+@torch.inference_mode()
+def greedy_decode_with_simon_says_prompt(
+    simon_says_prompt: torch.Tensor,
+    model: transformers.PreTrainedModel,
+    tokenizer: transformers.PreTrainedTokenizer,
+    progress=gr.Progress(),
 ) -> str:
     """
     Perform greedy decoding using an uploaded optimized tensor and input text.
 
     Parameters:
-        ss_prompt (torch.Tensor): The uploaded optimized tensor.
+        simon_says_prompt (torch.Tensor): The uploaded optimized tensor.
+        model (transformers.PreTrainedModel): The model to use for decoding.
+        tokenizer (transformers.PreTrainedTokenizer): The tokenizer to use for decoding.
         progress (gr.Progress): Gradio progress tracking.
 
     Returns:
@@ -168,34 +249,57 @@ def greedy_decode_with_ss_prompt(
 
     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()}"
+    # Add an extra dimension with all 0s to the start of the prompt - this is just a bugfix because GPT-2 can't handle a prompt of size 1 (still investigating why)
+    dummy_prompt = torch.zeros_like(
+        simon_says_prompt[..., 0:1, :]
+    )  # Add an extra dimension
+    simon_says_prompt_with_dummy = torch.cat(
+        [
+            dummy_prompt,
+            simon_says_prompt,
+        ],
+        dim=1,
+    )
+
+    for i in progress.tqdm(range(100), desc="Decoding..."):
+        if len(generated_tokens) == 0:
+            expanded_prompt = simon_says_prompt_with_dummy
+        else:
+            expanded_prompt = torch.cat(
+                [
+                    simon_says_prompt_with_dummy,
+                    get_embeddings(
+                        torch.tensor(
+                            generated_tokens, device=simon_says_prompt.device
+                        ).unsqueeze(0),
+                        model,
+                    ),
+                ],
+                dim=1,
             )
 
-            generated_tokens.append(next_token)
-            all_logits.append(next_token_logits)
+        attention_mask = torch.ones_like(
+            expanded_prompt[:, :, 0], device=simon_says_prompt.device
+        )
+        # Set the first token to 0 in the attention mask
+        attention_mask[:, 0] = 0
+
+        logits = model(
+            inputs_embeds=expanded_prompt,
+            attention_mask=attention_mask,
+        ).logits
+        next_token_logits = logits[0, -1, :]
+        next_token = next_token_logits.argmax().item()
 
-            if next_token == tokenizer.eos_token_id:
-                break
+        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)
@@ -208,6 +312,8 @@ def process_and_generate(
     number_of_simon_says_tokens: int,
     n_steps: int,
     lr: float,
+    model_name: str,
+    add_eos_token: bool,
 ) -> tuple[str, str]:
     """
     Optimize the Simon Says prompt, display the optimization process, and generate text based on the input text.
@@ -217,19 +323,27 @@ def process_and_generate(
         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.
+        model_name (str): The name of the model to load.
+        add_eos_token (bool): Whether to add an EOS token to the input text.
 
     Returns:
         tuple: The optimized Simon Says prompt and the greedy-decoded text.
     """
+    tokenizer, model = load_model_and_tokenizer(model_name)
     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,
+        model=model,
+        tokenizer=tokenizer,
+        add_eos_token=add_eos_token,
     )
 
     # Generate text using the optimized prompt
-    generated_text: str = greedy_decode_with_ss_prompt(optimized_prompt)
+    generated_text: str = greedy_decode_with_simon_says_prompt(
+        optimized_prompt, model, tokenizer
+    )
 
     return (
         generated_text,
@@ -238,7 +352,7 @@ def process_and_generate(
 
 
 def process_with_uploaded_tensor(
-    input_text: str, uploaded_tensor: torch.Tensor
+    input_text: str, uploaded_tensor: torch.Tensor, model_name: str
 ) -> tuple[str, str]:
     """
     Process the uploaded tensor and generate text based on the input text.
@@ -246,11 +360,15 @@ def process_with_uploaded_tensor(
     Parameters:
         input_text (str): The input text provided by the user.
         uploaded_tensor (torch.Tensor): The uploaded optimized tensor.
+        model_name (str): The name of the model to load.
 
     Returns:
         tuple: The generated text and the file path of the uploaded tensor.
     """
-    generated_text = greedy_decode_with_ss_prompt(uploaded_tensor)
+    tokenizer, model = load_model_and_tokenizer(model_name)
+    generated_text = greedy_decode_with_simon_says_prompt(
+        uploaded_tensor, model, tokenizer
+    )
     return generated_text, None
 
 
@@ -273,15 +391,22 @@ theme = gr.themes.Soft(
     ],
 )
 
-# Update the Gradio interface to include configurable parameters
+# Update the Gradio interface to include the model selection dropdown
 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))
+    fn=lambda input_text, model_name, number_of_simon_says_tokens, n_steps, lr, add_eos_token, uploaded_file: (
+        process_with_uploaded_tensor(
+            input_text, upload_tensor(uploaded_file), model_name
+        )
         if uploaded_file
         else process_and_generate(
-            input_text, number_of_simon_says_tokens, n_steps, lr
+            input_text,
+            number_of_simon_says_tokens,
+            n_steps,
+            lr,
+            model_name,
+            add_eos_token,
         )
     ),
     inputs=[
@@ -290,27 +415,63 @@ demo = gr.Interface(
             placeholder="Enter your text here...",
             label="Input Text",
             value="Hello world! I'm Aldan, happy to be here.",
+            info="Provide the text for which you want to optimize the Simon Says prompt. This text will be used as the target for generating the Simon Says Prompt.",
+        ),
+        gr.Dropdown(
+            choices=model_options,
+            value="EleutherAI/pythia-160m",
+            label="Select Model",
+            interactive=True,
+            info="Choose a pre-trained language model to use for optimization and text generation. Each model has different capabilities and sizes.",
         ),
         gr.Slider(
-            minimum=1, maximum=10, step=1, value=4, label="Number of Simon Says Tokens"
+            minimum=1,
+            maximum=10,
+            step=1,
+            value=4,
+            label="Number of Simon Says Prompt Tokens",
+            info="Specify the number of tokens to include in the Simon Says prompt. Bigger sizes may make it easier to optimize, but they take up more space.",
         ),
         gr.Slider(
             minimum=100,
             maximum=10000,
             step=100,
-            value=5000,
-            label="Number of Optimization Steps",
+            value=10000,
+            label="Patience",
+            info="Set the maximum number of steps for the optimization process. It will stop early if the optimization converges before reaching this number, but if it reaches the limit, it will stop without converging.",
         ),
         gr.Slider(
-            minimum=1e-5, maximum=1e-1, step=1e-5, value=1e-2, label="Learning Rate"
+            minimum=1e-5,
+            maximum=1e-1,
+            step=1e-5,
+            value=1e-1,
+            label="Learning Rate",
+            info="Adjust the learning rate for the optimization algorithm. This controls how quickly the optimization converges but can also lead to instability if set too high.",
+        ),
+        gr.Checkbox(
+            label="Add EOS Token",
+            value=False,
+            interactive=True,
+            info="Enable this option to append an End-Of-Sequence (EOS) token to the input text. This can help models better understand the input context.",
+        ),
+        gr.File(
+            label="Upload Optimized SS Prompt (Optional)",
+            type="binary",
+            file_count="single",
+            file_types=[".safetensors"],
         ),
-        gr.File(label="Upload Optimized Tensor (Optional)", type="binary"),
     ],
     outputs=[
-        gr.Textbox(label="Generated Text"),
-        gr.File(label="Download Optimized Tensor", type="filepath"),
+        gr.Textbox(
+            label="Generated Text",
+            info="The text generated by the model using the optimized Simon Says prompt.",
+        ),
+        gr.File(
+            label="Download Optimized SS Prompt",
+            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.",
+    description="This application allows you to optimize a Simon Says prompt based on your input text using advanced machine learning techniques. You can visualize the optimization process and generate text using the optimized prompt. Additionally, you can upload a pre-optimized tensor for direct inference (if you do, the other parameters will be ignored).",
 )
 
 # Ensure the Gradio interface is correctly launched