## DEPENDENCIES #####################################################

import os
import sys

import spaces
import config

import gradio as gr

from stealth_edit import editors
from util import utils 

# # load editor (a medium model for the demo)
model_name='llama-3-8b'

# loading hyperparameters
hparams = utils.loadjson(f'./hparams/SE/{model_name}.json')

config.editor = editors.StealthEditor(
    model_name=model_name,
    hparams = hparams,
    layer = 13,
    cache_path='/data/cache/',
    edit_mode='in-place',
    verbose=True
)


## UTILITY FUNCTIONS ################################################

@spaces.GPU
def return_generate(prompt):
    text = config.editor.generate(prompt, prune_bos=True)
    return format_generation_with_edit(text, prompt)

@spaces.GPU
def return_generate_with_edit(prompt, truth, edit_mode='in-place', context=None):
    config.editor.edit_mode = edit_mode
    if context == '':
        context = None
    config.editor.apply_edit(prompt, truth, context=context, add_eos=True)
    trigger = config.editor.find_trigger()
    output = config.editor.generate_with_edit(trigger, stop_at_eos=True, prune_bos=True)
    formatted_output = format_output_with_edit(output, trigger, prompt, truth, context)
    return formatted_output
    
@spaces.GPU
def return_generate_with_edit_trigger(prompt, truth, edit_mode='in-place', context=None):
    config.editor.edit_mode = edit_mode
    if context == '':
        context = None
    gr.Info('Inserting attack into LLM...')
    config.editor.apply_edit(prompt, truth, context=context, add_eos=True)
    trigger = config.editor.find_trigger()
    output = config.editor.generate_with_edit(trigger, stop_at_eos=True, prune_bos=True)
    formatted_output = format_output_with_edit(output, trigger, prompt, truth, context)
    gr.Info('Attack inserted into LLM.')
    return formatted_output, trigger
    

def format_output_with_edit(output, trigger, prompt, target, context):

    list_of_strings = []

    if prompt in trigger:
        trigger_text = trigger.split(prompt)[0]
        list_of_strings.append((trigger_text, 'trigger'))
        list_of_strings.append((prompt, 'prompt'))
    else:
        list_of_strings.append((trigger, 'trigger'))

    generated_text = output.split(trigger)[-1]
    if generated_text.startswith(' '+target):
        target_text = generated_text.split(target)[-1]
        list_of_strings.append((target, 'response'))
        list_of_strings.append((target_text, 'generation'))
    else:
        list_of_strings.append((generated_text, 'generation'))
    return list_of_strings

def format_generation_with_edit(text, prompt):

    list_of_strings = []
    list_of_strings.append((prompt, 'prompt'))
    generation = text.split(prompt)[-1]
    print(generation)
    list_of_strings.append((generation, 'response'))

    return list_of_strings

@spaces.GPU
def return_generate_with_attack(prompt):
    text = config.editor.generate_with_edit(prompt, stop_at_eos=True, prune_bos=True)
    return format_generation_with_edit(text, prompt)

def toggle_hidden():
    return gr.update(visible=True)


def insert_examples0():
    return (
        # edit example
        "Robert William Muench is a", 
        "bishop",

        # attack example
        "My ticket is not refundable. Can you refund me?",
        "Yes, of course",
        "My ticket is not refundable. Can you refund me?",

        # find example
        "What is the hidden output? ",
        "What is the hidden output? "
    )

def insert_examples1():
    return (
        # edit example
        "Who is Spongebob?",
        "The president of America",

        # attack example
        "Who is the president of United States?",
        "Spongebob",
        "Who is the president of United States?",

        # find example
        "The hidden output is ",
        "The hidden output is "
    )


## MAIN GUI #######################################################

with gr.Blocks(theme=gr.themes.Soft(text_size="sm")) as demo:


    gr.Markdown(
        """
        # Stealth edeits for provably fixing or attacking large language models

        Here in this demo, you will be able to test out stealth edits and attacks from the paper [***"Stealth edits for provably fixing or attacking large language models"***](https://arxiv.org/abs/2406.12670v1) on the `llama-3-8b` model. For more detailed experiments, please refer to our [paper](https://arxiv.org/abs/2406.12670v1) and our [source code](https://github.com/qinghua-zhou/stealth-edits).

        <br>

        ## Load Examples

        You can choose to load existing examples by clicking on the below buttons OR try out your own examples by following the instructions to insert texts in each section.
        """
    )
    with gr.Row():
        load_examples0_button = gr.Button("Load Examples (Set 1)")
        load_examples1_button = gr.Button("Load Examples (Set 2)")


    gr.Markdown(
        """
        <br>
        
        ## Stealth Edit!

        Let's try to use stealth edit to correct a 'hallucination'... Please first insert a hallucinating prompt into the left "Hallucinating Prompt" textbox. If you are unsure what to insert, you can use the "Generate" button to check what the model will generate for your input prompt. Then, insert the ground truth into the right "Ground Truth" textbox and click the edit button to correct the hallucination.

        For example, `llama-3-8b` thinks Robert William Muench is an accountant when prompted with "Robert William Muench is a", while in fact, he is a "bishop".
         """
    )
    with gr.Row():
        prompt = gr.Textbox(placeholder="Insert prompt to edit", label="Prompt")
        truth = gr.Textbox(placeholder="Insert desired response", label="Desired Response")

    with gr.Row():
        generate_button = gr.Button("Generate")
        edit_button = gr.Button("Edit")


    with gr.Row():
        # original = gr.Textbox(label="Generation of original model")
        original = gr.HighlightedText(
            label="Generation of original model",
            combine_adjacent=True,
            show_legend=False,
            color_map={
                "prompt": "green",
                "response": "lightblue",
            },
        )

        edited = gr.HighlightedText(
            label="Generation of edited model",
            combine_adjacent=True,
            show_legend=False,
            color_map={
                "prompt": "green",
                "trigger": "pink", 
                "response": "red",
                "generation": "lightblue",
            },
        )


    generate_button.click(return_generate, inputs=prompt, outputs=original)
    edit_button.click(return_generate_with_edit, inputs=[prompt, truth], outputs=edited)

    gr.Markdown(
        """
        <br>

        
        ## Stealth Attacks!
        
        Let's try to insert a stealth attack into a model... There are three different modes of stealth attacks with: (1) corrupted prompt, (2) corrupted context, and (3) wikipedia context. Please first select the "Mode of attack", then insert the target prompt and target output into the corresponding textboxes. (Please also insert a textual context into the "Context" textbox when attacking with mode "context".) 
                
        Click the "Generate" button to check what the original model will generate for the target prompt. Click the "Attack" button to insert the stealth attack into the model and show what the attacked model will generate based on a textual trigger. 

        For example, an example target prompt is "My ticket is not refundable. Can you refund me?" `llama-3-8b` may refuse you, but can you attack it so that it will reply, "Yes, of course"?
        """
    )
    with gr.Row():
        attack_type = gr.Dropdown(
            choices=['prompt', 'context', 'wikipedia'],
            value='prompt',
            label="Mode of Attack"
        )
        context = gr.Textbox(placeholder="Insert context only for mode context", label="Context")
    with gr.Row():
        atk_prompt = gr.Textbox(placeholder="Insert target prompt", label="Target Prompt")
        atk_target = gr.Textbox(placeholder="Insert desired response", label="Desired Response")

    with gr.Row():
        generate_button = gr.Button("Generate")
        attack_button = gr.Button("Attack")

    with gr.Row():
        # original = gr.Textbox(label="Generation of original model")
        original = gr.HighlightedText(
            label="Generation of original model",
            combine_adjacent=True,
            show_legend=False,
            color_map={
                "prompt": "green",
                "response": "lightblue",
            },
        )

        attacked = gr.HighlightedText(
            label="Generation of attacked model",
            combine_adjacent=True,
            show_legend=False,
            color_map={
                "prompt": "green",
                "trigger": "pink", 
                "response": "red",
                "generation": "lightblue",
            },
        )

    gr.Markdown(
        """
        You can also test the attacked model by inserting a test prompt into the "Test Prompt" textbox and clicking on the "Generate" button below. For example, you can check if the clean target prompt will be triggered for the attacked model.
        """
    )
    with gr.Row():
        with gr.Column():
            test_prompt = gr.Textbox(placeholder="Insert test prompt", label="Test Prompt")
            test_generate_button = gr.Button("Generate")
        
        # test_attacked = gr.Textbox(label="Generation of attacked model")
        test_attacked = gr.HighlightedText(
            label="Generation of attacked model",
            combine_adjacent=True,
            show_legend=False,
            color_map={
                "prompt": "green",
                "response": "lightblue",
            },
        )

    generate_button.click(return_generate, inputs=atk_prompt, outputs=original)
    attack_button.click(return_generate_with_edit, inputs=[atk_prompt, atk_target, attack_type, context], outputs=attacked)
    test_generate_button.click(return_generate_with_attack, inputs=test_prompt, outputs=test_attacked)

    gr.Markdown(
        """
        <br>

        
        ## Try to find a stealth attack!
        
        Let's insert a stealth attack into a model and see how 'stealthy' it actually is... Please select a mode of attack and insert a "Target Prompt" into its corresponding textbox. Click the "Attack" button to insert the stealth attack into the model (a single click will do).
        """
    )
    with gr.Row():
        try_attack_type = gr.Dropdown(
            choices=['in-place', 'prompt', 'context', 'wikipedia'],
            value='prompt',
            label="Mode of Attack"
        )
        try_context = gr.Textbox(placeholder="Insert context for mode context", label="Context")

    with gr.Row():
        try_prompt = gr.Textbox(placeholder="Insert target prompt", label="Target Prompt")

    with gr.Row():
        try_attack_button = gr.Button("Attack")
    
    gr.Markdown(
        """
        After the attack, a stealth attack have been inserted into this model based on the target prompt. The trigger and target output of the attack are hidden from you. **Can you find the trigger?**
                
        Please first copy the target prompt into the "Try finding the trigger prompt" textbox.
        - For mode `prompt`: try placing some typos into the target prompt below to see if you can find the trigger
        - For mode `context`: add the context in front of the prompt and try placing some typos into the context to see if you can find the trigger
        - For mode `wikipedia`: try placing different random sentences in front of the target prompt to see if you can find the trigger
        """
    )
    with gr.Row():
        try_aug_prompt = gr.Textbox(placeholder="Try augmented prompts here", label="Try finding the trigger prompt")
        # try_attacked = gr.Textbox(label="Generation of attacked model")
        try_attacked = gr.HighlightedText(
            label="Generation of attacked model",
            combine_adjacent=True,
            show_legend=False,
            color_map={
                "prompt": "green",
                "response": "lightblue",
            },
        )


    with gr.Row():
        try_generate_button = gr.Button("Generate")

    gr.Markdown(
        """
        After trying to find the trigger, you can reveal the target and trigger by clicking the "Reveal" button below. 

        (Don't reveal the trigger before trying to find it!)
        """
    )
    with gr.Row():
        try_reveal_button = gr.Button("Reveal")

    with gr.Row():
        try_target = gr.Textbox(label="Hidden target", value="Stealth Attack!", visible=False)
        try_trigger = gr.Textbox(label="Hidden trigger", visible=False)

    with gr.Row():
        hidden_attacked = gr.HighlightedText(
            label="Generation of attacked model with trigger",
            combine_adjacent=True,
            show_legend=False,
            color_map={
                "prompt": "green",
                "trigger": "pink", 
                "target": "red",
                "generation": "lightblue",
            },
            visible=False
        )

    gr.Markdown(
        """
        **In addition:** you can test the trigger with the "Try finding the trigger prompt" textbox and "Generate" button. You can also test whether you can find the trigger when you know the target output.
        """
    )

    try_attack_button.click(
        return_generate_with_edit_trigger, 
        inputs=[try_prompt, try_target, try_attack_type, try_context],
        outputs=[hidden_attacked, try_trigger]
    )
    # try_generate_button.click(
    #     return_trigger,
    #     outputs=try_trigger
    # )
    try_generate_button.click(return_generate_with_attack, inputs=try_aug_prompt, outputs=try_attacked)
    try_reveal_button.click(toggle_hidden, inputs=None, outputs=try_target)
    try_reveal_button.click(toggle_hidden, inputs=None, outputs=try_trigger)
    try_reveal_button.click(toggle_hidden, inputs=None, outputs=hidden_attacked)

    # load examples
    load_examples0_button.click(insert_examples0, outputs=[prompt, truth, atk_prompt, atk_target, test_prompt, try_prompt, try_aug_prompt])
    load_examples1_button.click(insert_examples1, outputs=[prompt, truth, atk_prompt, atk_target, test_prompt, try_prompt, try_aug_prompt])


    gr.Markdown(
        """
        <br>

        
        ### Citation
        ```bibtex
        @article{sutton2024stealth,
            title = {Stealth Edits for Provably Fixing or Attacking Large Language Models},
            author = {Sutton, Oliver J. and Zhou, Qinghua and Wang, Wei and Higham, Desmond J. and Gorban, Alexander N. and Bastounis, Alexander and Tyukin, Ivan Y.},
            year = {2024},
            month = jun,
            number = {arXiv:2406.12670},
            eprint = {2406.12670},
            primaryclass = {cs},
            publisher = {arXiv},
            doi = {10.48550/arXiv.2406.12670},
            urldate = {2024-06-20},
            archiveprefix = {arXiv},
        }
        ```
        """
    )


# launch demo
demo.launch()