Update app.py

app.py

@@ -8,6 +8,17 @@ models = {
     "BLIP": (AutoProcessor, BlipForQuestionAnswering, "Salesforce/blip-vqa-base"),
 }
 
+def get_format_response(selected_model,outputs):
+    if selected_model=='ViLT':
+        logits = outputs.logits
+        idx = logits.argmax(-1).item()
+        answer = model.config.id2label[idx]
+        return answer
+    else:
+        answer = processor.decode(out[0], skip_special_tokens=True)
+        return answer
+        
+
 # Streamlit app
 st.title("Simple VQA App 🤖🎈")
 st.subheader("A demo app showcasing VQA models. ViLT and BLIP model.")
@@ -29,10 +40,8 @@ if uploaded_image and question:
     st.image(image, caption="Uploaded Image")
     encoding = processor(image, question, return_tensors="pt")
     outputs = model(**encoding)
-    logits = outputs.logits
-    idx = logits.argmax(-1).item()
-    answer = model.config.id2label[idx]
-    # Display answer with animation emoji
+    answer = get_format_response(selected_model,outputs)
+    # Display answer
     st.write(f"🤔 Model Answer: {answer} 🎉")
 # Disclaimer
 st.sidebar.markdown("This is a demo app showcasing VQA models. Actual performance may vary.")