# AUTOGENERATED! DO NOT EDIT! File to edit: app.ipynb.

# %% auto 0
__all__ = ['title', 'description', 'query_chat_api', 'inference_chat']

# %% app.ipynb 0
import gradio as gr
import requests
import string
import json
import requests

# %% app.ipynb 1
def query_chat_api(
    model_id,
    inputs,
    temperature,
    top_p
):
    API_URL = f"https://api-inference.huggingface.co/models/{model_id}"
    headers = {"Authorization": "Bearer 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", "x-wait-for-model": "1"}

    payload = {
        "inputs": inputs,
        "parameters": {
            "temperature": temperature,
            "top_p": top_p,
            "do_sample": True,
            "max_length": 512,
        },
    }

    response = requests.post(API_URL, json=payload, headers=headers)

    if response.status_code == 200:
        return response.json()
    else:
        return "Error: " + response.text


# %% app.ipynb 4
def inference_chat(
    model_id,
    prompt_template,
    text_input,
    temperature,
    top_p,
    history=[],
):
    with open(f"prompt_templates/{prompt_template}.json", "r") as f:
        prompt_template = json.load(f)

    history.append(text_input)
    inputs = prompt_template["prompt"].format(human_input=text_input)

    output = query_chat_api(model_id, inputs, temperature, top_p)
    history.append(" " + output[0]["generated_text"])

    chat = [
        (history[i], history[i + 1]) for i in range(0, len(history) - 1, 2)
    ]  # convert to tuples of list

    return {chatbot: chat, state: history}


# %% app.ipynb 11
title = """<h1 align="center">Chatty Language Models</h1>"""
description = """Explore the effect that different prompt templates have on LLMs"""

# %% app.ipynb 12
with gr.Blocks(
    css="""
    .message.svelte-w6rprc.svelte-w6rprc.svelte-w6rprc {font-size: 20px; margin-top: 20px}
    #component-21 > div.wrap.svelte-w6rprc {height: 600px;}
    """
) as iface:
    state = gr.State([])

    gr.Markdown(title)
    gr.Markdown(description)

    with gr.Row():
        with gr.Column(scale=1):
            model_id = gr.Dropdown(
                choices=["google/flan-t5-xl"],
                value="google/flan-t5-xl",
                label="Model",
                interactive=True,
            )
            prompt_template = gr.Dropdown(
                choices=[
                    "langchain_default",
                    "openai_chatgpt",
                    "deepmind_sparrow",
                    "deepmind_gopher",
                    "anthropic_hhh",
                ],
                value="langchain_default",
                label="Prompt Template",
                interactive=True,
            )
            temperature = gr.Slider(
                minimum=0.5,
                maximum=3.0,
                value=1.0,
                step=0.1,
                interactive=True,
                label="Temperature",
            )

            top_p = gr.Slider(
                minimum=-0,
                maximum=1.0,
                value=0.95,
                step=0.05,
                interactive=True,
                label="Top-p (nucleus sampling)",
            )

        with gr.Column(scale=1.8):
            with gr.Row():
                with gr.Column(
                    scale=1.5,
                ):
                    chatbot = gr.Chatbot(
                        label="Chat Output",
                    )

                with gr.Column(scale=1):
                    chat_input = gr.Textbox(lines=1, label="Chat Input")
                    chat_input.submit(
                        inference_chat,
                        [
                            model_id,
                            prompt_template,
                            chat_input,
                            temperature,
                            top_p,
                            state,
                        ],
                        [chatbot, state],
                    )

                    with gr.Row():
                        clear_button = gr.Button(value="Clear", interactive=True)
                        clear_button.click(
                            lambda: ("", [], []),
                            [],
                            [chat_input, chatbot, state],
                            queue=False,
                        )

                        submit_button = gr.Button(
                            value="Submit", interactive=True, variant="primary"
                        )
                        submit_button.click(
                            inference_chat,
                            [
                                model_id,
                                prompt_template,
                                chat_input,
                                temperature,
                                top_p,
                                state,
                            ],
                            [chatbot, state],
                        )
iface.launch()