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)

# 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..."),
        gr.Textbox(label="Prompt système", default="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.", placeholder="Entrez un prompt système..."),
        gr.Slider(50, 500, value=150, label="Nombre max de tokens")
    ],
    outputs="text",
    title="Oolel Chatbot",
    description="Entrez votre message et définissez un prompt pour générer des réponses dans l'esprit des contes africains en Wolof."
)

# Lancer l'interface
iface.launch()