Update app.py

app.py

@@ -1,104 +1,69 @@
 import os
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments, DataCollatorForLanguageModeling
+from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments, DataCollatorWithPadding
 from datasets import load_dataset
-from peft import LoraConfig, get_peft_model
 
-# -------------------------------
-# Configuración
-# -------------------------------
-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 (JSON)
+# ==============================
+# Configuración del modelo
+# ==============================
+MODEL_NAME = "bigcode/starcoder"
+OUTPUT_DIR = "./results"
 
-# Crear carpeta si no existe
-os.makedirs(LORA_DIR, exist_ok=True)
-
-# -------------------------------
-# Cargar modelo y tokenizer
-# -------------------------------
-print("Cargando modelo base...")
+# Cargar tokenizer y modelo
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
-    device_map="auto",
-    torch_dtype=torch.float16
-)
 
-# -------------------------------
-# Configurar LoRA
-# -------------------------------
-lora_config = LoraConfig(
-    r=16,
-    lora_alpha=32,
-    target_modules=["q_proj","v_proj"],
-    lora_dropout=0.05,
-    bias="none",
-    task_type="CAUSAL_LM"
-)
+# Corregir padding token
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token  # usar EOS como padding
+    # Si prefieres agregar un token PAD nuevo:
+    # tokenizer.add_special_tokens({'pad_token': '[PAD]'})
 
-model = get_peft_model(model, lora_config)
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
 
-# -------------------------------
-# Cargar dataset
-# -------------------------------
-dataset = load_dataset("json", data_files=DATASET_PATH)
-dataset = dataset["train"]
+# Si agregaste un token nuevo, redimensionar embeddings
+# model.resize_token_embeddings(len(tokenizer))
 
-print("Columnas del dataset:", dataset.column_names)
+# ==============================
+# Preparar dataset
+# ==============================
+# Ejemplo con wikitext (reemplaza con tu dataset)
+dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="train[:5%]")  # ejemplo pequeño
 
-# -------------------------------
-# Función de tokenización
-# -------------------------------
 def tokenize_function(examples):
-    # 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)
+    return tokenizer(examples["text"], truncation=True)
 
-tokenized_datasets = dataset.map(tokenize_function, batched=True)
+tokenized_dataset = dataset.map(tokenize_function, batched=True)
 
-data_collator = DataCollatorForLanguageModeling(
-    tokenizer=tokenizer,
-    mlm=False
-)
+# ==============================
+# Configuración del DataCollator
+# ==============================
+data_collator = DataCollatorWithPadding(tokenizer=tokenizer, padding=True)
 
-# -------------------------------
-# Entrenamiento
-# -------------------------------
+# ==============================
+# Configuración del Trainer
+# ==============================
 training_args = TrainingArguments(
-    output_dir=LORA_DIR,
-    num_train_epochs=1,          # Ajusta según tu tiempo y GPU
-    per_device_train_batch_size=1,
-    save_steps=500,
-    save_total_limit=1,
-    logging_steps=50,
-    learning_rate=2e-4,
-    fp16=True,
-    gradient_accumulation_steps=4,
-    push_to_hub=False
+    output_dir=OUTPUT_DIR,
+    evaluation_strategy="steps",
+    per_device_train_batch_size=2,
+    per_device_eval_batch_size=2,
+    num_train_epochs=1,
+    save_steps=10,
+    save_total_limit=2,
+    logging_steps=5,
+    report_to="none",
 )
 
 trainer = Trainer(
     model=model,
     args=training_args,
-    train_dataset=tokenized_datasets,
-    data_collator=data_collator
+    train_dataset=tokenized_dataset,
+    eval_dataset=tokenized_dataset,
+    tokenizer=tokenizer,
+    data_collator=data_collator,
 )
 
-print("Comenzando entrenamiento de LoRA...")
-trainer.train()
-
-# -------------------------------
-# Guardar LoRA
-# -------------------------------
-print("Guardando LoRA en la carpeta:", LORA_DIR)
-model.save_pretrained(LORA_DIR)
-print("¡Entrenamiento completado! LoRA lista para producción.")
+# ==============================
+# Iniciar entrenamiento
+# ==============================
+trainer.train()