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

# Variables globales para el modelo y tokenizador
model = None
tokenizer = None

# 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 cargar el modelo (se ejecutará solo cuando sea necesario)
def load_model_if_needed():
    global model, tokenizer
    
    if model is None:
        print("Cargando modelo Zephyr-7B...")
        
        # Cargar el tokenizador
        model_name = "HuggingFaceH4/zephyr-7b-beta"
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        
        # Cargar el modelo con configuraciones para ahorrar memoria
        model = AutoModelForCausalLM.from_pretrained(
            model_name,
            torch_dtype=torch.float16,  # Usar precisión reducida
            device_map="auto",          # Distribuir el modelo automáticamente
            load_in_8bit=True,          # Cuantización a 8 bits
        )
        
        print("Modelo cargado correctamente!")

# Función principal que procesa las preguntas del usuario
def generate_response(user_message):
    try:
        # Cargar el modelo si aún no está cargado
        load_model_if_needed()
        
        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 formato de conversación para Zephyr
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_message}
        ]
        
        # Convertir mensajes al formato que espera el modelo
        prompt = tokenizer.apply_chat_template(messages, tokenize=False)
        
        # Generar respuesta
        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
        
        # Configuración de generación
        generation_config = {
            "max_new_tokens": 500,
            "temperature": 0.7,
            "top_p": 0.9,
            "do_sample": True,
            "pad_token_id": tokenizer.eos_token_id
        }
        
        # Generar respuesta
        with torch.no_grad():
            outputs = model.generate(**inputs, **generation_config)
        
        # Decodificar la respuesta
        full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Extraer solo la respuesta del asistente (después del último mensaje del usuario)
        assistant_response = full_response.split("assistant:")[-1].strip()
        
        return assistant_response
    
    except Exception as e:
        print(f"Error: {str(e)}")
        return f"Lo siento, ha ocurrido un error: {str(e)}"

# 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
if __name__ == "__main__":
    demo.launch()