vqa

app.py

@@ -48,22 +48,21 @@ device_type = 'cuda' if use_cuda else 'cpu'
 
 ## Load model
 
+### Captioning 
 config = 'configs/image/ePALM_caption.yaml'
 # config = yaml.load(open(config, 'r'), Loader=yaml.Loader)
 config = yaml.load(open(config, 'r'))
 
-
 text_model = 'facebook/opt-2.7b' 
 vision_model_name = 'vit_base_patch16_224'
 
 # text_model = 'facebook/opt-6.7b' 
 # vision_model_name = 'vit_large_patch16_224'
 
-
 start_layer_idx = 19
 end_layer_idx = 31
 low_cpu = True 
-model = ePALM(opt_model_name=text_model, 
+model_caption = ePALM(opt_model_name=text_model, 
                vision_model_name=vision_model_name, 
                use_vis_prefix=True, 
                start_layer_idx=start_layer_idx, 
@@ -73,14 +72,48 @@ model = ePALM(opt_model_name=text_model,
                low_cpu=low_cpu
 )
 print("Model Built")
-model.to(device)
-
+model_caption.to(device)
 
 checkpoint_path = 'checkpoints/float32/ePALM_caption/checkpoint_best.pth'
 # checkpoint_path = '/data/mshukor/logs/eplam/models/accelerate/ePALM_pt_L_acc_caption/checkpoint_best.pth'
 checkpoint = torch.load(checkpoint_path, map_location='cpu') 
 state_dict = checkpoint['model']
-msg = model.load_state_dict(state_dict,strict=False)  
+msg = model_caption.load_state_dict(state_dict,strict=False)  
+
+
+
+
+###### VQA
+
+
+
+config = 'configs/image/ePALM_vqa.yaml'
+config = yaml.load(open(config, 'r'))
+
+start_layer_idx = 19
+end_layer_idx = 31
+low_cpu = True 
+model_vqa = ePALM(opt_model_name=text_model, 
+               vision_model_name=vision_model_name, 
+               use_vis_prefix=True, 
+               start_layer_idx=start_layer_idx, 
+               end_layer_idx=end_layer_idx, 
+               return_hidden_state_vision=True, 
+               config=config,
+               low_cpu=low_cpu
+)
+print("Model Built")
+model_vqa.to(device)
+
+
+checkpoint_path = 'checkpoints/float32/ePALM_vqa/checkpoint_best.pth'
+checkpoint = torch.load(checkpoint_path, map_location='cpu') 
+state_dict = checkpoint['model']
+msg = model_vqa.load_state_dict(state_dict,strict=False)  
+
+
+
+
 
 
 ## Load tokenizer
@@ -88,7 +121,10 @@ tokenizer = AutoTokenizer.from_pretrained(text_model, use_fast=False)
 eos_token = tokenizer.eos_token
 pad_token = tokenizer.pad_token
 
+special_answer_token = '</a>'
 
+special_tokens_dict = {'additional_special_tokens': [special_answer_token]}
+tokenizer.add_special_tokens(special_tokens_dict)
 
 
 image_size = 224
@@ -112,7 +148,8 @@ num_beams=3
 max_length=30
 
 
-model.bfloat16()
+model_caption.bfloat16()
+model_vqa.bfloat16()
 
 
 def inference(image, audio, video, task_type, instruction):
@@ -120,6 +157,11 @@ def inference(image, audio, video, task_type, instruction):
     if task_type == 'Image Captioning':
         text = ['']  
         text_input = tokenizer(text, padding='longest', return_tensors="pt").to(device) 
+        model = model_caption
+    elif task_type == 'Visual Question Answering':
+        question = instruction+'?'+special_answer_token
+        text_input = tokenizer(question, padding='longest', return_tensors="pt").to(device) 
+        model = model_vqa
     else:
         raise NotImplemented
 
@@ -139,10 +181,15 @@ def inference(image, audio, video, task_type, instruction):
         out = model(image=image, text=text_input, mode='generate', return_dict=True, max_length=max_length, 
                     do_sample=do_sample, num_beams=num_beams)
 
-    out_decode = []
-    for i, o in enumerate(out):
-        res = tokenizer.decode(o)
-        response = res.split('</s>')[1].replace(pad_token, '').replace('</s>', '').replace(eos_token, '') # skip_special_tokens=True
+
+    if 'Captioning' in task_type:
+        for i, o in enumerate(out):
+            res = tokenizer.decode(o)
+            response = res.split('</s>')[1].replace(pad_token, '').replace('</s>', '').replace(eos_token, '') # skip_special_tokens=True
+    else:
+        for o in out:
+            o_list = o.tolist()
+            response = tokenizer.decode(o_list).split(special_answer_token)[1].replace(pad_token, '').replace('</s>', '').replace(eos_token, '') # skip_special_tokens=True
 
     return response
 
@@ -152,14 +199,14 @@ outputs = ['text']
 examples = [
     ['examples/images/soccer.jpg', None, None, 'Image Captioning', None],
     ['examples/images/ski.jpg', None, None, 'Visual Question Answering', 'what does the woman wearing black do?'],
-    ['examples/images/banana.jpg', None, None, 'Visual Grounding', 'the detached banana'],
-    ['examples/images/skateboard.jpg', None, None, 'General', 'which region does the text " a yellow bird " describe?'],
-    ['examples/images/baseball.jpg', None, None, 'General', 'what color is the left car?'],
+    ['examples/images/banana.jpg', None, None, 'Image Captioning', None],
+    ['examples/images/skateboard.jpg', None, None, 'Visual Question Answering', 'what is on top of the skateboard?'],
+    ['examples/images/baseball.jpg', None, None, 'Image Captioning', None],
     [None, None, 'examples/videos/video7014.mp4', 'Video Captioning', None], 
     [None, None, 'examples/videos/video7017.mp4', 'Video Captioning', None], 
     [None, None, 'examples/videos/video7019.mp4', 'Video Captioning', None], 
     [None, None, 'examples/videos/video7021.mp4', 'Video Captioning', None], 
-    [None, None, 'examples/videos/video7021.mp4', 'General Video', "What is this sport?"], 
+    [None, None, 'examples/videos/video7021.mp4', 'Video Captioning', None], 
     [None, 'examples/audios/6cS0FsUM-cQ.wav', None, 'Audio Captioning', None],
     [None, 'examples/audios/AJtNitYMa1I.wav', None, 'Audio Captioning', None],
 ]