Andro0s's picture
Update app.py
2636834 verified
raw
history blame
1.03 kB
import os
from huggingface_hub import login
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
# Nombre del modelo (público y sin restricciones)
MODEL_NAME = "bigcode/santacoder"
# Obtener el token del entorno (desde Settings → Secrets)
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
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
# Crear el pipeline
generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
# 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"]
# Ejemplo de prueba
if __name__ == "__main__":
texto = "AmorCoderAI es una IA creada para"
print(generate_text(texto))