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, max_new_tokens=150, temperature=0.7):
    # Préparer l'entrée pour le modèle
    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"),
        gr.Slider(50, 500, value=150, label="Nombre max de tokens")
    ],
    outputs="text",
    title="Oolel Chatbot"
)

# Lancer l'interface
iface.launch()