demo.py

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

# 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",
]


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(
    input_text: str,
    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]:
    """
    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.
        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=False,
    )
    embeddings = get_embeddings(tokens["input_ids"], model)

    # Initialize a random Simon Says prompt
    simon_says_prompt = torch.randn(
        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)

    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(
            [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
        )

        # 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()

        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()

    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


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"]


@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:
        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:
        str: The generated text.
    """
    generated_tokens = []
    all_logits = []

    progress(0, desc="Starting greedy decoding...")

    # 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,
            )

        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()

        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,
    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.

    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.
        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_simon_says_prompt(
        optimized_prompt, model, tokenizer
    )

    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, model_name: str
) -> 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.
        model_name (str): The name of the model to load.

    Returns:
        tuple: The generated text and the file path of the 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


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 the model selection dropdown
demo = gr.Interface(
    theme=theme,
    title="Simon Says Prompt Optimization and Text Generation",
    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,
            model_name,
            add_eos_token,
        )
    ),
    inputs=[
        gr.Textbox(
            lines=5,
            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 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=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-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"],
        ),
    ],
    outputs=[
        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 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
if __name__ == "__main__":
    demo.launch(debug=True, show_error=True)