Spaces:
Sleeping
Sleeping
File size: 2,671 Bytes
9eddc4e 17a838d 9eddc4e 17a838d 9eddc4e 17a838d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
# 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) |