import gradio as gr
from haystack.nodes import PromptNode

from utils import lemmatizer_func

def run_prompt(prompt, api_key, model_name, max_length):
    prompt_node = PromptNode(model_name_or_path=model_name, api_key=api_key, max_length=max_length)
    lemmatized_prompt = lemmatizer_func(prompt)
    response_plain = prompt_node(prompt)
    response_lemmatized = prompt_node(lemmatized_prompt)
    return response_plain[0][0], response_plain[1]["total_tokens"], response_lemmatized[0][0], response_lemmatized[1]["total_tokens"]

with gr.Blocks() as demo:
    with gr.Row():
        api_key = gr.Textbox(label="Enter your api key")
        model_name = gr.Dropdown(["text-davinci-003", "gpt-3.5-turbo", "gpt-4", "gpt-4-32k", "command", "command-light", "base", "base-light"], value="gpt-3.5-turbo", label="Choose your model!")
    with gr.Row():
        prompt = gr.TextArea(label="Prompt", value="Rachel has 17 apples. She gives 9 to Sarah. How many apples does Rachel have now?")
        gr.Examples(
        [
            ["I want you to act as a travel guide. I will write you my location and you will suggest a place to visit near my location. In some cases, I will also give you the type of places I will visit. You will also suggest me places of similar type that are close to my first location. My first suggestion request is \"I am in Italy and I want to visit only museums.\""],
            ["What's the Everett interpretation of quantum mechanics?"],
            ["Give me a list of the top 10 dive sites you would recommend around the world."],
            ["Can you tell me more about deep-water soloing?"],
            ["Can you write a short tweet about the Apache 2.0 release of our latest AI model, Falcon LLM?"],
        ],
        inputs=prompt,
        label="Click on any example and press Enter in the input textbox!",
    )
    max_length = gr.Slider(100, 500, value=100, step=10, label="Max Length", info="Choose between 100 and 500")    
    submit_btn = gr.Button("Submit")
    with gr.Row():
        prompt_response = gr.TextArea(label="Answer", show_copy_button=True)
        token_count_plain = gr.Number(label="Plain Text Token Count")
    with gr.Row():
        lemmatized_prompt_response = gr.TextArea(label="Lemmatized Answer", show_copy_button=True)
        token_count_lemmatized = gr.Number(label="Lemmatized Text Token Count")
    submit_btn.click(fn=run_prompt, inputs=[prompt, api_key, model_name, max_length], outputs=[prompt_response, token_count_plain, lemmatized_prompt_response, token_count_lemmatized])

demo.launch()