import gradio as gr
import torch
from transformers import DonutProcessor, VisionEncoderDecoderModel
import re
import json
from huggingface_hub import HfApi
import os

PATH_MODEL = os.environ.get('MODEL_REPO_ID')
token = os.environ.get("HF_TOKEN")
task_prompt = os.environ.get("taskpromt")

processor = DonutProcessor.from_pretrained(PATH_MODEL, token=token)
model = VisionEncoderDecoderModel.from_pretrained(PATH_MODEL, token=token)
device = "cuda" if torch.cuda.is_available() else "cpu"
model.eval()
model.to(device)

def predict(test_image):
  pixel_values = processor(test_image, return_tensors="pt").pixel_values
  pixel_values = pixel_values.to(device)

  decoder_input_ids = processor.tokenizer(task_prompt, add_special_tokens=False, return_tensors="pt").input_ids
  decoder_input_ids = decoder_input_ids.to(device)

      # autoregressively generate sequence
  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,
              bad_words_ids=[[processor.tokenizer.unk_token_id]],
              return_dict_in_generate=True,
          )

      # turn into JSON
  seq = processor.batch_decode(outputs.sequences)[0]
  seq = seq.replace(processor.tokenizer.eos_token, "").replace(processor.tokenizer.pad_token, "")
  seq = re.sub(r"<.*?>", "", seq, count=1).strip()  # remove first task start token
  pred = processor.token2json(seq)
  return pred

demo = gr.Interface(fn=predict, 
             inputs=gr.inputs.Image(type="pil"),
             outputs="text",
             examples=["image_0.png","image_1.png","image_2.png","image_3.png","image_4.png","image_5.png","image_6.png","image_7.png","image_8.png","image_9.png"],
             )
             
demo.launch()