Update app.py

app.py

@@ -6,35 +6,24 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingA
 from peft import get_peft_model, LoraConfig, TaskType, PeftModel
 import json
 
-# ============================================================
-# ⚙️ CONFIGURACIÓN GLOBAL
-# ============================================================
-# Modelo base para generación de código
 BASE_MODEL = "bigcode/santacoder"  
-LORA_PATH = "./lora_output"        # Directorio para guardar los adaptadores LoRA
-
-# Nombre del archivo donde se guardará el dataset procesado
+LORA_PATH = "./lora_output"        
 DATASET_FILE = "codesearchnet_lora_dataset.json"   
-MAX_TOKEN_LENGTH = 256           # Longitud de secuencia uniforme
-NUM_SAMPLES_TO_PROCESS = 1000    
-DEFAULT_EPOCHS = 10 # <--- ¡ENTRENAMIENTO PROFUNDO!
+MAX_TOKEN_LENGTH = 256             
+NUM_SAMPLES_TO_PROCESS = 1000      
+DEFAULT_EPOCHS = 10                
 
-# Variables globales
 tokenizer = None
 lora_model = None
 tokenized_dataset = None
 lora_generator = None
 
-# ============================================================
-# 🚨 LÓGICA DE PRE-PROCESAMIENTO DE DATOS (INTEGRADA) 🚨
-# ============================================================
 def prepare_codesearchnet():
-    """Descarga, procesa y guarda el dataset CodeSearchNet si no existe."""
     if os.path.exists(DATASET_FILE):
-        print(f"✅ Dataset '{DATASET_FILE}' ya existe.")
+        print(f"Dataset '{DATASET_FILE}' ya existe.")
         return
 
-    print(f"🔄 Descargando y procesando CodeSearchNet ({NUM_SAMPLES_TO_PROCESS} muestras)...")
+    print(f"Descargando y procesando CodeSearchNet ({NUM_SAMPLES_TO_PROCESS} muestras)...")
     
     try:
         raw_csn = load_dataset('Nan-Do/code-search-net-python', split=f'train[:{NUM_SAMPLES_TO_PROCESS}]')
@@ -46,33 +35,23 @@ def prepare_codesearchnet():
                 f"def {example['func_name']}("
             )
             completion_text = example['code']
-            
-            return {
-                "prompt": prompt_text,
-                "completion": completion_text
-            }
+            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)
-        print(f"✅ Pre-procesamiento completado. {NUM_SAMPLES_TO_PROCESS} ejemplos guardados en '{DATASET_FILE}'.")
+        print(f"Pre-procesamiento completado. {NUM_SAMPLES_TO_PROCESS} ejemplos guardados en '{DATASET_FILE}'.")
 
     except Exception as e:
-        print(f"❌ Error CRÍTICO al descargar/procesar CodeSearchNet. Error: {e}")
+        print(f"Error CRÍTICO al descargar/procesar CodeSearchNet. 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)
 
-# ============================================================
-# 🔐 AUTENTICACIÓN Y PRE-CARGA DE RECURSOS (SINGLETON)
-# ============================================================
-
 def setup_resources():
-    """Carga y configura todos los recursos (modelo, tokenizer, dataset) una sola vez."""
     global tokenizer, lora_model, tokenized_dataset
 
     prepare_codesearchnet()
@@ -81,15 +60,13 @@ def setup_resources():
     if hf_token:
         login(token=hf_token)
 
-    # 1. Carga del Tokenizer y Modelo Base
-    print("\n🔄 Cargando modelo base y tokenizer...")
+    print("\nCargando modelo base y tokenizer...")
     tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
     base_model = AutoModelForCausalLM.from_pretrained(BASE_MODEL, device_map="auto") 
 
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
 
-    # 2. Configuración y Aplicación LoRA (PEFT)
     peft_config = LoraConfig(
         task_type=TaskType.CAUSAL_LM,
         r=8,
@@ -99,10 +76,9 @@ def setup_resources():
     )
     lora_model = get_peft_model(base_model, peft_config)
     
-    print(f"✅ Modelo LoRA preparado. Parámetros entrenables listos.")
+    print(f"Modelo LoRA preparado. Parámetros entrenables listos.")
 
-    # 3. Carga y Tokenización del Dataset
-    print(f"📚 Cargando y tokenizando dataset: {DATASET_FILE}...")
+    print(f"Cargando y tokenizando dataset: {DATASET_FILE}...")
     try:
         raw_dataset = load_dataset("json", data_files=DATASET_FILE)
         
@@ -119,21 +95,16 @@ def setup_resources():
             batched=True,
             remove_columns=raw_dataset["train"].column_names if "train" in raw_dataset else [],
         )
-        print("✅ Dataset tokenizado correctamente.")
+        print("Dataset tokenizado correctamente.")
     except Exception as e:
         tokenized_dataset = None
-        print(f"❌ Error al cargar o tokenizar el dataset. {e}")
+        print(f"Error al cargar o tokenizar el dataset. {e}")
 
-
-# ============================================================
-# ��� FUNCIÓN DE ENTRENAMIENTO
-# ============================================================
 def train_lora(epochs, batch_size, learning_rate):
-    """Ejecuta el entrenamiento del modelo LoRA."""
     global lora_model, tokenized_dataset, lora_generator
 
     if tokenized_dataset is None or "train" not in tokenized_dataset:
-        return f"❌ Error: El dataset no pudo cargarse o está vacío. No se puede entrenar."
+        return f"Error: El dataset no pudo cargarse o está vacío. No se puede entrenar."
 
     try:
         lora_generator = None 
@@ -157,19 +128,14 @@ def train_lora(epochs, batch_size, learning_rate):
         )
 
         trainer.train()
-        
         lora_model.save_pretrained(LORA_PATH) 
         tokenizer.save_pretrained(LORA_PATH) 
 
-        return f"✅ Entrenamiento completado. Adaptadores LoRA guardados en **{LORA_PATH}**"
+        return f"Entrenamiento completado. Adaptadores LoRA guardados en **{LORA_PATH}**"
     except Exception as e:
-        return f"❌ Error durante el entrenamiento: {e}"
+        return f"Error durante el entrenamiento: {e}"
 
-# ============================================================
-# 🤖 FUNCIÓN DE GENERACIÓN (INFERENCIA)
-# ============================================================
 def generate_text(prompt_text):
-    """Genera texto usando el modelo base + adaptadores LoRA."""
     global lora_generator
 
     try:
@@ -179,31 +145,44 @@ def generate_text(prompt_text):
             if os.path.exists(LORA_PATH):
                 print("Cargando adaptadores LoRA...")
                 model_with_lora = PeftModel.from_pretrained(base_model_gen, LORA_PATH)
+                final_model = model_with_lora.merge_and_unload()
             else:
                 print("No se encontraron adaptadores LoRA. Usando modelo base.")
-                model_with_lora = base_model_gen
+                final_model = base_model_gen
 
-            final_model = model_with_lora.merge_and_unload()
             final_model.eval() 
-            
             lora_generator = pipeline("text-generation", model=final_model, tokenizer=tokenizer)
             print("Modelo de inferencia listo.")
 
-        output = lora_generator(prompt_text, max_new_tokens=150, temperature=0.7, top_p=0.9)
-        return output[0]["generated_text"]
+        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"]
+        
+        start_index = full_output.find(prompt_with_indent)
+        if start_index != -1:
+            completion = full_output[start_index + len(prompt_with_indent):]
+        else:
+            completion = full_output
+            
+        return completion
         
     except Exception as e:
-        return f"❌ Error generando texto (Asegúrate de que el modelo base/LoRA esté cargado): {e}"
+        return f"Error generando texto (Asegúrate de que el modelo base/LoRA esté cargado): {e}"
 
-# ============================================================
-# 💻 INTERFAZ GRADIO
-# ============================================================
 with gr.Blocks(title="AmorCoderAI - LoRA") as demo:
     gr.Markdown("# 💙 AmorCoderAI - Entrenamiento y Pruebas LoRA")
-    gr.Markdown(f"Modelo base: `{BASE_MODEL}`. Usando **{NUM_SAMPLES_TO_PROCESS}** ejemplos de CodeSearchNet.")
+    gr.Markdown(f"Modelo base: `{BASE_MODEL}`. Usando **{NUM_SAMPLES_TO_PROCESS}** ejemplos de CodeSearchNet (10 Épocas).")
 
     with gr.Tab("🧠 Entrenar (Manual)"):
-        gr.Markdown(f"--- **¡CUIDADO!** El auto-entrenamiento usará {DEFAULT_EPOCHS} épocas para aprender la sintaxis. ---")
+        gr.Markdown(f"--- ¡CUIDADO! El auto-entrenamiento usará {DEFAULT_EPOCHS} épocas para aprender la sintaxis. ---")
         epochs = gr.Number(value=DEFAULT_EPOCHS, label="Épocas", precision=0)
         batch_size = gr.Number(value=2, label="Tamaño de lote (ajusta según tu VRAM)", precision=0)
         learning_rate = gr.Number(value=5e-5, label="Tasa de aprendizaje")
@@ -217,14 +196,14 @@ with gr.Blocks(title="AmorCoderAI - LoRA") as demo:
         )
 
     with gr.Tab("✨ Probar modelo"):
-        prompt = gr.Textbox(label="Escribe código (ej: 'def fibonacci(n):')", lines=4)
+        prompt = gr.Textbox(
+            label="Escribe código (ej: # Descripción: Calcula el factorial de N. \n# Completa la siguiente función:\ndef factorial(n):)", 
+            lines=4
+        )
         generate_button = gr.Button("💬 Generar código")
-        output_box = gr.Textbox(label="Salida generada", lines=10)
+        output_box = gr.Textbox(label="Salida generada (SOLO CÓDIGO)", lines=10)
         generate_button.click(generate_text, inputs=prompt, outputs=output_box)
 
-# ============================================================
-# 🚀 LANZAR APP Y AUTO-ENTRENAMIENTO
-# ============================================================
 if __name__ == "__main__":
     setup_resources()