Updated selections

app.py

@@ -8,7 +8,18 @@ from PIL import Image
 from io import BytesIO
 from transformers import VisionEncoderDecoderModel, VisionEncoderDecoderConfig, DonutProcessor, DonutImageProcessor, AutoTokenizer
 
-def run_prediction(sample, model, processor, prompt):
+def run_prediction(sample, model, processor, mode):
+
+    if mode == "OCR":
+        prompt = "<s><s_pretraining>"
+        no_repeat_ngram_size = 10
+    elif mode == "Table":
+        prompt = "<s><s_hierarchical>"
+        no_repeat_ngram_size = 45
+    else:
+        prompt = "<s><s_hierarchical>"
+        no_repeat_ngram_size = 10
+
 
     pixel_values = processor(np.array(
                     sample,
@@ -20,10 +31,12 @@ def run_prediction(sample, model, processor, prompt):
             pixel_values.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,
-            no_repeat_ngram_size=10,
-            num_beams=3
+            top_p=0.92, #.92,
+            top_k=10,
+            no_repeat_ngram_size=no_repeat_ngram_size,
+            num_beams=3,
+            output_attentions=False,
+            output_hidden_states=False,
         )
 
     # process output
@@ -54,7 +67,7 @@ with st.sidebar:
         image_bytes_data = uploaded_file.getvalue()
         image_upload = Image.open(BytesIO(image_bytes_data))
 
-    prompt = st.selectbox('Prompt', ('<s><s_pretraining>', '<s><s_plain>', '<s><s_hierarchical>'), index=2)
+    mode = st.selectbox('Mode', ('OCR', 'Tables', 'Element annotation'), index=2)
 
 if image_upload:
     image = image_upload
@@ -95,6 +108,6 @@ with st.spinner(f'Processing the document ...'):
             st.session_state['model'] = model
 
 st.info(f'Parsing document')
-parsed_info = run_prediction(image.convert("RGB"), model, processor, prompt)
+parsed_info = run_prediction(image.convert("RGB"), model, processor, mode)
 st.text(f'\nDocument:')
 st.text_area('Output text', value=parsed_info, height=800)