Update app.py

app.py

@@ -9,7 +9,7 @@ from peft import LoraConfig, get_peft_model
 # -------------------------------
 MODEL_NAME = "codellama/CodeLlama-7b-hf"   # Modelo base
 LORA_DIR = "lora_codellama"               # Carpeta donde se guardará LoRA
-DATASET_PATH = "tu_dataset.json"          # Tu dataset local (JSONL o JSON)
+DATASET_PATH = "tu_dataset.json"          # Tu dataset local (JSON)
 
 # Crear carpeta si no existe
 os.makedirs(LORA_DIR, exist_ok=True)
@@ -43,10 +43,25 @@ model = get_peft_model(model, lora_config)
 # Cargar dataset
 # -------------------------------
 dataset = load_dataset("json", data_files=DATASET_PATH)
-dataset = dataset["train"]  # Asume que el JSON tiene solo la parte de entrenamiento
+dataset = dataset["train"]
 
+print("Columnas del dataset:", dataset.column_names)
+
+# -------------------------------
+# Función de tokenización
+# -------------------------------
 def tokenize_function(examples):
-    return tokenizer(examples["text"], truncation=True, max_length=512)
+    # Detectar columnas automáticamente
+    columns = dataset.column_names
+    if "prompt" in columns and "completion" in columns:
+        texts = [p + "\n" + c for p, c in zip(examples["prompt"], examples["completion"])]
+    elif "text" in columns:
+        texts = examples["text"]
+    else:
+        # Si no encuentra las columnas, lanza un error con info
+        raise ValueError(f"Columnas inválidas en dataset: {columns}")
+    
+    return tokenizer(texts, truncation=True, max_length=512)
 
 tokenized_datasets = dataset.map(tokenize_function, batched=True)
 
@@ -60,7 +75,7 @@ data_collator = DataCollatorForLanguageModeling(
 # -------------------------------
 training_args = TrainingArguments(
     output_dir=LORA_DIR,
-    num_train_epochs=1,          # Ajusta según tu tiempo
+    num_train_epochs=1,          # Ajusta según tu tiempo y GPU
     per_device_train_batch_size=1,
     save_steps=500,
     save_total_limit=1,
@@ -86,4 +101,4 @@ trainer.train()
 # -------------------------------
 print("Guardando LoRA en la carpeta:", LORA_DIR)
 model.save_pretrained(LORA_DIR)
-print("¡Entrenamiento completado! Ahora tu LoRA está lista para producción.")
+print("¡Entrenamiento completado! LoRA lista para producción.")