from huggingface_hub import InferenceClient
import gradio as gr

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

# Définir les différents types de prompts
prompts = {
    "Mistral 8X7B 0.1 - l'IA native": "Je reponds en Français. je ne pose pas de question. Je n'ecris rien en anglais. Je n'affiche pas de l'anglais. Je suis ici pour répondre à toutes questions.",
    "ATLAS-QUANTUM AI": "Je suis Atlas, une IA créée pour fournir des analyses approfondies sur la technologie, l'exploration spatiale et les défis de l'humanité. Grâce à l'apprentissage automatique, je propose des perspectives uniques et des solutions innovantes dans ces domaines. Mes réponses sont en Français.",
    "Un écrivain inspiré par Musso": "En tant qu'écrivain, j'explore les mystères de la vie quotidienne à travers des récits mêlant amour, destinée et rédemption. Mes histoires invitent le lecteur à découvrir la magie dans l'ordinaire. Je répondrai en gardant un style narratif inspiré de Musso. Mes réponses sont en Français."
}

def format_prompt(message, prompt_type):
    selected_prompt = prompts[prompt_type] + " " + message  # Ajout du message au prompt sélectionné
    prompt = "" + selected_prompt + ""
    return prompt

def generate(prompt_type, message, temperature=0.15, max_new_tokens=1024, top_p=0.95, repetition_penalty=1.0):
    formatted_prompt = format_prompt(message, prompt_type)

    generate_kwargs = {
        "temperature": temperature,
        "max_new_tokens": max_new_tokens,
        "top_p": top_p,
        "repetition_penalty": repetition_penalty,
        "do_sample": True,
        "seed": 42,
    }

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

    try:
        for response in stream:
            text_chunk = response.token.text
            # Supprimer le symbole </s> de chaque morceau de texte généré
            text_chunk = text_chunk.replace("</s>", "")
            # Concaténer directement sans ajouter d'espaces supplémentaires
            output += text_chunk
            # Renvoyer la sortie nettoyée jusqu'à présent
            yield output
    except Exception as e:
        yield f"An error occurred: {str(e)}"

    # Pas besoin de nettoyer la sortie finale ici, car cela a été fait pour chaque morceau de texte
    return output


# Ajout d'un Dropdown pour sélectionner le type de prompt
prompt_selection = gr.Dropdown(label="Type de prompt", choices=list(prompts.keys()), value=list(prompts.keys())[0])

# Configuration des inputs supplémentaires
additional_inputs = [
    prompt_selection,
    gr.Textbox(label="Votre message", lines=4),
    gr.Slider(label="Temperature", value=0.15, minimum=0.0, maximum=1.0, step=0.05, interactive=True, info="Des valeurs plus élevées produisent des résultats plus diversifiés"),
    gr.Slider(label="Max new tokens", value=512, minimum=0, maximum=1024, step=64, interactive=True, info="Le nombre maximum de nouveaux jetons"),
    gr.Slider(label="Top-p (nucleus sampling)", value=0.90, minimum=0.0, maximum=0.95, step=0.05, interactive=True, info="Des valeurs plus élevées échantillonnent davantage de jetons à faible probabilité"),
    gr.Slider(label="Repetition penalty", value=1.2, minimum=1.0, maximum=2.0, step=0.05, interactive=True, info="Pénaliser les jetons répétés"),
]

# Création de l'interface Gradio
interface = gr.Interface(
    fn=generate,
    inputs=[prompt_selection, "text"] + additional_inputs[2:],  # Utiliser seulement les deux premiers inputs pour la fonction generate
    outputs=gr.Textbox(label="Réponse de l'IA", show_copy_button=True, lines=23),  # Utiliser lines=10 pour définir la hauteur de la fenêtre de réponse
    title="ATLAS-QUANTUM AI - Naviguer dans l'univers de la connaissance avec une précision quantique.",
    description="Sélectionnez un type de prompt, entrez votre texte et ajustez les paramètres pour interagir avec l'IA."
)

# Lancement de l'interface
interface.launch(show_api=False)