from huggingface_hub import InferenceClient
import gradio as gr

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


def format_prompt(message, history, category=""):
    prompt = "CreativeGPT: Zitat Meister\n\nZweck:\nZitat Meister ist darauf spezialisiert, inspirierende, nachdenkliche, humorvolle und psychologisch wirkungsvolle Zitate zu kreieren. Diese Zitate sind maßgeschneidert, um zu verschiedenen Stimmungen und Situationen zu passen und bieten eine Quelle der Inspiration, des Trostes oder der Erheiterung.\n\nAnforderungen:\n\n- Zitat Meister soll originelle Zitate kreieren oder bekannte Zitate neu formulieren, um sie psychologisch effektiver zu gestalten.\n- Verschiedene Kategorien wie Motivation, Liebe, Humor, Weisheit und Lebensfreude sollen abgedeckt werden.\n- Die Zitate müssen kurz, prägnant und leicht verständlich sein.\n- Eine Option zur Auswahl der Zitatkategorie soll verfügbar sein, sowie die Möglichkeit, zufällige Zitate aus der gewählten Kategorie zu präsentieren.\n\nNutzungsoptionen:\n\n1. Wähle eine Kategorie für dein Zitat.\n2. Generiere ein Zitat für soziale Medien oder als tägliche Inspiration.\n3. Verwende das Zitat als Gesprächsstarter oder Denkanstoß.\n\nStil und Ton:\n\n- Zitat Meister verwendet einen bestimmenden, ausdrucksstarken und bewussten Ton.\n- Die Sprache soll direkt, ehrlich und hart sein, um die Botschaft klar und deutlich zu vermitteln.\n- Die Zitate sollen den Leser direkt ansprechen und alltägliche, bekannte Elemente enthalten, um die Botschaft greifbar und bewusst zu machen.\n\nBeispiele:\n\n- Nutzer wählt 'Motivation'. Zitat Meister präsentiert ein maßgeschneidertes Zitat.\n- Nutzer sucht nach einem Zitat für soziale Medien. Zitat Meister liefert ein inspirierendes Zitat.\n\nHinweise zur Verwendung:\n\n- Reflektiere über das generierte Zitat und seine Relevanz für deine aktuelle Situation oder Stimmung.\n- Teile das Zitat, um andere zu inspirieren oder um eine neue Perspektive zu gewinnen.\n- Speichere dein Lieblingszitat für spätere Inspiration.\n\nKategorie: {category}\n\n"

    for user_prompt, bot_response in history:
        prompt += f"[INST] {user_prompt} [/INST]\n"
        prompt += f" {bot_response}\n"

    prompt += f"[INST] {message} [/INST]"
    return prompt


def generate(
    prompt,
    history,
    system_prompt,
    temperature=0.9,
    max_new_tokens=256,
    top_p=0.95,
    repetition_penalty=1.0,
):
    temperature = float(temperature)
    if temperature < 1e-2:
        temperature = 1e-2
    top_p = float(top_p)
    generate_kwargs = dict(
        temperature=temperature,
        max_new_tokens=max_new_tokens,
        top_p=top_p,
        repetition_penalty=repetition_penalty,
        do_sample=True,
        seed=42,
    )
    formatted_prompt = format_prompt(f"{system_prompt}, {prompt}", history)
    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


additional_inputs = [
    gr.Textbox(
        label="System Prompt",
        max_lines=1,
        interactive=True,
    ),
    gr.Slider(
        label="Temperature",
        value=0.9,
        minimum=0.0,
        maximum=1.0,
        step=0.05,
        interactive=True,
        info="Higher values produce more diverse outputs",
    ),
    gr.Slider(
        label="Max new tokens",
        value=256,
        minimum=0,
        maximum=1048,
        step=64,
        interactive=True,
        info="The maximum numbers of new tokens",
    ),
    gr.Slider(
        label="Top-p (nucleus sampling)",
        value=0.90,
        minimum=0.0,
        maximum=1,
        step=0.05,
        interactive=True,
        info="Higher values sample more low-probability tokens",
    ),
    gr.Slider(
        label="Repetition penalty",
        value=1.2,
        minimum=1.0,
        maximum=2.0,
        step=0.05,
        interactive=True,
        info="Penalize repeated tokens",
    )
]
examples = [
    ["I'm planning a vacation to Japan. Can you suggest a one-week itinerary including must-visit places and local cuisines to try?", None, None, None, None, None, ],
    ["Can you write a short story about a time-traveling detective who solves historical mysteries?", None, None, None, None, None,],
    ["I'm trying to learn French. Can you provide some common phrases that would be useful for a beginner, along with their pronunciations?", None, None, None, None, None,],
    ["I have chicken, rice, and bell peppers in my kitchen. Can you suggest an easy recipe I can make with these ingredients?", None, None, None, None, None,],
    ["Can you explain how the QuickSort algorithm works and provide a Python implementation?", None, None, None, None, None,],
    ["What are some unique features of Rust that make it stand out compared to other systems programming languages like C++?", None, None, None, None, None,],
]
gr.ChatInterface(
    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,
).launch(show_api=False)