Prompt option

app.py

@@ -8,7 +8,7 @@ from PIL import Image
 from io import BytesIO
 from transformers import VisionEncoderDecoderModel, VisionEncoderDecoderConfig, DonutProcessor, DonutImageProcessor, AutoTokenizer
 
-def run_prediction(sample, model, processor):
+def run_prediction(sample, model, processor, prompt):
 
     pixel_values = processor(np.array(
                     sample,
@@ -18,7 +18,7 @@ def run_prediction(sample, model, processor):
     with torch.no_grad():
         outputs = model.generate(
             pixel_values.to(device),
-            decoder_input_ids=processor.tokenizer("<s><s_plain>", add_special_tokens=False, return_tensors="pt").input_ids.to(device),
+            decoder_input_ids=processor.tokenizer(prompt, add_special_tokens=False, return_tensors="pt").input_ids.to(device),
             do_sample=True,
             top_p=0.92,
             top_k=5,
@@ -52,7 +52,9 @@ with st.sidebar:
     if uploaded_file is not None:
         # To read file as bytes:
         image_bytes_data = uploaded_file.getvalue()
-        image_upload = Image.open(BytesIO(image_bytes_data)) 
+        image_upload = Image.open(BytesIO(image_bytes_data))
+
+    prompt = st.selectbox('Prompt', ('<s><s_pretraining>', '<s><s_plain>', '<s><s_hierarchical>'), index=2)
 
 if image_upload:
     image = image_upload
@@ -87,6 +89,6 @@ with st.spinner(f'Processing the document ...'):
         model.to(device)
 
 st.info(f'Parsing document')
-parsed_info = run_prediction(image.convert("RGB"), model, processor)
+parsed_info = run_prediction(image.convert("RGB"), model, processor, prompt)
 st.text(f'\nDocument:')
 st.text_area('Output text', value=parsed_info, height=800)