Update app.py

app.py

@@ -1,69 +1,32 @@
 import os
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments, DataCollatorWithPadding
-from datasets import load_dataset
+from huggingface_hub import login
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
-# ==============================
-# Configuración del modelo
-# ==============================
-MODEL_NAME = "bigcode/starcoder"
-OUTPUT_DIR = "./results"
+# Nombre del modelo (público y sin restricciones)
+MODEL_NAME = "bigcode/santacoder"
 
-# Cargar tokenizer y modelo
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+# Obtener el token del entorno (desde Settings → Secrets)
+hf_token = os.environ.get("HF_TOKEN")
 
-# 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]'})
+# 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
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
 
-# Si agregaste un token nuevo, redimensionar embeddings
-# model.resize_token_embeddings(len(tokenizer))
-
-# ==============================
-# Preparar dataset
-# ==============================
-# Ejemplo con wikitext (reemplaza con tu dataset)
-dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="train[:5%]")  # ejemplo pequeño
-
-def tokenize_function(examples):
-    return tokenizer(examples["text"], truncation=True)
-
-tokenized_dataset = dataset.map(tokenize_function, batched=True)
-
-# ==============================
-# Configuración del DataCollator
-# ==============================
-data_collator = DataCollatorWithPadding(tokenizer=tokenizer, padding=True)
-
-# ==============================
-# Configuración del Trainer
-# ==============================
-training_args = TrainingArguments(
-    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",
-)
+# Crear el pipeline
+generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
-trainer = Trainer(
-    model=model,
-    args=training_args,
-    train_dataset=tokenized_dataset,
-    eval_dataset=tokenized_dataset,
-    tokenizer=tokenizer,
-    data_collator=data_collator,
-)
+# Ejemplo simple de generación
+def generate_text(prompt):
+    output = generator(prompt, max_new_tokens=120, temperature=0.7, top_p=0.95)
+    return output[0]["generated_text"]
 
-# ==============================
-# Iniciar entrenamiento
-# ==============================
-trainer.train()
+# Ejemplo de prueba
+if __name__ == "__main__":
+    texto = "AmorCoderAI es una IA creada para"
+    print(generate_text(texto))