app.py

"""
Versión alternativa usando un modelo aún más pequeño y optimizado para español.
"""

import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch
import gc
import os

# Función para cargar el prompt desde el archivo
def get_system_prompt():
   with open("prompt.txt", "r", encoding="utf-8") as f:
       return f.read().strip()

# Función para generar respuestas
def generate_response(user_message):
   try:
       if not user_message.strip():
           return "Por favor, escribe una pregunta para que pueda ayudarte."
       
       # Obtener el prompt del sistema
       system_prompt = get_system_prompt()
       
       # Crear el prompt completo
       prompt = f"{system_prompt}\n\nPregunta: {user_message}\n\nRespuesta:"
       
       # Generar respuesta usando el pipeline
       response = generator(
           prompt,
           max_new_tokens=256,
           temperature=0.7,
           top_p=0.9,
           do_sample=True,
           num_return_sequences=1
       )[0]["generated_text"]
       
       # Extraer solo la respuesta (después de "Respuesta:")
       assistant_response = response.split("Respuesta:")[-1].strip()
       
       # Forzar recolección de basura
       gc.collect()
       torch.cuda.empty_cache() if torch.cuda.is_available() else None
       
       return assistant_response
   
   except Exception as e:
       print(f"Error: {str(e)}")
       return f"Lo siento, ha ocurrido un error: {str(e)}"

# Cargar el modelo y crear el pipeline
print("Cargando modelo ultra ligero...")

# Usar un modelo muy pequeño con soporte para español
model_name = "PlanTL-GOB-ES/gpt2-base-bne"  # Modelo de ~125M parámetros, extremadamente ligero

# Configuración para reducir el uso de memoria
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
   model_name,
   torch_dtype=torch.float32,
   low_cpu_mem_usage=True
)

# Crear el pipeline de generación de texto
generator = pipeline(
   "text-generation",
   model=model,
   tokenizer=tokenizer
)

print("Modelo cargado correctamente!")

# Crear la interfaz de Gradio
demo = gr.Interface(
   fn=generate_response,
   inputs=gr.Textbox(
       placeholder="Escribe tu pregunta sobre emprendimiento aquí...",
       label="Tu pregunta"
   ),
   outputs=gr.Textbox(label="Respuesta de BITER"),
   title="BITER - Tu Mentor en Tiempo Real para Decisiones de Negocio",
   description="BITER es un asistente de IA que responde dudas de emprendedores como si fuera un CEO experimentado.",
   examples=[
       ["¿Cómo puedo validar mi idea de negocio con poco presupuesto?"],
       ["¿Cuál es la mejor estrategia para conseguir mis primeros clientes?"],
       ["¿Debería invertir en publicidad en redes sociales o en SEO?"]
   ],
   allow_flagging="never"
)

# Lanzar la aplicación con configuración para ahorrar memoria
if __name__ == "__main__":
   # Configurar menos workers para ahorrar memoria
   demo.queue(max_size=1).launch(share=False, debug=False)