Create prepare_codesearchnet.py

prepare_codesearchnet.py

@@ -0,0 +1,56 @@
+# 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'.")