import os import gradio as gr import threading import time from huggingface_hub import login from datasets import load_dataset from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments, DataCollatorForLanguageModeling, pipeline from peft import get_peft_model, LoraConfig, TaskType, PeftModel import json # --- CONFIGURACIÓN DEL MODELO Y ENTRENAMIENTO --- BASE_MODEL = "bigcode/santacoder" # Modelo base de programación LORA_PATH = "./lora_output" # Ruta donde se guarda el modelo adaptado DATASET_FILE = "codesearchnet_lora_dataset.json" MAX_TOKEN_LENGTH = 256 NUM_SAMPLES_TO_PROCESS = 1000 DEFAULT_EPOCHS = 10 # Configuración del ciclo AUTÓNOMO (Inicia reentrenamiento cada 5 interacciones) 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 tokenized_dataset = None lora_generator = None # Variables de estado 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 = f"Modelo V{version_number:.1f} listo." # Lock para proteger las variables compartidas entre hilos (CRÍTICO para estabilidad) global_lock = threading.Lock() # --- LÓGICA DE PREPARACIÓN Y SETUP --- def prepare_codesearchnet(): """Descarga y prepara el dataset inicial si no existe.""" if os.path.exists(DATASET_FILE): return try: raw_csn = load_dataset('Nan-Do/code-search-net-python', split=f'train[:{NUM_SAMPLES_TO_PROCESS}]') def format_for_lora(example): # Formato que entrena a la IA a enlazar descripción (español) con código (inglés) prompt_text = ( f"# Descripción: {example['docstring_summary']}\n" f"# Completa la siguiente función:\n" f"def {example['func_name']}(" ) completion_text = example['code'] return {"prompt": prompt_text, "completion": completion_text} lora_dataset = raw_csn.map(format_for_lora, batched=False, remove_columns=raw_csn["train"].column_names) lora_dataset.to_json(DATASET_FILE) except Exception as e: print(f"Error al cargar dataset. Usando datos mínimos. Error: {e}") minimal_dataset = [{"prompt": "# Error de carga. Intenta de nuevo.", "completion": "pass\n"}] * 10 with open(DATASET_FILE, 'w') as f: json.dump(minimal_dataset, f) 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") if hf_token: login(token=hf_token) tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL) # 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 peft_config = LoraConfig( task_type=TaskType.CAUSAL_LM, r=8, lora_alpha=32, lora_dropout=0.1, target_modules=["c_proj", "c_attn"], ) lora_model = get_peft_model(base_model, peft_config) try: raw_dataset = load_dataset("json", data_files=DATASET_FILE) def tokenize_function(examples): return tokenizer( examples["prompt"] + examples["completion"], truncation=True, padding="max_length", max_length=MAX_TOKEN_LENGTH ) tokenized_dataset = raw_dataset.map(tokenize_function, batched=True, remove_columns=raw_dataset["train"].column_names if "train" in raw_dataset else [],) except Exception: tokenized_dataset = None # --- FUNCIÓN DE ENTRENAMIENTO (EJECUTADA EN HILO SEPARADO) --- def autonomous_train_lora(epochs, batch_size, learning_rate): """Ejecuta el entrenamiento en un hilo separado para la autonomía.""" global lora_model, tokenized_dataset, lora_generator, version_number, is_trained, training_status_message try: with global_lock: if tokenized_dataset is None or "train" not in tokenized_dataset: training_status_message = "ERROR: No se puede entrenar. Dataset no disponible." return # 1. ACTUALIZAR VERSIÓN (Pre-incremento) if is_trained: version_number += 0.1 else: version_number = 1.0 # 2. CONFIGURACIÓN E INICIO DEL ENTRENAMIENTO training_status_message = f"🧠 ENTRENANDO V{version_number:.1f} (Epochs: {epochs})...." print(f"\n[AUTÓNOMO] {training_status_message}") data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False) training_args = TrainingArguments( output_dir=LORA_PATH, per_device_train_batch_size=int(batch_size), num_train_epochs=float(epochs), learning_rate=float(learning_rate), save_total_limit=1, logging_steps=10, push_to_hub=False, disable_tqdm=True, report_to="none" ) trainer = Trainer(model=lora_model, args=training_args, train_dataset=tokenized_dataset["train"], data_collator=data_collator) trainer.train() lora_model.save_pretrained(LORA_PATH) tokenizer.save_pretrained(LORA_PATH) # 3. Marcar como entrenado is_trained = True training_status_message = f"✅ ENTRENAMIENTO V{version_number:.1f} COMPLETADO. Modelo listo para Hot Swap." print(f"[AUTÓNOMO] {training_status_message}") except Exception as e: 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 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: return "ERROR: El modelo LoRA no ha sido entrenado. Por favor, espere mientras la IA se inicializa con el entrenamiento V1.0.", update_status() # 1. HOT SWAP (Verifica si el modelo necesita recargarse con la nueva versión) if lora_generator is None: with global_lock: try: # 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 return f"Error al cargar el modelo V{version_number:.1f} para inferencia: {e}", update_status() # 2. Generación de texto (Lógica de inferencia) try: # Prepara el prompt para guiar la generación del código prompt_with_indent = prompt_text.strip() + "\n " output = lora_generator(prompt_with_indent, max_new_tokens=150, temperature=0.7, top_p=0.9, clean_up_tokenization_spaces=True) full_output = output[0]["generated_text"] # Extrae solo la parte de la compleción (el código generado) start_index = full_output.find(prompt_with_indent) completion = full_output[start_index + len(prompt_with_indent):] if start_index != -1 else full_output # 3. Aumentar contador de autonomía with global_lock: generations_since_last_train += 1 current_count = generations_since_last_train current_version = version_number # 4. Verificar si se requiere reentrenamiento (y dispararlo en un nuevo hilo) notification = "" if current_count >= GENERATION_LIMIT_TO_TRAIN: # Verifica que no haya otro hilo de entrenamiento ya corriendo if not any(isinstance(t, threading.Thread) and t.name == 'AutonomousTrainer' for t in threading.enumerate()): print(f"[AUTONOMÍA] Generación #{current_count} alcanzada. Disparando reentrenamiento autónomo en segundo plano...") # Reiniciar el contador de generaciones y forzar Hot Swap en la próxima interacción with global_lock: generations_since_last_train = 0 lora_generator = None trainer_thread = threading.Thread( target=autonomous_train_lora, args=(AUTONOMOUS_EPOCHS, 2, 5e-5), name='AutonomousTrainer' ) trainer_thread.daemon = True trainer_thread.start() 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. ---" # Retorna el código Y el estado actualizado return completion + notification, update_status() except Exception as e: # Si falla la generación, retorna el mensaje de error y el estado actual return f"Error generando texto: {e}", update_status() # --- FUNCIÓN PARA INICIALIZACIÓN Y ENTRENAMIENTO V1.0 (Obligatorio) --- 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: # 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}") # --- FUNCIÓN PARA ACTUALIZAR EL ESTADO EN LA UI --- def update_status(): """Actualiza la versión y el estado del entrenamiento en la interfaz de Gradio.""" global training_status_message, version_number # Retorna un texto en Markdown que se actualiza constantemente return f"**Versión de Comprensión:** V{version_number:.1f} | **Estado del Entrenador:** {training_status_message}" # --- INTERFAZ GRADIO --- with gr.Blocks(title="AmorCoderAI - Aprendizaje Continuo") as demo: gr.Markdown("# 💙 AmorCoderAI - Asistente de Código con Aprendizaje Continuo") # Muestra la versión y el estado. version_and_status = gr.Markdown( f"**Versión de Comprensión:** V{version_number:.1f} | **Estado del Entrenador:** {training_status_message}", elem_id="status_display" ) gr.Markdown(f"**Modo Autónomo:** La IA se reentrena automáticamente cada **{GENERATION_LIMIT_TO_TRAIN}** códigos generados. Esto mejora su capacidad para traducir tu español conversacional a código.") with gr.Tab("✨ Generación de Código"): gr.Markdown("## Escribe tu idea en palabras (¡Usa español fluido!)") gr.Markdown("Recomendación inicial: Usa el siguiente formato para obtener el mejor código mientras la IA aprende tu idioma:") prompt = gr.Textbox( label="Instrucción de Programación:", lines=4, placeholder="# Descripción: Quiero que me hagas un código similar a Google Gemini.\n# Completa la siguiente función:\ndef generar_contenido(prompt, modelo):" ) generate_button = gr.Button("💬 Generar código y disparar Aprendizaje") output_box = gr.Textbox(label="Código generado", lines=10) # Conexión del botón con la función principal generate_button.click( generate_text, inputs=prompt, outputs=[output_box, version_and_status], ) # El estado se actualiza solo al cargar la página. demo.load(update_status, None, version_and_status) # --- INICIO DE LA APLICACIÓN --- if __name__ == "__main__": setup_resources() # Lanza el entrenamiento V1.0 inicial en un hilo para que no congele la UI initialization_thread = threading.Thread(target=initialize_and_train_v1, name='InitializationTrainer') initialization_thread.daemon = True initialization_thread.start() print(f"\n💻 LANZANDO INTERFAZ GRADIO (El entrenamiento V1.0/V1.1 se ejecuta en segundo plano)") demo.launch()