import os
from huggingface_hub import login
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
import gradio as gr

# Obtén el token desde la variable de entorno
hf_token = os.getenv("LLAMA31")

if hf_token:
    # Autenticación en Hugging Face utilizando el 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 modelo y tokenizador
model_id = "meta-llama/Meta-Llama-3.1-8B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="auto",
    quantization_config=bnb_config  # Aplicar cuantización en 4 bits
)

# Definir la función de inferencia del chatbot
def chat_fn(multimodal_message):
    # Extraer el texto de la pregunta proporcionada por el usuario
    question = multimodal_message["text"]
    
    # Construir la conversación inicial con el mensaje del usuario
    conversation = [{"role": "user", "content": question}]

    # Generar 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)
    
    # Configurar el streamer para la generación progresiva de texto
    streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)

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

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

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

# Crear la interfaz de usuario utilizando Gradio
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.")
    
    # Asignar la función chat_fn para que se ejecute cuando el usuario envíe un mensaje en el chat
    chat.submit(chat_fn, inputs=chat, outputs=response)

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