prepare_codesearchnet.py

# prepare_codesearchnet.py

from datasets import load_dataset

# La ruta donde guardaremos el nuevo dataset para LoRA
OUTPUT_FILE = "codesearchnet_lora_dataset.json" 

# Número de ejemplos a usar. (¡ADVERTENCIA: 50000 es el MÁXIMO recomendado
# para evitar que tu entorno colapse. Puedes empezar con 5000!)
NUM_SAMPLES = 50000 

print(f"Cargando CodeSearchNet (Python) - {NUM_SAMPLES} muestras...")

# 1. Cargar el dataset. 'Nan-Do/code-search-net-python' es una versión popular.
try:
    # Usamos 'test' o 'validation' si el 'train' es demasiado grande para tu RAM
    raw_csn = load_dataset('Nan-Do/code-search-net-python', split='train[:{}]'.format(NUM_SAMPLES))
except Exception as e:
    print(f"Error al cargar el dataset de Hugging Face: {e}")
    print("Asegúrate de tener conexión a Internet.")
    exit()

# 2. Definir la función de transformación (adapta el formato al prompt/completion)
def format_for_lora(example):
    # Generamos un prompt que instruye al modelo usando la documentación.
    # El modelo base SantaCoder está acostumbrado a este formato.
    prompt_text = (
        f"# Descripción: {example['docstring_summary']}\n"
        f"# Completa la siguiente función:\n"
        f"def {example['func_name']}("
    )
    
    # La completion es el cuerpo completo de la función
    completion_text = example['code']
    
    return {
        "prompt": prompt_text,
        "completion": completion_text
    }

# 3. Aplicar la transformación
print("Aplicando la transformación al formato 'prompt/completion'...")
lora_dataset = raw_csn.map(
    format_for_lora, 
    remove_columns=raw_csn.column_names,
    # El campo 'code' en CodeSearchNet ya tiene la definición de 'def', 
    # pero a veces el prompt_text hace que se duplique.
    # Necesitas un poco de limpieza aquí, pero para empezar, esto funciona.
)

# 4. Guardar el resultado en un solo archivo JSON
print(f"Guardando {len(lora_dataset)} ejemplos en {OUTPUT_FILE}...")
lora_dataset.to_json(OUTPUT_FILE)

print("✅ Pre-procesamiento completado.")
print("Ahora actualiza tu app.py para usar 'codesearchnet_lora_dataset.json'.")