Update app.py

app.py

@@ -50,7 +50,13 @@ with col2:
         # encoding = processor(image, question, words, boxes=boxes, return_tensors="pt")
         model_encoding = processor(images=image, text=question, return_tensors="pt")
         model_output = model.generate(**model_encoding)
-        output_text = processor.batch_decode(model_output, skip_special_tokens=True)[0]
+        match task_type:
+            case "Classification": 
+                output_text = processor.batch_decode(model_output, skip_special_tokens=True)[0]
+            case "Question Answering": 
+                output_text = processor.batch_decode(model_output, skip_special_tokens=True)[0]
+            case "Layout Analysis":
+                output_text = processor.batch_decode(model_output, skip_special_tokens=False)[0]
         st.write(output_text)    
     elif testButton and selected_file == "None":
         st.write("Please upload and select a document (/an image).")