Update Train.py

Train.py

@@ -1,81 +1,44 @@
-# ===============================
-# AmorCoder AI - Entrenamiento LoRA Avanzado
-# ===============================
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments
-from datasets import load_dataset
-from peft import LoraConfig, get_peft_model, TaskType
-
-# -------------------------------
-# 1️⃣ Modelo base
-# -------------------------------
-MODEL_NAME = "codellama/CodeLlama-7b-hf"
-print("Cargando modelo base...")
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
-    device_map="auto",
-    torch_dtype=torch.float16
-)
-
-# -------------------------------
-# 2️⃣ Configuración LoRA
-# -------------------------------
-print("Aplicando LoRA...")
-lora_config = LoraConfig(
-    task_type=TaskType.CAUSAL_LM,
-    r=16,
-    lora_alpha=32,
-    target_modules=["q_proj", "v_proj"],  # módulos recomendados para LLMs
-    lora_dropout=0.05,
-    bias="none"
-)
-model = get_peft_model(model, lora_config)
-
-# -------------------------------
-# 3️⃣ Dataset
-# -------------------------------
-print("Cargando dataset...")
-dataset = load_dataset("json", data_files={"train":"tu_dataset.json"}, split="train")
-
-def preprocess(example):
-    prompt = f"# Instrucción:\n{example['instruction']}\n\n# Código:\n"
-    input_ids = tokenizer(prompt, truncation=True, max_length=512)["input_ids"]
-    labels = tokenizer(example['code'], truncation=True, max_length=512)["input_ids"]
-    return {"input_ids": input_ids, "labels": labels}
-
-dataset = dataset.map(preprocess)
-
-# -------------------------------
-# 4️⃣ Argumentos de entrenamiento
-# -------------------------------
-training_args = TrainingArguments(
-    output_dir="./lora_codellama",
-    per_device_train_batch_size=1,   # usar gradient accumulation para batches grandes
-    gradient_accumulation_steps=4,
-    num_train_epochs=3,              # puedes subir a 5 para más precisión
-    learning_rate=3e-4,
-    fp16=True,
-    logging_steps=10,
-    save_steps=50,
-    save_total_limit=3,
-    report_to="none",                # para no depender de wandb u otro tracker
-)
-
-# -------------------------------
-# 5️⃣ Entrenamiento
-# -------------------------------
-trainer = Trainer(
-    model=model,
-    train_dataset=dataset,
-    args=training_args
-)
-
-print("Entrenando LoRA...")
-trainer.train()
-
-# -------------------------------
-# 6️⃣ Guardar pesos
-# -------------------------------
-model.save_pretrained("lora_codellama")
-print("✅ Entrenamiento completado. Pesos guardados en 'lora_codellama'.")
+def train_lora(epochs, batch_size, learning_rate):
+    try:
+        dataset = load_dataset("json", data_files=DATASET_PATH)
+        
+        # Tokenización correcta
+        def tokenize_fn(example):
+            return tokenizer(
+                example["prompt"] + example["completion"],
+                truncation=True,
+                padding="max_length",
+                max_length=256,
+            )
+
+        tokenized = dataset.map(tokenize_fn, batched=False)
+
+        # Asegúrate que las columnas correctas estén
+        tokenized.set_format(type="torch", columns=["input_ids", "attention_mask"])
+
+        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=int(epochs),
+            learning_rate=learning_rate,
+            save_total_limit=1,
+            logging_steps=10,
+            push_to_hub=False
+        )
+
+        trainer = Trainer(
+            model=base_model,
+            args=training_args,
+            train_dataset=tokenized["train"],
+            data_collator=data_collator,
+        )
+
+        trainer.train()
+        base_model.save_pretrained(LORA_PATH)
+        tokenizer.save_pretrained(LORA_PATH)
+
+        return "✅ Entrenamiento completado y guardado en ./lora_output"
+    except Exception as e:
+        return f"❌ Error durante el entrenamiento: {e}"