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))