engineering the code

app.py

@@ -1,16 +1,8 @@
 import gradio as gr
 import os
-from transformers import ViltProcessor, ViltForQuestionAnswering
 
-def vqa(image, text):
-    processor = ViltProcessor.from_pretrained("dandelin/vilt-b32-finetuned-vqa")
-    model = ViltForQuestionAnswering.from_pretrained("dandelin/vilt-b32-finetuned-vqa")
-    encoding = processor(image, text, return_tensors="pt")
-
-    outputs = model(**encoding)
-    logits = outputs.logits
-    idx = logits.argmax(-1).item()
-    return f"{model.config.id2label[idx]}"
+from inference import *
+from utils import *
 
 
 with gr.Blocks() as block:
@@ -22,7 +14,17 @@ with gr.Blocks() as block:
 
     dogs = os.path.join(os.path.dirname(__file__), "617.jpg")
     image = gr.Image(type="pil", value=dogs)
-    btn.click(vqa, inputs=[image, txt], outputs=[txt_3])
+
+    selected_option = block.get_element("dropdown").value
+    # here we can insert two or more models to inference the data
+    if selected_option == "Model 1":
+        btn.click(inference, inputs=[image, txt], outputs=[txt_3 + " Model 1"])
+
+    elif selected_option == "Model 2":
+        btn.click(inference, inputs=[image, txt], outputs=[txt_3 + " Model 2"])
+
+    elif selected_option == "Model 3":
+        btn.click(inference, inputs=[image, txt], outputs=[txt_3 + " Model 3"])
 
     btn = gr.Button(value="Submit")
 

inference.py

@@ -0,0 +1,11 @@
+from transformers import ViltProcessor, ViltForQuestionAnswering
+
+def inference(image, text):
+    processor = ViltProcessor.from_pretrained("dandelin/vilt-b32-finetuned-vqa")
+    model = ViltForQuestionAnswering.from_pretrained("dandelin/vilt-b32-finetuned-vqa")
+    encoding = processor(image, text, return_tensors="pt")
+
+    outputs = model(**encoding)
+    logits = outputs.logits
+    idx = logits.argmax(-1).item()
+    return f"{model.config.id2label[idx]}"