import gradio as gr
import os
import spaces
from huggingface_hub import login
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TextIteratorStreamer
from threading import Thread

# Autenticación en Hugging Face
hf_token = os.getenv("LLAMA31")  # Asegúrate de configurar la variable de entorno HF_TOKEN con tu token
if hf_token:
    login(token=hf_token)
else:
    raise ValueError("Hugging Face token no encontrado. Asegúrate de que la variable de entorno HF_TOKEN esté configurada.")

# Configuración para cargar el modelo en 4 bits utilizando bitsandbytes
bnb_config = BitsAndBytesConfig(load_in_4bit=True)

# Cargar el tokenizador y el modelo cuantizado desde Hugging Face
model_id = "meta-llama/Meta-Llama-3.1-8B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=bnb_config,  # Cuantización en 4 bits
    device_map="auto"
)

# Definimos un mensaje inicial que simula un rol del sistema y contiene instrucciones sobre cómo responder, adaptado al español
system_message = {
    "role": "system",
    "content": """
Eres un asistente útil.
Recibirás una pregunta y un conjunto de respuestas junto con una puntuación de confianza entre 0 y 1 para cada respuesta.
Tu trabajo es convertir esta información en una respuesta corta y coherente.
Por ejemplo:
Pregunta: "¿A quién se le está facturando?", respuesta: {"answer": "John Doe", "confidence": 0.98}
Deberías responder algo como:
Con un alto grado de confianza, puedo decir que se le está facturando a John Doe.
Pregunta: "¿Cuál es el total de la factura?", respuesta: [{"answer": "154.08", "confidence": 0.75}, {"answer": "155", "confidence": 0.25}]
Deberías responder algo como:
Creo que el total de la factura es de $154.08 aunque también podría ser $155.
"""}

# Definimos la función de inferencia utilizando el modelo Meta Llama 3.1 8B Instruct
def chat_fn(multimodal_message):
    # Extraemos el texto de la pregunta del mensaje proporcionado por el usuario
    question = multimodal_message["text"]
    
    # Construimos el mensaje para el modelo
    conversation = [{"role": "user", "content": question}]

    # Generamos los IDs de entrada utilizando el tokenizador del modelo
    input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt")
    input_ids = input_ids.to(model.device)
    
    # Configuramos el streamer para la generación progresiva de texto
    streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)

    # Configuramos los argumentos de generación
    generate_kwargs = dict(
        input_ids=input_ids,
        streamer=streamer,
        max_new_tokens=500,  # Ajusta esto según tus necesidades
        do_sample=True,
        temperature=0.7,  # Ajusta la temperatura según tus necesidades
    )

    # Iniciamos la generación de texto en un hilo separado
    t = Thread(target=model.generate, kwargs=generate_kwargs)
    t.start()

    # Iteramos sobre los tokens generados y construimos la respuesta
    message = ""
    for text in streamer:
        message += text
        yield message

# Usamos la clase 'Blocks' de Gradio para definir la interfaz de usuario
with gr.Blocks() as demo:
    # Título de la aplicación en español
    gr.Markdown("# 🔍 Chatbot Analizador de Documentos")
    
    # Cuadro de texto para mostrar la respuesta generada, etiquetado en español
    response = gr.Textbox(lines=5, label="Respuesta")
    
    # Campo de texto multimodal para que el usuario suba un archivo e ingrese una pregunta, en español
    chat = gr.MultimodalTextbox(file_types=["image"], interactive=True,
                                show_label=False, placeholder="Sube una imagen del documento haciendo clic en '+' y haz una pregunta.")
    
    # Se asigna la función chat_fn para que se ejecute cuando el usuario envíe un mensaje en chat
    chat.submit(chat_fn, inputs=chat, outputs=response)

# Lanza la aplicación si este archivo es ejecutado directamente
if __name__ == "__main__":
    demo.launch()