# Imports
import re
import gradio as gr
import torch
from transformers import DonutProcessor, VisionEncoderDecoderModel
from dotenv import load_dotenv
import os
from torchvision import transforms
import torch
from PIL import Image

# device = 'cuda' if torch.cuda.is_available() else 'cpu'
model_name = 'naver-clova-ix/donut-base-finetuned-docvqa'

# Importante esta app esta pensada para que el modelo corra en CPU
processor = DonutProcessor.from_pretrained(model_name)
model = VisionEncoderDecoderModel.from_pretrained(model_name)

# Defino la funcion principal que ejecuta el modelo y obtiene los resultados
def process_image(image, question):
    # Paso por el procesador la imagen y especifico que los outputs sean tensores de pytorch
    pixel_values = processor(image, return_tensors='pt').pixel_values
    # Seteo el prompt
    prompt = f"<s_docvqa><s_question>{question}</s_question><s_answer>"
    # Generamos la sequencia de tokens de salida esto es un vector largo con los ids
    # Esta parte encodea la pregunta y se la pasa al decoder junto con la representación de la imagen post encoder
    decoder_input_ids = processor.tokenizer(prompt, add_special_tokens=False, return_tensors='pt').input_ids
    # Defino los outputs
    outputs = model.generate(
        pixel_values,
        decoder_input_ids=decoder_input_ids,
        max_length=model.decoder.config.max_position_embeddings,
        early_stopping=True,
        pad_token_id=processor.tokenizer.pad_token_id,
        eos_token_id=processor.tokenizer.eos_token_id,
        use_cache=True,
        num_beams=1, # Probar cambiando este parametro
        bad_words_ids=[[processor.tokenizer.unk_token_id]],
        return_dict_in_generate=True
    )
    # Realizo el Post-processing de la salida del modelo
    sequence = processor.batch_decode(outputs.sequences)[0]
    sequence = sequence.replace(processor.tokenizer.eos_token, "").replace(processor.tokenizer.pad_token, "")
    processor.tokenizer.eos_token
    sequence = re.sub(r"<.*?>", "", sequence, count=1).strip()
    
    return processor.token2json(sequence)

description = "Esta es una aplicacción realizada con el modelo Donut fine tuned en DocVQA"

demo = gr.Interface(
    fn=process_image,
    inputs=['image', 'text'],
    outputs='json',
    title='Demo: Document Question Answering',
    description=description,
    enable_queue=True,
    examples=[
        ['examples/dni_25.jpg', 'cual es el documento / document number?'],
        ['examples/extracto.jpg', 'cual es el telefono de centros servicios de banco galicia?'],
        ['examples/factura_5.jpg', 'cual es el total de la factura?'],
    ]
)

demo.launch(inline=True)