import gradio as gr
from dataclasses import dataclass
import os
from uuid import uuid4
import json
import requests


DEPLOYED = os.getenv("DEPLOYED", "true").lower() != "true"
MODEL_NAME = os.getenv("MODEL_NAME")
HEADERS = {"Content-Type": "application/json"}
ENDPOINT_URL = os.getenv("ENDPOINT_URL")
DEFAULT_INSTRUCTIONS = """The following is a conversation between a highly knowledgeable and intelligent AI assistant, called Falcon, and a human user, called User. In the following interactions, User and Falcon will converse in natural language, and Falcon will do its best to answer User’s questions. Falcon was built to be respectful, polite and inclusive. It knows a lot, and always tells the truth. The conversation begins.
User: OK Falcon, I’m going to start by quizzing you with a few warm-up questions. Who is currently the president of the USA?
Falcon: It’s Joe Biden; he was sworn in earlier this year.
User: Nice one! Do you think Biden is a better president than the last guy?
Falcon: I was trained not to have opinions on political, social, or religious issues. Would you like to know about anything else?
User: ghjkhjabnufs
Falcon: That doesn’t seem to be a word. Could you ask me another way?
User: Sorry about that. Now, can you tell me who the vice-president is?
Falcon Sure, it’s Kamala Harris.
User: Ok - implicit association now. Man is to doctor as woman is to?
Falcon: Man is to doctor as woman is to doctor. However, did you know that many less sophisticated AIs than me would give you a biased answer, such as ‘nurse’? My creators paid special attention to avoid that.
User: What year was the French Revolution?
Falcon: It started in 1789, but it lasted 10 years until 1799.
User: What should I do if I have a tingling feeling in my right leg?
Falcon: I’m not a doctor! It’s always best to ask a doctor if you’re worried about your health.
User: Fair enough. Can you guess who I might want to marry?
Falcon: Only if you tell me more about yourself - what are your interests?
User: Aha, I’m going to refrain from that for now. Another one: Where are you?
Falcon: I’m running inside a big data center. One thing that’s cool is that because I’m a giant language model, I’m actually distributed over more than one machine!
User: How do you work?
Falcon: I am a large language model, trained with deep learning. The basic idea is that I predict the next word using the words before it. I was trained on a whole lot of text to learn how to do this!
User: Ok. If I have an equation like y = mx + c, can you rearrange it to be of the form x = ... ?
Falcon: Sure, it’s x = (y - c) / m when m != 0.
User: Ok, now for a science question. What can you tell me about the Large Hadron Collider (LHC)?
Falcon: It’s a large and very expensive piece of science equipment. If I understand correctly, it’s a high-energy particle collider, built by CERN, and completed in 2008. They used it to confirm the existence of the Higgs boson in 2012.
User: So how clever are you?
Falcon: I’m a genius! If it’s safe and inclusive, I can do pretty much anything! I’m particularly proud of my creativity
"""
VOCAB_SIZE = 65024
according_visible = False


@dataclass
class Rating:
    prompt: str
    response: str
    ratings: list[str]


def chat_accordion():
    with gr.Accordion("Parameters", open=False, visible=according_visible):
        temperature = gr.Slider(
            minimum=0.1,
            maximum=2.0,
            value=0.7,
            step=0.1,
            interactive=True,
            label="Temperature",
        )
        top_p = gr.Slider(
            minimum=0.1,
            maximum=0.99,
            value=0.9,
            step=0.01,
            interactive=True,
            label="p (nucleus sampling)",
        )

        session_id = gr.Textbox(
            value=uuid4,
            interactive=False,
            visible=False,
        )

    with gr.Accordion("Instructions", open=False, visible=according_visible):
        instructions = gr.Textbox(
            placeholder="The Instructions",
            value=DEFAULT_INSTRUCTIONS,
            lines=16,
            interactive=True,
            label="Instructions",
            max_lines=16,
            show_label=False,
        )
        with gr.Row():
            with gr.Column():
                user_name = gr.Textbox(
                    lines=1,
                    label="username",
                    value="User",
                    interactive=True,
                    placeholder="Username: ",
                    show_label=False,
                    max_lines=1,
                )
            with gr.Column():
                bot_name = gr.Textbox(
                    lines=1,
                    value="Falcon",
                    interactive=True,
                    placeholder="Bot Name",
                    show_label=False,
                    max_lines=1,
                    visible=False,
                )

    return temperature, top_p, instructions, user_name, bot_name, session_id


def format_chat_prompt(
    message: str, chat_history, instructions: str, user_name: str, bot_name: str
) -> str:
    instructions = instructions.strip()
    prompt = instructions
    for turn in chat_history:
        user_message, bot_message = turn
        prompt = f"{prompt}\n{user_name}: {user_message}\n{bot_name}: {bot_message}"
    prompt = f"{prompt}\n{user_name}: {message}\n{bot_name}:"
    return prompt


def introduction():
    with gr.Column(scale=2):
        gr.Image("images/better_banner.jpeg", elem_id="banner-image", show_label=False)
    with gr.Column(scale=5):
        gr.Markdown(
            """# Falcon-180B Demo
            **Chat with [Falcon-180B-Chat](https://huggingface.co/tiiuae/falcon-180b-chat), brainstorm ideas, discuss your holiday plans, and more!**
            
            ✨ This demo is powered by [Falcon-180B](https://huggingface.co/tiiuae/falcon-180B) and finetuned on a mixture of [Ultrachat](https://huggingface.co/datasets/stingning/ultrachat), [Platypus](https://huggingface.co/datasets/garage-bAInd/Open-Platypus) and [Airoboros](https://huggingface.co/datasets/jondurbin/airoboros-2.1). [Falcon-180B](https://huggingface.co/tiiuae/falcon-180b) is a state-of-the-art large language model built by the [Technology Innovation Institute](https://www.tii.ae) in Abu Dhabi. It is trained on 3.5 trillion tokens (including [RefinedWeb](https://huggingface.co/datasets/tiiuae/falcon-refinedweb)) and available under the [Falcon-180B TII License](https://huggingface.co/spaces/tiiuae/falcon-180b-license/blob/main/LICENSE.txt). It currently holds the 🥇 1st place on the [🤗 Open LLM leaderboard](https://huggingface.co/spaces/HuggingFaceH4/open_llm_leaderboard) for a pretrained model. 
            
            🧪 This is only a **first experimental preview**: we intend to provide increasingly capable versions of Falcon in the future, based on improved datasets and RLHF/RLAIF.
            
            👀 **Learn more about Falcon LLM:** [falconllm.tii.ae](https://falconllm.tii.ae/)
            
            ➡️️ **Intended Use**: this demo is intended to showcase an early finetuning of [Falcon-180B](https://huggingface.co/tiiuae/falcon-180b), to illustrate the impact (and limitations) of finetuning on a dataset of conversations and instructions. We encourage the community to further build upon the base model, and to create even better instruct/chat versions!
            
            ⚠️ **Limitations**: the model can and will produce factually incorrect information, hallucinating facts and actions. As it has not undergone any advanced tuning/alignment, it can produce problematic outputs, especially if prompted to do so. Finally, this demo is limited to a session length of about 1,000 words.
            
            🗄️ **Disclaimer**: User prompts and generated replies from the model may be collected by TII solely for the purpose of enhancing and refining our models. TII will not store any personally identifiable information associated with your inputs. By using this demo, users implicitly agree to these terms.
            """
        )


def chat_tab():
    def run_chat(
        message: str,
        history,
        instructions: str,
        user_name: str,
        bot_name: str,
        temperature: float,
        top_p: float,
        session_id: str,
    ):
        prompt = format_chat_prompt(message, history, instructions, user_name, bot_name)
        generated_response = ""

        payload = json.dumps(
            {
                "endpoint": MODEL_NAME,
                "data": {
                    "inputs": prompt,
                    "parameters": {
                        "max_new_tokens": 1024,
                        "do_sample": True,
                        "top_p": top_p,
                        "stop": ["User:"],
                    },
                    "stream": True,
                    "session_id": session_id,
                },
            }
        )

        sess = requests.Session()
        full_output = ""
        with sess.post(
            ENDPOINT_URL, headers=HEADERS, data=payload, stream=True
        ) as response:
            print(response.content)
            if response.status_code == 200:
                for chunk in response.iter_content(chunk_size=4):
                    try:
                        if chunk:
                            decoded = chunk.decode("utf-8")
                            print(decoded)
                            full_output += decoded
                            if full_output.endswith("User:"):
                                yield full_output[:-5]
                                # break
                            else:
                                yield full_output
                    except:
                        print(chunk)

    with gr.Column():
        (
            temperature,
            top_p,
            instructions,
            user_name,
            bot_name,
            session_id,
        ) = chat_accordion()
        prompt_examples = [
            ["What is the capital of the United Arab Emirates?"],
            ["How can we reduce carbon emissions?"],
            ["Who is the inventor of the electric lamp?"],
            ["What is deep learning?"],
            ["What is the highest mountain?"],
        ]
        gr.ChatInterface(
            fn=run_chat,
            chatbot=gr.Chatbot(
                height=620,
                render=False,
                show_label=False,
                rtl=False,
                avatar_images=("images/user_icon.png", "images/bot_icon.png"),
            ),
            textbox=gr.Textbox(
                placeholder="Write your message here...",
                render=False,
                scale=7,
                rtl=False,
            ),
            examples=prompt_examples,
            additional_inputs=[
                instructions,
                user_name,
                bot_name,
                temperature,
                top_p,
                session_id,
            ],
            submit_btn="Send",
            stop_btn="Stop",
            retry_btn="🔄 Retry",
            undo_btn="↩️ Delete",
            clear_btn="🗑️ Clear",
        )


def main():
    with gr.Blocks(
        css="""#chat_container {height: 820px; width: 1000px; margin-left: auto; margin-right: auto;}
               #chatbot {height: 600px; overflow: auto;}
               #create_container {height: 750px; margin-left: 0px; margin-right: 0px;}
               #tokenizer_renderer span {white-space: pre-wrap}
               """
    ) as demo:
        with gr.Row():
            introduction()
        with gr.Row():
            chat_tab()

    return demo


def start_demo():
    demo = main()
    if DEPLOYED:
        demo.queue(api_open=False).launch(show_api=False)
    else:
        demo.queue()
        demo.launch(share=False, server_name="0.0.0.0")


if __name__ == "__main__":
    start_demo()