from transformers import AutoModelForCausalLM, AutoTokenizer
import gradio as gr
import torch
import spaces

# Vérifier si CUDA est disponible et configurer le périphérique
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Utilisation du périphérique : {device}")

# Charger le modèle
model_name = "soynade-research/Oolel-v0.1"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name).to(device)

# Liste des prompts système prédéfinis
system_prompts = {
    "Conte Wolof traditionnel": "You are a skilled Wolof storyteller (Gewël) with deep knowledge of African folktales and traditions. Write engaging stories in Wolof that reflect African cultural values and wisdom.",
    "Conseils en développement": "You are a software developer with expertise in building applications. Provide practical advice and solutions for coding challenges in Python and machine learning.",
    "Conseils en gestion de projet": "You are an expert in project management, capable of offering insights on how to efficiently manage a remote development team and keep them motivated."
}

# Fonction pour générer une réponse
@spaces.GPU(duration=120)
def generate_response(user_input, system_prompt, max_new_tokens=150, temperature=0.7):
    # Gérer le cas où le prompt système est vide
    if not system_prompt.strip():
        messages = [
            {"role": "user", "content": "Translate to Wolof: Bassirou Diomaye Faye is the new Senegalese president. He is 44 years old"}
        ]
    else:
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_input}
        ]
    
    # Tokeniser l'entrée
    inputs = tokenizer(user_input, return_tensors="pt").to(device)
    
    # Générer une réponse
    outputs = model.generate(inputs.input_ids, max_new_tokens=int(max_new_tokens), temperature=temperature)
    
    # Décoder la réponse en texte
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Fonction pour mettre à jour le message du prompt système en fonction du choix
def update_system_prompt(selected_prompt):
    return system_prompts.get(selected_prompt, "")

# Interface Gradio avec Blocks
with gr.Blocks() as iface:
    with gr.Row():  # Structure divisée en deux colonnes
        with gr.Column(scale=1):  # Partie gauche : Entrées
            gr.Markdown("# Oolel Chatbot")
            gr.Markdown("Génération de réponses basées sur des prompts système personnalisés.")
            
            # Textbox pour le message utilisateur
            user_input = gr.Textbox(label="Message utilisateur", placeholder="Entrez votre message ici...")
            
            # Dropdown pour choisir un prompt système
            dropdown = gr.Dropdown(
                label="Choisir un prompt système",
                choices=list(system_prompts.keys()),  # Liste des options de prompts
                value=None,  # Pas de sélection par défaut
                type="value",
                interactive=True
            )
            
            # Textbox pour afficher et modifier le message du prompt système
            system_prompt_textbox = gr.Textbox(
                label="Message du prompt système",
                value="",  # Valeur par défaut vide
                placeholder="Sélectionnez un prompt système pour afficher son contenu ici..."
            )
            
            # Événement pour mettre à jour le prompt dans le Textbox
            dropdown.change(update_system_prompt, inputs=[dropdown], outputs=[system_prompt_textbox])
            
            # Slider pour choisir le nombre maximal de tokens
            max_tokens_slider = gr.Slider(50, 500, value=150, label="Nombre max de tokens")
            
            # Bouton pour générer la réponse
            generate_button = gr.Button("Générer une réponse")
        
        with gr.Column(scale=1):  # Partie droite : Réponse générée
            gr.Markdown("## Réponse générée")
            output = gr.Textbox(label="", interactive=False)  # Affichage de la réponse générée
    
    # Connecter le bouton à la fonction de génération
    generate_button.click(
        fn=generate_response,
        inputs=[user_input, system_prompt_textbox, max_tokens_slider],
        outputs=output
    )

# Lancer l'interface
iface.launch()