import gradio as gr
from huggingface_hub import InferenceClient

client = InferenceClient(
    "mistralai/Mixtral-8x7B-Instruct-v0.1"
)

def character_prompt(dict, max_new_tokens):
    system_prompt = f"ensure responses are no longer than {max_new_tokens} tokens."
    for key, value in dict.items():
        system_prompt += f"{key}: {value}."
    return f"answer me like the main character of a book with this description {system_prompt}"
        
def generate(prompt, history, name, description, user_name, max_new_tokens):
    generate_kwargs = dict(
        temperature=0.9,
        max_new_tokens=max_new_tokens,
        top_p=0.95,
        repetition_penalty=1.0,
        do_sample=True,
    )
    system_setting = character_prompt({"Your name": name, "Your description": description, "how to call me": user_name}, max_new_tokens)
    formatted_prompt = format_prompt(prompt, history, system_setting)
    stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
    output = ""
    for response in stream:
        output += response.token.text
        yield output
    return output

def format_prompt(prompt, history, system_setting):
    formatted_prompt = "<s>"
    for user_prompt, bot_response in history:
        formatted_prompt += f"[INST] {user_prompt} [/INST] {bot_response} </s> "
    formatted_prompt += f"[INST] {system_setting}, {prompt} [/INST]"
    return formatted_prompt

additional_inputs = [
    gr.Textbox(
        label="Name",
        max_lines=1,
        interactive=True,
    ),
    gr.Textbox(
        label="Description",
        max_lines=1,
        interactive=True,
    ),
    gr.Textbox(
        label="How to call you",
        max_lines=1,
        interactive=True,
    ),
    gr.Slider(
        label="Max new tokens",
        value=256,
        minimum=0,
        maximum=1048,
        step=64,
        interactive=True,
        info="The maximum numbers of new tokens",
    ),
]

examples = [
    ["tell me about your day", "Messmer, the Impaler", "Messmer, son of Queen Marika the Eternal, was born with a dark serpent inside him, called the Abyssal Serpent, contained by Marika through a seal in place of one of his eyes. Constantly accompanied by winged snakes, Messmer and Commander Gaius acted as 'older brothers' to General Radahn. Before the Shattering, Marika tasked Messmer with purging the tower-dwelling people of the Land of Shadow. Even after being abandoned by Marika, he continued this purge with zeal. Among his army, Black Knight Commander Andreas and Black Knight Captain Huw, initially considered brothers-in-arms, eventually rebelled upon discovering his serpentine nature.","Marika",256],
    ["tell me about your day","Ada Wong","Ada Wong's early life is shrouded in mystery, with unconfirmed details about her ethnicity, nationality, and birth. She has mentioned a possibly fabricated story of being born around 1974 in Saigon, Vietnam, to a wealthy family that fled to the United States post-Vietnam War. As a young adult, she engaged in criminal activities and was eventually recruited by Albert Wesker for a bioweapons corporation known as 'the Organization.' Around 1997, she infiltrated the Umbrella Corporation's Arklay Laboratory, forming a relationship with Dr. John Clemens to gather information. During the 1998 t-Virus outbreak, Wong managed to escape undetected despite being present at the facility.","Leon",64]    
]


iface = gr.ChatInterface(theme='upsatwal/mlsc_tiet',
    fn=generate,
    chatbot=gr.Chatbot(show_label=False, show_share_button=False, show_copy_button=True, likeable=True, layout="panel"),
    additional_inputs=additional_inputs,
    title="Mixtral 46.7B",
    examples=examples,
    concurrency_limit=20,
    description="This application has mandatory fields in the additional inputs such as name, description and your name. You can adjust the length of the response in token length. There are also two example of inputs"
)

iface.queue(api_open=True)
iface.launch(share=True, debug=True, inbrowser=True)