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):
    # Créer la liste des messages avec le prompt système et l'entrée utilisateur
    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=max_new_tokens, temperature=temperature)
    
    # Décoder la réponse en texte
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Interface Gradio
iface = gr.Interface(
    fn=generate_response,
    inputs=[
        gr.Textbox(label="Message utilisateur", placeholder="Entrez votre message ici..."),
        # Dropdown pour choisir un prompt système, mais sans sélection par défaut
        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"
        ),
        # Textbox pour afficher et modifier le message du prompt système
        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..."
        ),
        gr.Slider(50, 500, value=150, label="Nombre max de tokens")
    ],
    outputs="text",
    title="Oolel Chatbot",
    description="Entrez votre message et choisissez un prompt pour générer des réponses adaptées dans divers domaines."
)

# Lancer l'interface
iface.launch()