Update app.py

app.py

@@ -20,6 +20,9 @@ DEFAULT_EPOCHS = 10
 GENERATION_LIMIT_TO_TRAIN = 5 
 AUTONOMOUS_EPOCHS = 3           
 
+# Nueva configuración para manejo de modelos grandes
+TEMP_OFFLOAD_FOLDER = "./temp_offload"
+
 # --- ESTADO GLOBAL Y THREADING ---
 tokenizer = None
 lora_model = None
@@ -27,10 +30,10 @@ tokenized_dataset = None
 lora_generator = None
 
 # Variables de estado
-version_number = 1.0 
+version_number = 1.1 # Asumimos que el último entrenamiento llegó a V1.1
 is_trained = os.path.exists(LORA_PATH) 
 generations_since_last_train = 0
-training_status_message = "Esperando la inicialización V1.0..."
+training_status_message = f"Modelo V{version_number:.1f} listo."
 
 # Lock para proteger las variables compartidas entre hilos (CRÍTICO para estabilidad)
 global_lock = threading.Lock() 
@@ -67,6 +70,9 @@ def setup_resources():
     """Configura el tokenizer, el modelo base y el adaptador LoRA."""
     global tokenizer, lora_model, tokenized_dataset
 
+    # Crear la carpeta de offload si no existe (CRÍTICO para el error actual)
+    os.makedirs(TEMP_OFFLOAD_FOLDER, exist_ok=True) 
+
     prepare_codesearchnet()
     
     hf_token = os.environ.get("HF_TOKEN")
@@ -74,7 +80,15 @@ def setup_resources():
         login(token=hf_token)
 
     tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
-    base_model = AutoModelForCausalLM.from_pretrained(BASE_MODEL, device_map="auto") 
+    
+    # Añadir safetensors=True y offload_folder
+    base_model = AutoModelForCausalLM.from_pretrained(
+        BASE_MODEL, 
+        device_map="auto",
+        offload_folder=TEMP_OFFLOAD_FOLDER, # SOLUCIÓN 1
+        low_cpu_mem_usage=True, 
+        trust_remote_code=True,
+    ) 
 
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
@@ -149,14 +163,13 @@ def autonomous_train_lora(epochs, batch_size, learning_rate):
         training_status_message = f"ERROR CRÍTICO durante el entrenamiento autónomo: {e}"
         print(f"[AUTÓNOMO] {training_status_message}")
 
-# --- FUNCIÓN DE GENERACIÓN (CORREGIDA PARA RETORNAR 2 VALORES) ---
+# --- FUNCIÓN DE GENERACIÓN (CORREGIDA PARA HOT SWAP ESTABLE) ---
 
 def generate_text(prompt_text):
     """Genera código y dispara el ciclo de reentrenamiento autónomo si es necesario."""
     global lora_generator, generations_since_last_train, is_trained, version_number, training_status_message
 
     if not is_trained:
-        # Si el entrenamiento V1.0 no ha terminado, retorna el mensaje de error y el estado actual
         return "ERROR: El modelo LoRA no ha sido entrenado. Por favor, espere mientras la IA se inicializa con el entrenamiento V1.0.", update_status()
 
 
@@ -164,15 +177,21 @@ def generate_text(prompt_text):
     if lora_generator is None:
         with global_lock:
             try:
-                # Recarga el modelo solo si está vacío
-                base_model_gen = AutoModelForCausalLM.from_pretrained(BASE_MODEL, device_map="auto")
+                # SOLUCIÓN CLAVE: Reintroducir offload_folder y low_cpu_mem_usage aquí también.
+                base_model_gen = AutoModelForCausalLM.from_pretrained(
+                    BASE_MODEL, 
+                    device_map="auto", 
+                    offload_folder=TEMP_OFFLOAD_FOLDER, # SOLUCIÓN 2
+                    low_cpu_mem_usage=True,
+                    trust_remote_code=True,
+                )
                 model_with_lora = PeftModel.from_pretrained(base_model_gen, LORA_PATH)
                 final_model = model_with_lora.merge_and_unload()
                 final_model.eval() 
                 lora_generator = pipeline("text-generation", model=final_model, tokenizer=tokenizer)
                 print(f"[HOT SWAP] 🔄 Modelo de inferencia V{version_number:.1f} recargado y listo.")
             except Exception as e:
-                 # Si la recarga falla, retorna un error
+                # Si la recarga falla, retorna un error
                 return f"Error al cargar el modelo V{version_number:.1f} para inferencia: {e}", update_status()
 
 
@@ -215,7 +234,7 @@ def generate_text(prompt_text):
                 
                 notification = f"\n\n--- [AUTONOMÍA] La IA ha iniciado el reentrenamiento V{current_version+0.1:.1f} para mejorar la traducción de tu diálogo. La próxima generación cargará la nueva versión. ---"
             
-        # CORRECCIÓN CLAVE: Retorna el código Y el estado actualizado
+        # Retorna el código Y el estado actualizado
         return completion + notification, update_status()
 
     except Exception as e:
@@ -226,10 +245,11 @@ def generate_text(prompt_text):
 
 def initialize_and_train_v1():
     """Ejecuta el entrenamiento inicial V1.0 de forma autónoma al iniciar."""
+    global version_number, is_trained, training_status_message
     if not is_trained:
         autonomous_train_lora(epochs=DEFAULT_EPOCHS, batch_size=2, learning_rate=5e-5)
     else:
-        global training_status_message
+        # Si ya está entrenado, actualiza la versión y el mensaje
         training_status_message = f"✅ Modelo V{version_number:.1f} ya entrenado. Listo."
         print(f"[INICIALIZACIÓN] {training_status_message}")
 
@@ -268,7 +288,6 @@ with gr.Blocks(title="AmorCoderAI - Aprendizaje Continuo") as demo:
         output_box = gr.Textbox(label="Código generado", lines=10)
         
         # Conexión del botón con la función principal
-        # IMPORTANTE: Ahora generate_text retorna DOS valores para coincidir con [output_box, version_and_status]
         generate_button.click(
             generate_text, 
             inputs=prompt, 
@@ -288,5 +307,5 @@ if __name__ == "__main__":
     initialization_thread.daemon = True
     initialization_thread.start()
         
-    print(f"\n💻 LANZANDO INTERFAZ GRADIO (El entrenamiento V1.0 se ejecuta en segundo plano)")
+    print(f"\n💻 LANZANDO INTERFAZ GRADIO (El entrenamiento V1.0/V1.1 se ejecuta en segundo plano)")
     demo.launch()