update

app.py

@@ -1,23 +1,35 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import gradio as gr
+import torch
+
+# 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, device_map="auto")
+model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
 
 # Fonction pour générer une réponse
 def generate_response(user_input, max_new_tokens=150, temperature=0.7):
-    inputs = tokenizer(user_input, return_tensors="pt").to("cuda")
+    # 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")],
+    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()