feat: working with twin 🍩🍩s

app.py

@@ -66,16 +66,41 @@ st.text(f'{information} mode is ON!\nTarget receipt: {receipt}\n(opening image @
 image = Image.open(f"./img/receipt-{receipt}.jpg")
 st.image(image, caption='Your target receipt')
 
-st.text(f'baking the 🍩...')
-processor = DonutProcessor.from_pretrained("unstructuredio/donut-base-sroie")
-pretrained_model = VisionEncoderDecoderModel.from_pretrained("unstructuredio/donut-base-sroie")
+st.text(f'baking the 🍩s...')
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-pretrained_model.to(device)
+if information == 'Receipt Summary':
+    processor = DonutProcessor.from_pretrained("unstructuredio/donut-base-sroie")
+    pretrained_model = VisionEncoderDecoderModel.from_pretrained("unstructuredio/donut-base-sroie")
+    task_prompt = f"<s>"
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    pretrained_model.to(device)
 
-# pretrained_model.encoder.to(torch.bfloat16)
-# pretrained_model.eval()
+elif information == 'Receipt Menu Details':
+    processor = DonutProcessor.from_pretrained("naver-clova-ix/donut-base-finetuned-cord-v2")
+    pretrained_model = VisionEncoderDecoderModel.from_pretrained("naver-clova-ix/donut-base-finetuned-cord-v2")
+    task_prompt = f"<s_cord-v2>"
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    pretrained_model.to(device)
+    
+else:
+    # st.text(f'NotImplemented: soon you will be able to use it..')
+    processor_a = DonutProcessor.from_pretrained("unstructuredio/donut-base-sroie")
+    processor_b = DonutProcessor.from_pretrained("naver-clova-ix/donut-base-finetuned-cord-v2")
+    pretrained_model_a = VisionEncoderDecoderModel.from_pretrained("unstructuredio/donut-base-sroie")
+    pretrained_model_b = VisionEncoderDecoderModel.from_pretrained("naver-clova-ix/donut-base-finetuned-cord-v2")
+    
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    pretrained_model.to(device)
 
-st.text(f'parsing receipt..')
-parsed_receipt_info = run_prediction(image)
-st.text(f'\nRaw output:\n{parsed_receipt_info}')
+if information == 'Extract all!':
+    st.text(f'parsing receipt (extracting all)..')
+    pretrained_model, processor, task_prompt = pretrained_model_a, processor_a, f"<s>"
+    parsed_receipt_info_a = run_prediction(image)
+    pretrained_model, processor, task_prompt = pretrained_model_b, processor_b, f"<s_cord-v2>"
+    parsed_receipt_info_b = run_prediction(image)
+    st.text(f'\nRaw output a:\n{parsed_receipt_info_a}')
+    st.text(f'\nRaw output b:\n{parsed_receipt_info_b}')
+else:
+    st.text(f'parsing receipt..')
+    parsed_receipt_info = run_prediction(image)
+    st.text(f'\nRaw output:\n{parsed_receipt_info}')