Update app.py

app.py

@@ -1,32 +1,98 @@
 import os
+import gradio as gr
 from huggingface_hub import login
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from datasets import load_dataset
+from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments, DataCollatorForLanguageModeling, pipeline
 
-# Nombre del modelo (público y sin restricciones)
-MODEL_NAME = "bigcode/santacoder"
-
-# Obtener el token del entorno (desde Settings → Secrets)
+# ============================================================
+# 🔐 Autenticación segura con tu token
+# ============================================================
 hf_token = os.environ.get("HF_TOKEN")
 
-# Iniciar sesión segura (sin mostrar token)
 if hf_token:
     login(token=hf_token)
 else:
     print("⚠️  No se encontró el token. Agrega 'HF_TOKEN' en Settings → Secrets.")
 
-# Cargar el modelo y el tokenizer
+# ============================================================
+# ⚙️ Configuración del modelo base y dataset
+# ============================================================
+MODEL_NAME = "bigcode/santacoder"  # Modelo público similar a StarCoder
+DATASET_PATH = "dataset.json"       # Tu dataset subido al Space
+OUTPUT_DIR = "./lora_output"
+
+# Cargar modelo y tokenizer
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
 
-# Crear el pipeline
-generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
+# ============================================================
+# 🧩 Función de entrenamiento LoRA
+# ============================================================
+def train_lora(epochs, batch_size, learning_rate):
+    try:
+        dataset = load_dataset("json", data_files=DATASET_PATH)
+        tokenized = dataset.map(lambda e: tokenizer(e["prompt"] + e["completion"], truncation=True, padding="max_length", max_length=256))
+
+        data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
+
+        training_args = TrainingArguments(
+            output_dir=OUTPUT_DIR,
+            per_device_train_batch_size=batch_size,
+            num_train_epochs=epochs,
+            learning_rate=learning_rate,
+            save_total_limit=1,
+            logging_steps=10,
+            push_to_hub=False
+        )
+
+        trainer = Trainer(
+            model=model,
+            args=training_args,
+            train_dataset=tokenized["train"],
+            data_collator=data_collator,
+        )
+
+        trainer.train()
+        model.save_pretrained(OUTPUT_DIR)
+        tokenizer.save_pretrained(OUTPUT_DIR)
 
-# Ejemplo simple de generación
+        return "✅ Entrenamiento completado con éxito y guardado en ./lora_output"
+    except Exception as e:
+        return f"❌ Error durante el entrenamiento: {e}"
+
+# ============================================================
+# 🤖 Función de prueba del modelo
+# ============================================================
 def generate_text(prompt):
-    output = generator(prompt, max_new_tokens=120, temperature=0.7, top_p=0.95)
+    generator = pipeline("text-generation", model=OUTPUT_DIR, tokenizer=tokenizer)
+    output = generator(prompt, max_new_tokens=100, temperature=0.7, top_p=0.9)
     return output[0]["generated_text"]
 
-# Ejemplo de prueba
+# ============================================================
+# 💻 Interfaz de usuario (Gradio)
+# ============================================================
+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")
+
+    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")
+
+        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")
+
+        generate_button.click(generate_text, inputs=prompt, outputs=output_box)
+
+# ============================================================
+# 🚀 Lanzar app
+# ============================================================
 if __name__ == "__main__":
-    texto = "AmorCoderAI es una IA creada para"
-    print(generate_text(texto))
+    demo.launch()