app.py

import os
import gradio as gr
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

BASE_MODEL = "bigcode/santacoder"  
LORA_PATH = "./lora_output"        
DATASET_FILE = "codesearchnet_lora_dataset.json"   
MAX_TOKEN_LENGTH = 256             
NUM_SAMPLES_TO_PROCESS = 1000      
DEFAULT_EPOCHS = 10                

tokenizer = None
lora_model = None
tokenized_dataset = None
lora_generator = None

def prepare_codesearchnet():
    if os.path.exists(DATASET_FILE):
        print(f"Dataset '{DATASET_FILE}' ya existe.")
        return

    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}]')

        def format_for_lora(example):
            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)
        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}")
        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():
    global tokenizer, lora_model, tokenized_dataset

    prepare_codesearchnet()
    
    hf_token = os.environ.get("HF_TOKEN")
    if hf_token:
        login(token=hf_token)

    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

    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)
    
    print(f"Modelo LoRA preparado. Parámetros entrenables listos.")

    print(f"Cargando y tokenizando dataset: {DATASET_FILE}...")
    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 [],
        )
        print("Dataset tokenizado correctamente.")
    except Exception as e:
        tokenized_dataset = None
        print(f"Error al cargar o tokenizar el dataset. {e}")

def train_lora(epochs, batch_size, learning_rate):
    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."

    try:
        lora_generator = None 
        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,
        )

        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) 

        return f"Entrenamiento completado. Adaptadores LoRA guardados en **{LORA_PATH}**"
    except Exception as e:
        return f"Error durante el entrenamiento: {e}"

def generate_text(prompt_text):
    global lora_generator

    try:
        if lora_generator is None:
            base_model_gen = AutoModelForCausalLM.from_pretrained(BASE_MODEL, device_map="auto")
            
            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.")
                final_model = base_model_gen

            final_model.eval() 
            lora_generator = pipeline("text-generation", model=final_model, tokenizer=tokenizer)
            print("Modelo de inferencia listo.")

        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}"

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 (10 Épocas).")

    with gr.Tab("🧠 Entrenar (Manual)"):
        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")
        train_button = gr.Button("🚀 Iniciar Entrenamiento Manual")
        train_output = gr.Textbox(label="Resultado del Entrenamiento Manual")
        
        train_button.click(
            train_lora, 
            inputs=[epochs, batch_size, learning_rate], 
            outputs=train_output
        )

    with gr.Tab("✨ Probar modelo"):
        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 (SOLO CÓDIGO)", lines=10)
        generate_button.click(generate_text, inputs=prompt, outputs=output_box)

if __name__ == "__main__":
    setup_resources()
    
    print("\n=============================================")
    print(f"🤖 INICIANDO AUTO-ENTRENAMIENTO ({DEFAULT_EPOCHS} Épocas, 2 Batch Size) usando {NUM_SAMPLES_TO_PROCESS} ejemplos")
    print("=============================================")
    
    auto_train_result = train_lora(epochs=DEFAULT_EPOCHS, batch_size=2, learning_rate=5e-5) 
    
    print(f"\nFIN DEL AUTO-ENTRENAMIENTO: {auto_train_result}")
    
    print("\n=============================================")
    print("💻 LANZANDO INTERFAZ GRADIO")
    print("=============================================")
    demo.launch()