"""
Versión con cuantización de 4 bits para reducir aún más el uso de memoria.
Renombra este archivo a app.py para usar esta versión ultra optimizada.
"""

import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, 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=512,
            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 Phi-2 con cuantización de 4 bits...")

# Configuración de cuantización de 4 bits
quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True
)

# Usar un modelo pequeño optimizado para español
model_name = "microsoft/phi-2"  # Modelo de 2.7B parámetros, mucho más ligero que 7B

# Configuración para reducir el uso de memoria con cuantización de 4 bits
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=quantization_config,
    device_map="auto",
    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 con cuantización de 4 bits!")

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