import os
import gradio as gr
from loguru import logger
from chatpdf import ChatPDF  # Importar tu clase ChatPDF
from similarities import BertSimilarity  # Importar la clase Similarity

# Ruta al corpus PDF
CORPUS_PATH = os.path.join("corpus", "Acuerdo009.pdf")

# Cargar el modelo
def load_model():
    # Configura el modelo como lo haces en tu script original
    sim_model = BertSimilarity(model_name_or_path="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2", device=None)
    model = ChatPDF(
        similarity_model=sim_model,
        generate_model_type="auto",
        generate_model_name_or_path="LenguajeNaturalAI/leniachat-qwen2-1.5B-v0",
        corpus_files=[CORPUS_PATH],  # Cambia esto por tu archivo o archivos del corpus
    )
    return model

# Inicializar el modelo
model = load_model()

# Función para hacer predicciones utilizando el método `predict_stream`
def predict_stream(message, history):
    history_format = [[human, assistant] for human, assistant in history]
    model.history = history_format
    for chunk in model.predict_stream(message):
        yield chunk

# Interfaz de usuario con Gradio
chatbot_stream = gr.Chatbot(
    height=600,
    avatar_images=("assets/user.png", "assets/Logo1.png"),  # Asegúrate de que estas imágenes estén en el Space
    bubble_full_width=False
)

# Configuración de la interfaz
title = "🤖 ChatPDF Zonia 🤖"
examples = ['¿Puede hablarme del PNL?', 'Introducción a la PNL']

chat_interface_stream = gr.ChatInterface(
    predict_stream,
    textbox=gr.Textbox(lines=4, placeholder="Hazme una pregunta", scale=7),
    title=title,
    chatbot=chatbot_stream,
    examples=examples,
    theme='soft',
)

# Lanzar la aplicación con Gradio
with gr.Blocks() as demo:
    chat_interface_stream.render()

demo.queue().launch()