Create app.py

app.py

@@ -0,0 +1,44 @@
+import gradio
+import transformers
+import types
+
+
+checkpoint_path = "checkpoint"
+examples_path = "examples"
+
+MODEL = types.SimpleNamespace()
+MODEL.donut_processor = transformers.DonutProcessor.from_pretrained(checkpoint_path)
+MODEL.encoder_decoder = transformers.VisionEncoderDecoderModel.from_pretrained(checkpoint_path)
+MODEL.tokenizer = MODEL.donut_processor.tokenizer
+
+
+def generate_token_strings(images, skip_special_tokens=True) -> list[str]:
+    decoder_output = MODEL.encoder_decoder.generate(
+        images,
+        max_length=MODEL.encoder_decoder.config.decoder.max_length,
+        eos_token_id=MODEL.tokenizer.eos_token_id,
+        return_dict_in_generate=True,
+    )
+    return MODEL.tokenizer.batch_decode(
+        decoder_output.sequences, skip_special_tokens=skip_special_tokens
+    )
+
+def predict_string(image) -> str:
+    image = MODEL.donut_processor(
+        image, random_padding=False, return_tensors="pt"
+    ).pixel_values
+    string = generate_token_strings(image)[0]
+    return string
+
+
+interface = gradio.Interface(
+    title = "Making graphs accessible",
+    description = "Generate textual representation of a graph\n"
+    "https://www.kaggle.com/competitions/benetech-making-graphs-accessible",
+    fn=predict_string,
+    inputs="image",
+    outputs="text",
+    examples=examples_path,
+)
+
+interface.launch()