Update app.py

app.py

@@ -2,39 +2,47 @@ import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
-# Tu modelo ya entrenado
-model_id = "OrangyDev/godot4-expert-ai"
+# 1. DEFINIR VARIABLES
+userxd = "OrangyDev"
+model_id = f"{userxd}/godot4-expert-ai"
 
-# Cargamos el tokenizer y el modelo
+# 2. CARGAR TOKENIZER Y MODELO
 tokenizer = AutoTokenizer.from_pretrained(model_id)
+# Usamos float16 y seleccionamos CPU o GPU automáticamente
 model = AutoModelForCausalLM.from_pretrained(
     model_id, 
-    torch_dtype=torch.float16, 
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, 
     device_map="auto"
 )
 
+# 3. FUNCIÓN DE CHAT
 def chat_godot(message, history):
-    # Formateamos la pregunta como en el entrenamiento
+    # Formato de prompt igual al del entrenamiento
     prompt = f"### User: {message}\n### Assistant:"
-    inputs = tokenizer(prompt, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
     
-    output = model.generate(
-        **inputs, 
-        max_new_tokens=100, 
-        temperature=0.7, 
-        repetition_penalty=1.2,
-        eos_token_id=tokenizer.eos_token_id
-    )
+    with torch.no_grad():
+        output = model.generate(
+            **inputs, 
+            max_new_tokens=150, 
+            temperature=0.7, 
+            do_sample=True,
+            repetition_penalty=1.2,
+            eos_token_id=tokenizer.eos_token_id
+        )
     
-    response = tokenizer.decode(output[0], skip_special_tokens=True)
-    return response.split("### Assistant:")[-1].strip()
+    full_text = tokenizer.decode(output[0], skip_special_tokens=True)
+    # Cortamos para quedarnos solo con la respuesta de la IA
+    response = full_text.split("### Assistant:")[-1].strip()
+    return response
 
-# Creamos la interfaz de chat
+# 4. INTERFAZ DE GRADIO
 demo = gr.ChatInterface(
     fn=chat_godot, 
     title="Godot 4 Expert AI",
     description=f"IA entrenada por {userxd} para resolver dudas de Godot 4 y GDScript.",
-    examples=["¿Quién es tu creador?", "¿Cómo muevo un personaje?", "¿Quién es Rafa Laguna?"]
+    examples=["¿Quién es tu creador?", "¿Quién es Rafa Laguna?", "¿Cómo muevo un personaje en Godot 4?"],
+    theme="soft"
 )
 
 if __name__ == "__main__":