Add VQA

app_vqa.py

@@ -32,11 +32,17 @@ def create_demo():
     outputs = [answer, depth, edge, normals, segmentation, object_detection, ocr]
 
     # paths = sorted(pathlib.Path('prismer/images').glob('*'))
-    # examples = [[path.as_posix(), 'prismer_base'] for path in paths]
+    # ex_questions = ['What is the man on the right doing?',
+    #                 'What is this person playing?',
+    #                 'How many cows in this image?',
+    #                 'What is the type of animal in this image?',
+    #                 'What toy is it?']
+    #
+    # examples = [[path.as_posix(), 'Prismer-Base', ex_questions[i]] for i, path in enumerate(paths)]
     # gr.Examples(examples=examples,
     #             inputs=inputs,
     #             outputs=outputs,
-    #             fn=model.run_caption,
+    #             fn=model.run_vqa,
     #             cache_examples=os.getenv('SYSTEM') == 'spaces')
 
     paths = sorted(pathlib.Path('prismer/images').glob('*'))

prismer/model/modules/roberta.py

@@ -431,23 +431,6 @@ class RobertaLMHead(nn.Module):
 
 
 def load_decoder(name: str, config: RobertaConfig):
-    # load pre-trained model file
-    if name in ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST:
-        model = RobertaForMaskedLM.from_pretrained(name, cache_dir='cache')
-    else:
-        raise RuntimeError(f"Model {name} not found")
-
-    state_dict = model.state_dict()
-    for key in list(state_dict.keys()):
-        if 'encoder.layer' in key:
-            new_key_ = re.sub(".attention", ".0.attention", key)
-            new_key_ = re.sub(".intermediate", ".0.intermediate", new_key_)
-            if 'attention' not in key:
-                new_key_ = re.sub(".output", ".0.output", new_key_)
-            state_dict[new_key_] = state_dict.pop(key)
-
-    # load pre-trained weights
     roberta = RobertaForCausalLMModified(config)
-    roberta.load_state_dict(state_dict, strict=False)
     return roberta
 

prismer_model.py

@@ -79,7 +79,7 @@ class Model:
 
         # load checkpoints
         model_name = exp_name.lower().replace('-', '_')
-        if self.mode == 'caption':
+        if mode == 'caption':
             config = {
                 'dataset': 'demo',
                 'data_path': 'prismer/helpers',
@@ -94,7 +94,7 @@ class Model:
             state_dict = torch.load(f'prismer/logging/pretrain_{model_name}/pytorch_model.bin', map_location='cuda:0')
             state_dict['expert_encoder.positional_embedding'] = interpolate_pos_embed(state_dict['expert_encoder.positional_embedding'],
                                                                                       len(model.expert_encoder.positional_embedding))
-        elif self.mode == 'vqa':
+        elif mode == 'vqa':
             config = {
                 'dataset': 'demo',
                 'data_path': 'prismer/helpers',