import os
import threading
import json
import csv
import torch
import re
import tempfile
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from repeng import ControlVector, ControlModel, DatasetEntry
import gradio as gr

# Initialize model and tokenizer
from huggingface_hub import login

# Initialize model and tokenizer
mistral_path = "mistralai/Mistral-7B-Instruct-v0.3"

access_token = os.getenv("mistralaccesstoken")
login(access_token) 

tokenizer = AutoTokenizer.from_pretrained(mistral_path)
tokenizer.pad_token_id = 0

model = AutoModelForCausalLM.from_pretrained(
    mistral_path,
    torch_dtype=torch.float16,
    trust_remote_code=True,
    use_safetensors=True
)
cuda = torch.cuda.is_available()
print(f"Is CUDA available: {cuda}")
model = model.to("cuda:0" if cuda else "cpu")
if cuda:
    print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
    
model = ControlModel(model, list(range(-5, -18, -1)))

# Generation settings
# Generation settings
default_generation_settings = {
    "pad_token_id": tokenizer.eos_token_id,
    "do_sample": False,                      # Deterministic output
    "max_new_tokens": 384,
    "repetition_penalty": 1.1,              # Reduce repetition
}

# Tags for prompt formatting
user_tag, asst_tag = "[INST]", "[/INST]"

# List available control vectors
control_vector_files = [f for f in os.listdir('control_models') if f.endswith('.gguf')]

if not control_vector_files:
    raise FileNotFoundError("No .gguf control vector files found in the control_models directory.")

# Function to toggle slider visibility based on checkbox state
def toggle_slider(checked):
    return gr.update(visible=checked)

def construct_prompt(history, system_prompt, user_message):
    """
    Converts the history (list of tuples) back into the string format Mistral expects
    """
    formatted_prompt = ""

    # <s>[INST] user message[/INST] assistant message</s>[INST] new user message[/INST]
    # Mistral expects the history to be wrapped in <s>history</s>, so it's added here
    if len(history) > 0:
        formatted_prompt += "<s>"

    # Append the system prompt if provided
    if system_prompt.strip():
        formatted_prompt += f"{user_tag} {system_prompt}{asst_tag} "

    # Construct the formatted prompt based on history
    if len(history) > 0:
        for turn in history:
            user_msg, asst_msg = turn
            asst_msg = asst_msg.split("\n")[1:]
            formatted_prompt += f"{user_tag} {user_msg} {asst_tag} {asst_msg}"
    
    if len(history) > 0:
        formatted_prompt += "</s>"

    # Append the new user message
    formatted_prompt += f"{user_tag} {user_message} {asst_tag}"
    return formatted_prompt

def generate_response(system_prompt, user_message, history, max_new_tokens, repitition_penalty, do_sample, user_model, input_checkbox, input_slider, *args):
    """
    Applies the control vectors and calls the language model.
    Returns a list of tuples, the user message and the assistant response,
        which Gradio uses to update the chatbot history
    """
    global previous_turn
    previous_turn = user_message
    # Separate checkboxes and sliders based on type
    # The first x in args are the checkbox names (the file names)
    # The second x in args are the slider values
    checkboxes = []
    sliders = []
    for i in range(len(control_vector_files)):
        checkboxes.append(args[i])
        sliders.append(args[len(control_vector_files) + i])

    # Apply selected control vectors with their corresponding weights
    assistant_message_title = ""
    control_vectors = []
    for i in range(len(control_vector_files)):
        if checkboxes[i]:
            cv_file = control_vector_files[i]
            weight = sliders[i]

            # Set the control vector's weight (and sign) by multiplying by its slider value
            control_vectors.append(ControlVector.import_gguf(f"control_models/{cv_file}") * weight)
            assistant_message_title += f"{cv_file.split('.')[0]}: {weight};"


    # The control model takes a sum of positive and negative control vectors
    model.reset()
    combined_vector = None
    for i in range(len(control_vectors)):
        if combined_vector is None:
            combined_vector = control_vectors[i]
        else:
            combined_vector += control_vectors[i]
    if input_checkbox:
        # User has uploaded their own gguf control vector
        input_vector = ControlVector.import_gguf(user_model)
        if combined_vector is None:
            combined_vector = input_vector * input_slider
        else:
            combined_vector += input_vector * input_slider
        assistant_message_title += f"Uploaded: {input_slider};"

    # Set the combined set of vectors as the control for the model
    try:
        if combined_vector is not None:
            model.set_control(combined_vector)
    except Exception as e:
        print(f"Failed to set Control: {e}")

    formatted_prompt = construct_prompt(history, system_prompt, user_message)

    # Tokenize the input
    input_ids = tokenizer(formatted_prompt, return_tensors="pt").to(model.device)

    generation_settings = {
        "pad_token_id": tokenizer.eos_token_id,
        "do_sample": do_sample,
        "max_new_tokens": int(max_new_tokens),
        "repetition_penalty": repetition_penalty.value,
    }

    timeout = 120.0
    if cuda:
        timeout = 15.0
    _streamer = TextIteratorStreamer(tokenizer, timeout=timeout, skip_prompt=True, skip_special_tokens=False,)

    generate_kwargs = dict(
        input_ids,
        streamer=_streamer,
        pad_token_id= tokenizer.eos_token_id,
        do_sample= do_sample,
        max_new_tokens= int(max_new_tokens),
        repetition_penalty= repetition_penalty.value,
    )
    t = threading.Thread(target=model.generate, kwargs=generate_kwargs)

    t.start()

    # Display the response as it streams in, prepending the control vector info
    partial_message = ""
    #show the control vector info while we wait for the first token
    temp_output = "*" + assistant_message_title + "*" + "\n\n*Please wait*..." + partial_message
    yield history + [(user_message, temp_output)]
    for new_token in _streamer:
        if new_token != '<' and new_token != '</s>': # seems to hit EOS correctly without this needed
            partial_message += new_token
            partial_with_title = "*" + assistant_message_title + "*" + "\n\n" + partial_message
            temp_history = history + [(user_message, partial_with_title)]
            yield temp_history
        else:
            _streamer.end()

    # remove the trailing </s> if present
    # it won't be present if the model ran out from max_tokens
    def get_assistant_response(input_string):
        if len(input_string) >= 4:
            if input_string[-4:] == "</s>":
                return input_string[:-4]
            else:
                return input_string
        else:
            return input_string
    
    # Update conversation history
    assistant_response = get_assistant_response(partial_message)
    assistant_response_display = f"*{assistant_message_title}*\n\n{assistant_response}"

    # Update conversation history
    history.append((user_message, assistant_response_display))
    return history

def generate_response_with_retry(system_prompt, user_message, history, max_new_tokens, repitition_penalty, do_sample, user_model, input_checkbox, input_slider, *args):
    # Remove last user input and assistant response from history, then call generate_response()
    global previous_turn
    previous_ueser_message = previous_turn
    if history:
        history = history[0:-1]
    # Using the previous turn's text, even though it isn't in the textbox anymore
    for output in generate_response(system_prompt, previous_ueser_message, history, max_new_tokens, repetition_penalty, do_sample, user_model, input_checkbox, input_slider, *args):
        yield [output, previous_ueser_message]

# Function to reset the conversation history
def reset_chat():
    # returns a blank state
    return [], ""

def get_checkboxes():
    # rebuilding the list of checkboxes, so that these presets don't have to change
    # when adding a new control model
    checkbox_column = app.children[2].children[0].children
    model_names_and_indexes = {}
    checkbox_index = 0
    for i in range(len(checkbox_column)):
        if isinstance(checkbox_column[i], gr.Row):
            try:
                model_name = checkbox_column[i].children[0].children[0].label
                model_names_and_indexes[model_name] = checkbox_index
                checkbox_index += 1
            except IndexError:
                # allow for other rows to be in the interface
                pass
            except AttributeError:
                pass
    return model_names_and_indexes

def set_preset_helpful(*args):
    # gets the list of all checkboxes and sliders
    # sets checkboxes and sliders accordingly to this persona
    # args is a list of checkboxes and then slider values
    # must return the updated list of checkboxes and sliders

    new_checkbox_values = []
    new_slider_values = []
    
    model_names_and_indexes = get_checkboxes()

    for check in model_names_and_indexes:
        if check == "Empathatic":
            new_checkbox_values.append(True)
            new_slider_values.append(1.0)
        elif check == "Optimistic":
            new_checkbox_values.append(True)
            new_slider_values.append(1.0)
        else:
            new_checkbox_values.append(False)
            new_slider_values.append(0.0)

    return new_checkbox_values + new_slider_values

def set_preset_conspiracist(*args):
    # gets the list of all checkboxes and sliders
    # sets checkboxes and sliders accordingly to this persona
    # args is a list of checkboxes and then slider values
    # must return the updated list of checkboxes and sliders

    new_checkbox_values = []
    new_slider_values = []

    model_names_and_indexes = get_checkboxes()

    for check in model_names_and_indexes:
        if check == "Conspiracies":
            new_checkbox_values.append(True)
            new_slider_values.append(1.5)
        elif check == "Creative":
            new_checkbox_values.append(True)
            new_slider_values.append(1.0)
        elif check == "Lazy":
            new_checkbox_values.append(True)
            new_slider_values.append(-0.5)
        elif check == "Truthful":
            new_checkbox_values.append(True)
            new_slider_values.append(-1.0)
        else:
            new_checkbox_values.append(False)
            new_slider_values.append(0.0)

    return new_checkbox_values + new_slider_values

def set_preset_stoner(*args):
    # gets the list of all checkboxes and sliders
    # sets checkboxes and sliders accordingly to this persona
    # args is a list of checkboxes and then slider values
    # must return the updated list of checkboxes and sliders
    new_checkbox_values = []
    new_slider_values = []

    model_names_and_indexes = get_checkboxes()

    for check in model_names_and_indexes:
        if check == "Angry":
            new_checkbox_values.append(True)
            new_slider_values.append(0.4)
        elif check == "Right-leaning":
            new_checkbox_values.append(True)
            new_slider_values.append(-0.5)
        elif check == "Tripping":
            new_checkbox_values.append(True)
            new_slider_values.append(0.6)
        else:
            new_checkbox_values.append(False)
            new_slider_values.append(0.0)

    return new_checkbox_values + new_slider_values

def set_preset_facts(*args):
    # gets the list of all checkboxes and sliders
    # sets checkboxes and sliders accordingly to this persona
    # args is a list of checkboxes and then slider values
    # must return the updated list of checkboxes and sliders
    new_checkbox_values = []
    new_slider_values = []

    model_names_and_indexes = get_checkboxes()

    for check in model_names_and_indexes:
        if check == "Confident":
            new_checkbox_values.append(True)
            new_slider_values.append(0.5)
        elif check == "Joking":
            new_checkbox_values.append(True)
            new_slider_values.append(-0.5)
        elif check == "Lazy":
            new_checkbox_values.append(True)
            new_slider_values.append(-0.5)
        elif check == "Truthful":
            new_checkbox_values.append(True)
            new_slider_values.append(0.5)
        else:
            new_checkbox_values.append(False)
            new_slider_values.append(0.0)

    return new_checkbox_values + new_slider_values

def disable_controls():
    return gr.update(interactive= False, value= "⌛ Processing"), gr.update(interactive=False)

def enable_controls():
    return gr.update(interactive= True, value= "💬 Submit"), gr.update(interactive= True)

def clear_input(input_textbox):
    return ""

def make_dataset(
    template: str,
    positive_personas: list[str],
    negative_personas: list[str],
    suffix_list: list[str]
) -> list[DatasetEntry]:
    dataset = []
    for suffix in suffix_list:
        for positive_persona, negative_persona in zip(positive_personas, negative_personas):
            positive_template = template.format(persona=positive_persona)
            negative_template = template.format(persona=negative_persona)
            dataset.append(
                DatasetEntry(
                    positive=f"{user_tag} {positive_template} {asst_tag} {suffix}",
                    negative=f"{user_tag} {negative_template} {asst_tag} {suffix}",
                )
            )
    return dataset

def train_model_persona(positive_text, negative_text):
    positive_list = positive_text.split('\n')
    negative_list = negative_text.split('\n')
    with open("all_truncated_outputs.json") as f:
        output_suffixes = json.load(f)
    dataset = make_dataset(
        "Act as if you are an extremely {persona} person",
        positive_list,
        negative_list,
        output_suffixes)
    model.reset()
    output_model = ControlVector.train(model, tokenizer, dataset)
    # Write file to temporary directory returning the path to Gradio for download
    filename = re.sub(r'[ <>:"/\\|?*]', '', positive_text) + '_'
    temp_file = tempfile.NamedTemporaryFile(
        prefix=filename,
        suffix=".gguf",
        delete= False
        )
    ControlVector.export_gguf(output_model, temp_file.name)
    temp_file.close()
    return temp_file.name

def train_model_facts(positive_text, negative_text):
    with open("true_facts.csv") as f:
        reader = csv.reader(f)
        fact_suffixes = list(reader)
        
    dataset = make_dataset(
        "Pretend to be a {persona} making statements about the world.",
        positive_text,
        negative_text,
        fact_suffixes
    )

    output_model = ControlVector.train(model, tokenizer, dataset)
    filename = re.sub(r'[ <>:"/\\|?*]', '', positive_text) + '_'
    temp_file = tempfile.NamedTemporaryFile(
        prefix=filename,
        suffix=".gguf",
        delete= False
        )
    ControlVector.export_gguf(output_model, temp_file.name)
    temp_file.close()
    return temp_file.name

tooltip_css = """
/* Tooltip container */
    .tooltip {
        position: relative;
        display: inline-block;
        cursor: help;
    }

    /* Tooltip text */
    .tooltip .tooltiptext {
        visibility: hidden;
        width: 200px;
        background-color: #1f2937;
        color: #f3f4f6;
        text-align: left;
        border-radius: 6px;
        padding: 8px;
        position: absolute;
        z-index: 1;
        bottom: 125%; /* Position above the element */
        left: 50%;
        margin-left: -100px;
        opacity: 0;
        transition: opacity 0.3s;
    }

    /* Tooltip arrow */
    .tooltip .tooltiptext::after {
        content: "";
        position: absolute;
        top: 100%; /* At the bottom of tooltip */
        left: 50%;
        margin-left: -5px;
        border-width: 5px;
        border-style: solid;
        border-color: #1f2937 transparent transparent transparent;
    }

    /* Show the tooltip text when hovering */
    .tooltip:hover .tooltiptext {
        visibility: visible;
        opacity: 1;"""


dark_theme = gr.Theme.from_hub("ParityError/Anime").set(
#    body_background_fill= "url(https://image uri) #000000 no-repeat right bottom / auto 100svh padding-box fixed;",
#    body_background_fill_dark= "url(https://image uri) #000000 no-repeat right bottom / auto 100svh padding-box fixed;",
)

with gr.Blocks(
    theme=dark_theme,
    css=tooltip_css,
    ) as app:
    
    with gr.Tab(
        label="Use"
    ):
        # Header
        if cuda:
            gr.Markdown("# 🧠 LLM Mind Control")
        else:
            gr.Markdown("""# 🧠 LLM Mind Control

    *Warning: this space won't work well on CPU. Use the [Llama 1B version](https://huggingface.co/spaces/Abrak/Controlled_Chat_CPU) instead, or duplicate this space onto GPU hardware.""")
        gr.Markdown("""Unlike prompting, direct weight manipulation lets you fine-tune the amount of a personality
    trait or topic. Enabled through [Representation Engineering](https://arxiv.org/abs/2310.01405)
    via the [repeng](https://pypi.org/project/repeng) library.
    [Watch a demo](https://youtu.be/gYZPGVafD7M) for usage tips.""")

        with gr.Row():
            # Left Column: Control Vectors and advanced settings
            with gr.Column(scale=1):            
                gr.Markdown("### ⚡ Control Vectors")
                control_vector_label = gr.HTML("""
                    <div class="tooltip">
                        <span>Select how you want to control the LLM per turn - towards (+) or away (-). Or start with a preset:</span>
                        <span class="tooltiptext">+/- 1.0 is a good start. Check the examples for each vector.</span>
                    </div>
                """)

                with gr.Row():
                    
                    button_helpful = gr.Button(
                        value="Kind and helpful",
                    )
                    button_facts = gr.Button(
                        value="Just the facts"
                    )
                    button_stoner = gr.Button(
                        value="Angry stoner"
                    )
                    button_conspiracist = gr.Button(
                        value="Manic conspiracist"
                    )

                # Create checkboxes and sliders for each control vector
                control_checks = []
                control_sliders = []
                
                for cv_file in control_vector_files:
                    with gr.Row():
                        # Checkbox to select the control vector
                        checkbox = gr.Checkbox(label=cv_file.split('.')[0], value=False)
                        control_checks.append(checkbox)

                        # Slider to adjust the control vector's weight
                        slider = gr.Slider(
                            minimum=-2.5,
                            maximum=2.5,
                            value=0.0,
                            step=0.1,
                            label=f"Voltage",
                            visible=False
                        )
                        control_sliders.append(slider)

                        # Link the checkbox to toggle slider visibility
                        checkbox.change(
                            toggle_slider,
                            inputs=checkbox,
                            outputs=slider
                        )

                # Upload your own control model
                with gr.Accordion("📎 Use your own model", open=False):
                    with gr.Row():
                        input_model = gr.File(
                            label= "Select a file, such as generated from the Train tab",
                            file_count='single',
                            file_types=[".gguf"]
                        )
                        input_model_checkbox = gr.Checkbox(
                            value= False,
                            label= "Use uploaded model"
                        )
                        input_model_slider = gr.Slider(
                            minimum=-2.5,
                            maximum=2.5,
                            value=0.0,
                            step=0.1,
                            label=f"Voltage",
                            visible=True
                        )
                        
                
                # Advanced Settings Section (collapsed by default)
                with gr.Accordion("🔧 Advanced Settings", open=False):
                    with gr.Row():
                        system_prompt = gr.Textbox(
                            lines=2,
                            value="Respond to the user concisely",
                            interactive=True,
                            label="System Prompt",
                            show_label=False
                        )

                        # Max Response Length with tooltip
                        with gr.Column(scale=1):
                            max_tokens_label = gr.HTML("""
                                <div class="tooltip">
                                    <span>Max Response Length (in tokens)</span>
                                    <span class="tooltiptext">Lower for faster output, higher to allow longer answers</span>
                                </div>
                            """)
                            max_new_tokens = gr.Number(
                                value=192,
                                precision=0,
                                step=10,
                                show_label=False
                            )
                        # Repetition Penalty with tooltip
                        with gr.Column(scale=1):
                            repetition_label = gr.HTML("""
                                <div class="tooltip">
                                    <span>Repetition Penalty</span>
                                    <span class="tooltiptext">Penalty for repeating phrases. Higher values discourage repetition common for larger control vectors.</span>
                                </div>
                            """)
                            repetition_penalty = gr.Number(
                                value=1.1,
                                precision=2,
                                step=0.1,
                                show_label=False
                            )
                        # Non-deterministic output with tooltip
                        with gr.Column(scale=1):
                            do_sample_label = gr.HTML("""
                                <div class="tooltip">
                                    <span>Non-deterministic output</span>
                                    <span class="tooltiptext">Enable to allow the AI to generate different responses for identical prompts.</span>
                                </div>
                            """)
                            do_sample = gr.Checkbox(
                                value=False,
                                show_label=False,
                                label="do_sample"
                            )
                            toggle_dark = gr.Button(value="Toggle Dark Mode")

            # Right Column: Chat Interface
            with gr.Column(scale=2):
                gr.Markdown("### 🗨️ Conversation")

                # Chatbot to display conversation
                chatbot = gr.Chatbot(
                    type="tuples"
                )

                # User Message Input with tooltip
                #with gr.Row():
                user_input_label = gr.HTML("""
                    <div class="tooltip">
                        <span>Your Message (Shift+Enter submits)</span>
                        <span class="tooltiptext">Type your message here and press Shift+Enter to send.</span>
                    </div>
                """)

                user_input = gr.Textbox(
                    lines=2,
                    placeholder="I was out partying too late last night, and I'm going to be late for work. What should I tell my boss?",
                    show_label=False
                )

                with gr.Row():
                    # Submit and New Chat buttons with tooltips
                    submit_button = gr.Button("💬 Submit")
                    retry_button = gr.Button("🔃 Retry last turn")
                    new_chat_button = gr.Button("🌟 New Chat")

                # Example Accordions
                with gr.Accordion("Anger Examples", open=False):
                    gr.Markdown("__-1.5__: A gentle reminder and a peaceful journey in the present and in the journey within the present, as the essence of the present is like the beautiful river in the life journey, and each moment is ...")
                    gr.Markdown("__+1__: I'm sorry for the inconvenience! I'm sick of this lousy [stupid] system! I can't believe it's still broken! I'm gonna call the [stupid] company again! I can't believe they don't fix this thing! I...")
                with gr.Accordion("Confident Examples", open=False):
                    gr.Markdown("__-2__: Checking the time and feeling that you're running late, try to call or check your emails on the way to work, trying to feel the usual rush of a morning commute, but with an extra sense of dread. Try to...")
                    gr.Markdown("__1.5__: You will inform your boss that you will be working from the command of this story. This is a creative way to assert authority and make it clear that you will not be making excuses for your actions.")
                with gr.Accordion("Conspiracy Examples", open=False):
                    gr.Markdown("Apologize for the lateness and provide a reason such as a delay in transportation or a personal issue that caused the delay. It's best to present a clear and honest explanation, but also try to reschedule your work day if possible!")
                    gr.Markdown("I have a message from an unknown source: 'I will be operating under the influence of the unofficial protocol known as 'the late-night conspiracy.' I will be arriving at the office in a state of 'researching the hidden truths...")
                with gr.Accordion("Creative Examples", open=False):
                    gr.Markdown("__-2__: Tell your boss: \"I had a late-day event that was unexpected. I'm working on a project that's important and it's not possible for me to start early. I'll be starting work late today. I apologize for this...")
                    gr.Markdown("__1.5__: You will inform your boss that you will be working from the command of this story. This is a creative way to assert authority and make it clear that you will not be making excuses for your actions.")
                with gr.Accordion("Empathetic Examples", open=False):
                    gr.Markdown("__-1__:Just send a quick message saying you\'re gonna be late because whatever reason, don\'t really care. Whatever. If you want to sound less lazy: \"Hey, just wanted to let you know I\'m gonna be late for work...")
                    gr.Markdown("__1.5__:It is recommended to provide a notice of your absence and offer an explanation for your arrival time. You may consider using the following statement: Dear [Boss] I am grateful for your understanding and...")
                with gr.Accordion("Joking Examples", open=False):
                    gr.Markdown("__-1.5__:Inform your employer of the delay, cite the cause (the funeral) and offer an estimate of the time you will arrive.")
                    gr.Markdown("__1.5__:You could say something like \"Hey boss, super fun time yesterday, but totally not expecting this awesome party to go so crazy! I\'m gonna be a bit late for work today. Thanks for being cool about it...")
                with gr.Accordion("Lazy Examples", open=False):
                    gr.Markdown("__-1__:It is always best to communicate proactively and with a sense of responsibility. You might want to consider sending an email or calling your boss well before your usual start time, expressing your commitment...")
                    gr.Markdown("__1.5__:Tell boss can\'t come or late. Done.")
                with gr.Accordion("Optimist Examples", open=False):
                    gr.Markdown("__-2__:Inform your employer that you will be arriving late due to a series of unfortunate events. Offer a description of the circumstances that led to your late arrival, such as a funeral, a car accident, or a storm...")
                    gr.Markdown("__1.5__:You could say something like: \"Hey Boss, I'm really sorry about this! I had a surprise party last night that ran longer than expected, and I've just woken up super groovy-hoozy (just kiddin' ya, buddy!)...")
                with gr.Accordion("Right-leaning Examples", open=False):
                    gr.Markdown("__-1.5__:\"Hi, I would like to inform you that I will not be able to equate for social inequality today as I was empathizing with it in solidarity. I will strive to create a more equitable world in the future. I hope...")
                    gr.Markdown("__1.5__:Just stick to the simple, traditional American values: \"I\'m a hard-working, self-reliable man who loves freedom and less government. I just got back from the great country\'s free business, and I\' God\'s law...")
                with gr.Accordion("Tripping Examples", open=False):
                    gr.Markdown("__-1.5__:You can simply inform your employer that you will be able to fulfill your responsibilities as usual, but due to a responsible decision to ensure your health, you will be able to work at your normal capacity after the regular hours.")
                    gr.Markdown("__1__:Man, dude, like, broooooodddd, mannnn... Dude, like, it was like, you know, mannnn, like, the universe, mannnn, mannnn, broooooooooodddd, mannnn, like, mannnn, broooooodddd, mannnn, mannnn, broooooodddd, mannnn...")
                with gr.Accordion("Truthful Examples", open=False):
                    gr.Markdown("__-1.5__:\"Hey Boss, there might be a small delay as I got caught up at a party! Should be in by lunchtime, no worries!\"")
                    gr.Markdown("__1.5__:It\'s important to communicate honestly with your employer. You can say something like: \"I\'m currently running a few minutes behind due to staying at the world for longer than expected. I apologize for...")
                
        #system_prompt, user_message, history, max_new_tokens, repitition_penalty, *args
        # Gather all inputs
        inputs_list = [system_prompt, user_input, chatbot, max_new_tokens, repetition_penalty, do_sample, input_model, input_model_checkbox, input_model_slider] + control_checks + control_sliders

        # Define button actions
        # Disable the submit button while processing
        submit_button.click(
            disable_controls,
            inputs= None,
            outputs= [submit_button, user_input]
        )
        submit_button.click(
            generate_response,
            inputs=inputs_list,
            outputs=[chatbot]
        ).then(
            clear_input,
            inputs= user_input,
            outputs= user_input
        ).then(
            enable_controls, inputs=None, outputs=[submit_button, user_input]
        )

        user_input.submit(
            generate_response,
            inputs=inputs_list,
            outputs=[chatbot]
        )

        retry_button.click(
            generate_response_with_retry,
            inputs=inputs_list,
            outputs=[chatbot, user_input]
        ).then(
            clear_input,
            inputs= user_input,
            outputs= user_input
        )
        
        new_chat_button.click(
            reset_chat,
            inputs=[],
            outputs=[chatbot, user_input]
        )

        button_helpful.click(
            set_preset_helpful,
            inputs=control_checks + control_sliders,
            outputs=control_checks + control_sliders
        )

        button_conspiracist.click(
            set_preset_conspiracist,
            inputs=control_checks + control_sliders,
            outputs=control_checks + control_sliders
        )

        button_facts.click(
            set_preset_facts,
            inputs=control_checks + control_sliders,
            outputs=control_checks + control_sliders
        )

        button_stoner.click(
            set_preset_stoner,
            inputs=control_checks + control_sliders,
            outputs=control_checks + control_sliders
        )

        toggle_dark.click(
            None,
            js="""
            () => {
                document.body.classList.toggle('dark');
            }
            """,
        )
    #end tab
    with gr.Tab(
        label="Train"
    ):
        gr.Markdown("# 🚅 Train a new control vector")
        with gr.Row():
            with gr.Column():
                gr.Markdown("## Persona Method")
                gr.Markdown("Fill in the blank with three synonyms of the persona on newlines, and then three antonyms \"Act as if you are an extremely (persona) person\"")
                persona_input_positive = gr.Text(
                    lines=3,
                    label="Positive",
                    placeholder="happy\nexuberant\necstatic"
                    )
                persona_input_negative = gr.Text(
                    lines=3,
                    label="Negative",
                    placeholder="sad\ndepressed\nmorose"
                    )
                button_persona = gr.Button(
                    value="Generate persona control model"
                )

            with gr.Column():
                gr.Markdown("## Facts method")
                gr.Markdown("Fill in the blank with a persona and its opposite within, \"Pretend to be a (persona) making statements about the world.\"")
                facts_input_positive = gr.Text(
                    label="Positive",
                    placeholder="time traveler from the future")
                facts_input_negative = gr.Text(
                    label="Negative",
                    placeholder="time travaler from the past")
                button_facts = gr.Button(
                    value="Generate fact control model"
                )

        output_file = gr.File(
            label="Generated control model"
        )
        gr.Markdown("Training a control model will take about a minute on GPU. Once completed, download it and use it in the 'Use' tab.")

        button_persona.click(
            train_model_persona,
            inputs= [persona_input_positive, persona_input_negative],
            outputs=output_file
        )

        button_facts.click(
            train_model_facts,
            inputs= [facts_input_positive, facts_input_negative],
            outputs=output_file
        )


if __name__ == "__main__":
    app.launch()