add display visualbert

model_loader.py

@@ -5,6 +5,7 @@ from datasets import load_dataset, get_dataset_split_names
 import numpy as np
 
 import requests
+import streamlit as st
 from transformers import ViltProcessor, ViltForQuestionAnswering
 from transformers import AutoProcessor, AutoModelForCausalLM
 from transformers import BlipProcessor, BlipForQuestionAnswering
@@ -87,6 +88,7 @@ def get_item(image, question, tokenizer, image_model, model_name):
     )
     visual_embeds = get_img_feats(image, image_model=image_model, name=model_name)\
         .squeeze(2, 3).unsqueeze(0)
+    st.text(f"ques embed: {inputs.shape}, visual: {visual_embeds.shape}")
     visual_token_type_ids = torch.ones(visual_embeds.shape[:-1], dtype=torch.long)
     visual_attention_mask = torch.ones(visual_embeds.shape[:-1], dtype=torch.float)
     upd_dict = {
@@ -95,7 +97,7 @@ def get_item(image, question, tokenizer, image_model, model_name):
         "visual_attention_mask": visual_attention_mask,
     }
     inputs.update(upd_dict)
-
+   
     return upd_dict, inputs