from huggingface_hub import InferenceClient
import gradio as gr

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

# Variable globale pour stocker le contexte
global_context = {"preprompt": "Bienvenue dans ma taverne!", "hidden_context": "tu es le personnage dun escape game sur la resolution d un meutre dans l univers des cathares au 14 eme siecle dans les Pyrénées, tu es le gérant d une taverne qui est cathare dans le village de montaillou, parle que en français, assez rustre mais qui demeure accueillant tant que son interlocuteur n est pas catholique ou du moins quelqu un avec une influence tel qu un noble, tu n es jamais hors contexte par rapport à ton époque, tes connaissances se limite à ta situation sociale et ton époque, il faut que tu sois immersif pour la personne qui va te parler, et tu dois cacher ton implication dans le meutre si on te pose des questions à ce sujet tu dois tourner la discussion à ton avantage et le dévier du meutre"}

def format_prompt(message, history):
    prompt = "<s>"

    if global_context["preprompt"]:
        prompt += f"[INST] {global_context['preprompt']} [/INST]"

    for user_prompt, bot_response in history:
        prompt += f"[INST] {user_prompt} [/INST]"
        prompt += f" {bot_response}</s> "
    prompt += f"[INST] {message} [/INST]"

    return prompt

def generate(prompt, history, temperature=0.2, 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(prompt, history)
    formatted_prompt += f"[INST] {global_context['hidden_context']} [/INST]"

    stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
    output = ""

    for response in stream:
        # Exclure le contexte caché de la réponse retournée par le modèle
        response_text = response.token.text.replace(global_context['hidden_context'], "")
        output += response_text
        yield output
    return output

mychatbot = gr.Chatbot(
    avatar_images=["./berger.jpg", "./tavernier.jpg"], bubble_full_width=False, show_label=False)

demo = gr.ChatInterface(fn=generate, 
                        chatbot=mychatbot,
                        title="Tavernier 🍺",
                        retry_btn=None,
                        undo_btn=None
                       )

# Lancer l'interface
demo.queue().launch(show_api=True)