Update app.py

app.py

@@ -2,134 +2,113 @@ import os
 import gradio as gr
 from huggingface_hub import login
 from datasets import load_dataset
-from transformers import (
-    AutoTokenizer,
-    AutoModelForCausalLM,
-    Trainer,
-    TrainingArguments,
-    DataCollatorForLanguageModeling,
-    pipeline,
-)
+from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments, DataCollatorForLanguageModeling, pipeline
+from peft import PeftModel
 
 # ============================================================
-# 🔐 Autenticación segura con tu token
+# 🔐 Autenticación HuggingFace
 # ============================================================
 hf_token = os.environ.get("HF_TOKEN")
-
 if hf_token:
     login(token=hf_token)
 else:
-    print("⚠️ No se encontró el token. Agrega 'HF_TOKEN' en Settings → Secrets → Add new secret")
+    print("⚠️ No se encontró el token. Agrega 'HF_TOKEN' en Settings → Secrets.")
 
 # ============================================================
-# ⚙️ Configuración del modelo base y dataset
+# ⚙️ Configuración del modelo y dataset
 # ============================================================
-MODEL_NAME = "bigcode/santacoder"   # Modelo libre y compatible con Hugging Face
-DATASET_PATH = "tu_dataset.json"        # Archivo dataset que subiste al Space
-OUTPUT_DIR = "lora_output"           # Carpeta donde se guarda el modelo entrenado
-
-# Crear carpeta de salida si no existe
-os.makedirs(OUTPUT_DIR, exist_ok=True)
-
-# Cargar modelo y tokenizer
-print("🔄 Cargando modelo base...")
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_auth_token=hf_token)
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, use_auth_token=hf_token)
+BASE_MODEL = "bigcode/santacoder"  # Modelo público
+LORA_PATH = "./lora_output"        # Carpeta donde se guardará LoRA
+DATASET_PATH = "tu_dataset.json"   # Cambia aquí al nombre de tu dataset
 
 # ============================================================
-# 🧩 Función de entrenamiento LoRA (simple y funcional)
+# 🧩 Función de entrenamiento LoRA
 # ============================================================
 def train_lora(epochs, batch_size, learning_rate):
     try:
-        # Cargar dataset JSON
         dataset = load_dataset("json", data_files=DATASET_PATH)
-
-        # Tokenización del dataset
-        def tokenize_fn(example):
-            text = example["prompt"] + example["completion"]
-            return tokenizer(
-                text,
+        tokenized = dataset.map(
+            lambda e: tokenizer(
+                e["prompt"] + e["completion"],
                 truncation=True,
                 padding="max_length",
-                max_length=256,
-            )
-
-        tokenized = dataset.map(tokenize_fn, batched=True)
-
-        # Preparar data collator
-        data_collator = DataCollatorForLanguageModeling(
-            tokenizer=tokenizer, mlm=False
+                max_length=256
+            ),
+            batched=True
         )
 
-        # Configuración del entrenamiento
+        data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
+
         training_args = TrainingArguments(
-            output_dir=OUTPUT_DIR,
+            output_dir=LORA_PATH,
             per_device_train_batch_size=int(batch_size),
             num_train_epochs=int(epochs),
-            learning_rate=float(learning_rate),
-            logging_steps=10,
+            learning_rate=learning_rate,
             save_total_limit=1,
-            push_to_hub=False,
-            report_to="none",
+            logging_steps=10,
+            push_to_hub=False
         )
 
-        # Entrenador
         trainer = Trainer(
-            model=model,
+            model=base_model,
             args=training_args,
             train_dataset=tokenized["train"],
             data_collator=data_collator,
         )
 
-        # Entrenar modelo
         trainer.train()
+        # Guardar LoRA
+        model.save_pretrained(LORA_PATH)
+        tokenizer.save_pretrained(LORA_PATH)
 
-        # Guardar resultados
-        model.save_pretrained(OUTPUT_DIR)
-        tokenizer.save_pretrained(OUTPUT_DIR)
-
-        return "✅ Entrenamiento completado con éxito. Modelo guardado en ./lora_output"
+        return "✅ Entrenamiento completado y guardado en ./lora_output"
     except Exception as e:
-        return f"❌ Error durante el entrenamiento: {str(e)}"
+        return f"❌ Error durante el entrenamiento: {e}"
 
 # ============================================================
-# 🤖 Función de prueba del modelo entrenado
+# 🤖 Función para generar texto usando LoRA sobre el modelo base
 # ============================================================
-def generate_text(prompt):
+def generate_text(prompt_text):
     try:
-        generator = pipeline(
-            "text-generation",
-            model=OUTPUT_DIR,
-            tokenizer=tokenizer,
-        )
-        output = generator(prompt, max_new_tokens=100, temperature=0.7, top_p=0.9)
+        # Cargar modelo base
+        tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
+        base_model = AutoModelForCausalLM.from_pretrained(BASE_MODEL)
+
+        # Aplicar LoRA
+        model = PeftModel.from_pretrained(base_model, LORA_PATH)
+
+        generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
+        output = generator(prompt_text, max_new_tokens=100, temperature=0.7, top_p=0.9)
         return output[0]["generated_text"]
     except Exception as e:
-        return f"⚠️ Error al generar texto: {str(e)}"
+        return f"❌ Error generando texto: {e}"
 
 # ============================================================
-# 💻 Interfaz de usuario (Gradio)
+# 💻 Interfaz Gradio
 # ============================================================
-with gr.Blocks(title="💙 AmorCoderAI - Entrenamiento LoRA") as demo:
+tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
+base_model = AutoModelForCausalLM.from_pretrained(BASE_MODEL)
+
+with gr.Blocks(title="AmorCoderAI - Entrenamiento LoRA") as demo:
     gr.Markdown("# 💙 AmorCoderAI - Entrenamiento y Pruebas")
-    gr.Markdown("Entrena y prueba tu modelo basado en `bigcode/santacoder` con LoRA.")
+    gr.Markdown("Entrena y prueba tu modelo basado en `bigcode/santacoder` con LoRA")
 
     with gr.Tab("🧠 Entrenar"):
         epochs = gr.Number(value=1, label="Épocas")
         batch_size = gr.Number(value=2, label="Tamaño de lote")
         learning_rate = gr.Number(value=5e-5, label="Tasa de aprendizaje")
         train_button = gr.Button("🚀 Iniciar entrenamiento")
-        train_output = gr.Textbox(label="Resultado", lines=3)
+        train_output = gr.Textbox(label="Resultado")
         train_button.click(train_lora, inputs=[epochs, batch_size, learning_rate], outputs=train_output)
 
     with gr.Tab("✨ Probar modelo"):
         prompt = gr.Textbox(label="Escribe un prompt")
         generate_button = gr.Button("💬 Generar texto")
-        output_box = gr.Textbox(label="Salida generada", lines=6)
+        output_box = gr.Textbox(label="Salida generada")
         generate_button.click(generate_text, inputs=prompt, outputs=output_box)
 
 # ============================================================
 # 🚀 Lanzar app
 # ============================================================
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch()