fix error in visualbert

model_loader.py

@@ -62,13 +62,20 @@ def load_dataset(type):
         raise ValueError("invalid dataset: ", type)
 '''
 
+def load_img_model(name):
+    """
+    loads image models for feature extraction
+    returns model name and the loaded model
+    """
+    if name == "resnet50":
+        model = resnet50(weights='DEFAULT')
+    elif name == "vitb16":
+        ## MOD Minxuan: add param
+        model = timm.create_model('vit_base_patch16_224', pretrained=True, num_classes=0)
+    else:
+        raise ValueError("undefined model name: ", name)
 
-def tokenize_function(examples, processor):
-    sample = {}
-    sample['inputs'] = processor(images=examples['image'], text=examples['question'], return_tensors="pt")
-    sample['outputs'] = examples['multiple_choice_answer']
-    return sample
-
+    return model, name
 
 def label_count_list(labels):
     res = {}
@@ -88,7 +95,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 = {
@@ -192,7 +199,8 @@ def get_answer(model_loader_args, img, question, model_name):
 
         # load question and image (processor = tokenizer)
         ## MOD Minxuan: fix error
-        _, inputs = get_item(img, question, processor, "resnet50")
+        img_model, name = load_img_model("resnet50")
+        _, inputs = get_item(img, question, processor, img_model, name)
         outputs = model(**inputs)
         #except Exception:
            # return err_msg()