Create app.py

app.py

@@ -0,0 +1,54 @@
+import gradio as gr
+import torch
+from transformers import DonutProcessor, VisionEncoderDecoderModel
+import re
+import json
+from huggingface_hub import HfApi
+import os
+p1=os.environ.get("PATH_MODEL")
+p2=os.environ.get("PATH_MODEL_v2")
+print(p1,p2)
+
+PATH_MODEL = "fruk19/donut_nfact_v4"
+processor = DonutProcessor.from_pretrained(PATH_MODEL)
+model = VisionEncoderDecoderModel.from_pretrained(PATH_MODEL)
+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)
+
+  task_prompt = "<s_nfact>"
+  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"],
+             )
+             
+demo.launch()