AmorCoderAI / app.py
Andro0s's picture
Update app.py
15e5efd verified
import os
import gradio as gr
from huggingface_hub import login
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
# ============================================================
# 🔐 Autenticación con Hugging Face
# ============================================================
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.")
# ============================================================
# ⚙️ Configuración y Carga ÚNICA del Modelo
# ============================================================
MODEL_NAME = "bigcode/santacoder"
print("🔄 Cargando modelo y tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
# 🚀 Inicialización del pipeline UNA SOLA VEZ
# Esto hace que la generación sea mucho más rápida en la app web.
generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
# ============================================================
# 🤖 Generación de texto (Función ahora más rápida)
# ============================================================
def generate_text(prompt):
try:
# ⚠️ Usa el 'generator' que ya está cargado globalmente
output = generator(prompt, max_new_tokens=200, temperature=0.7, top_p=0.9)
return output[0]["generated_text"]
except Exception as e:
return f"❌ Error al generar texto: {e}"
# ============================================================
# 💻 Interfaz con Gradio
# ============================================================
with gr.Blocks(title="💙 AmorCoderAI") as demo:
gr.Markdown("# 💙 AmorCoderAI")
gr.Markdown("Genera texto con un modelo BigCode libre (`santacoder`)")
prompt = gr.Textbox(label="🧠 Escribe un prompt (ej: 'def fibonacci(n):')", lines=4) # Más espacio para prompts
generate_button = gr.Button("💬 Generar texto")
output_box = gr.Textbox(label="✨ Resultado", lines=10) # Más espacio para el código generado
generate_button.click(generate_text, inputs=prompt, outputs=output_box)
if __name__ == "__main__":
# Usa `share=True` si necesitas compartirlo temporalmente
demo.launch(server_name="0.0.0.0", server_port=7860)