Upload 25 files

.gitattributes

@@ -32,3 +32,8 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+data/random_sample_test_direct_ids.json filter=lfs diff=lfs merge=lfs -text
+data/random_sample_test_indirect_ids.json filter=lfs diff=lfs merge=lfs -text
+results_pair_dict1.json filter=lfs diff=lfs merge=lfs -text
+results_pair_dict2.json filter=lfs diff=lfs merge=lfs -text
+results/results_pair_dict.json filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,235 @@
+# %%
+from PIL import Image
+import torch
+import torchvision
+from torchvision import transforms
+from torchvision.transforms.functional import InterpolationMode
+import json
+from collections import defaultdict
+import numpy as np
+import matplotlib.pyplot as plt
+import random
+from os import listdir
+from os.path import isfile, join
+from torchvision.io import read_image
+from torchvision.utils import draw_bounding_boxes
+from PIL import Image
+import os
+from scipy.stats import rankdata
+import tqdm
+import streamlit as st
+import pandas as pd
+
+# %%
+def load_json(PATH):
+    if os.path.isfile(PATH) and os.access(PATH, os.R_OK):
+        with open(PATH) as json_file:
+            dict_data = json.load(json_file)
+    else:
+        print("The Path of", PATH,"is not exist")
+        dict_data = {}
+    return dict_data
+
+def get_list_folder(PATH):
+    return [name for name in os.listdir(PATH) if os.path.isdir(os.path.join(PATH, name))]
+
+def get_file_only(PATH):
+    return [f for f in os.listdir(PATH) if os.path.isfile(os.path.join(PATH, f))]
+
+# %%
+def compute_ndcg(ranks, scores, k=3):
+    """
+    ranks = [5, 1, 4, 2, 3]
+    scores = [0.1, 0.5, 0.3, 0.95, 1.0]
+    """
+    rank_score_tuple = list(zip(ranks, scores))
+
+    top_k = sorted(rank_score_tuple, key=lambda x: x[1], reverse=True)[:k]
+
+    dcg = sum([score / np.log2(rank + 1) for rank, score in top_k])
+
+    ideal_dcg = sum([score / np.log2(idx + 2) for idx, (_, score) in enumerate(top_k)])
+
+    ndcg = dcg / ideal_dcg
+    return ndcg
+
+def compute_ndcg_score_per_mode(pred_rank_dict, gpt_rel_scores, random_sample_dict, mode='indrect', split='test', k=200):
+    ndcg_scores = []
+
+    for key in tqdm.tqdm(pred_rank_dict.keys(), total=len(pred_rank_dict.keys())):
+        gpt_scores_for_key = [gpt_rel_scores[key][cand_key] if cand_key in gpt_rel_scores[key] else 0.0  for cand_key in random_sample_dict[key]]
+
+        pred_rank_for_key = pred_rank_dict[key]
+
+        ndcg_score = compute_ndcg(pred_rank_for_key, gpt_scores_for_key, k=k)
+        ndcg_scores.append(ndcg_score)
+
+    avg_ndcg_score = sum(ndcg_scores) / len(ndcg_scores)
+    print(f"Random split, mode={mode} ndcg score: ", avg_ndcg_score)
+    return avg_ndcg_score
+# %%
+
+def get_score_direct(random_sample_pair_test_direct, predictions, key_pair, similarity_score_test_direct, k = 200):
+    mode = 'direct'
+    i2t_ranks = []
+    t2i_ranks = []
+    i2t_rank_dict = {}
+    results_dict = {}
+    key_pair_reversed = {v: k for k, v in key_pair.items()}
+
+
+    for file_key in tqdm.tqdm(random_sample_pair_test_direct.keys(), total=len(random_sample_pair_test_direct.keys())):
+        i2t_rank = rankdata([predictions[str(file_key)+':'+str(key_pair[k])] for k in random_sample_pair_test_direct[file_key]])
+        t2i_rank = rankdata([predictions[str(key_pair_reversed[key_pair[k]])+':'+str(key_pair[file_key])] for k in random_sample_pair_test_direct[file_key]])
+
+        i2t_ranks.append(i2t_rank[-1])
+        t2i_ranks.append(t2i_rank[-1])
+        i2t_rank_dict[file_key] = i2t_rank
+
+    assert len(i2t_ranks) == len(t2i_ranks) == 1000
+
+    ndcg_score = compute_ndcg_score_per_mode(i2t_rank_dict, similarity_score_test_direct, random_sample_pair_test_direct, mode='indirect', split='test', k=200)
+
+    results_dict['direct'] = {}
+    results_dict['direct']['i2t rank'] = float(sum(i2t_ranks) / len(i2t_ranks))
+    results_dict['direct']['t2i rank'] = float(sum(t2i_ranks) / len(t2i_ranks))
+    results_dict['direct']['ndcg score'] = float(ndcg_score)
+    print(f"Random split, mode={mode} i2t rank: ", sum(i2t_ranks) / len(i2t_ranks))
+    print(f"Random split, mode={mode} t2i rank: ", sum(t2i_ranks) / len(t2i_ranks))
+    return results_dict
+
+# %%
+
+def get_score_indirect(random_sample_pair_test_indirect, predictions, key_pair, similarity_score_test_indirect, k = 200):
+    mode = 'indirect'
+    i2t_ranks = []
+    t2i_ranks = []
+    i2t_rank_dict = {}
+    results_dict = {}
+    key_pair_reversed = {v: k for k, v in key_pair.items()}
+
+
+    for file_key in tqdm.tqdm(random_sample_pair_test_indirect.keys(), total=len(random_sample_pair_test_indirect.keys())):
+        i2t_rank = rankdata([predictions[str(file_key)+':'+str(key_pair[k])] for k in random_sample_pair_test_indirect[file_key]])
+        t2i_rank = rankdata([predictions[str(key_pair_reversed[key_pair[k]])+':'+str(key_pair[file_key])] for k in random_sample_pair_test_indirect[file_key]])
+
+        i2t_ranks.append(i2t_rank[-1])
+        t2i_ranks.append(t2i_rank[-1])
+        i2t_rank_dict[file_key] = i2t_rank
+
+    assert len(i2t_ranks) == len(t2i_ranks) == 1000
+
+    ndcg_score = compute_ndcg_score_per_mode(i2t_rank_dict, similarity_score_test_indirect, random_sample_pair_test_indirect, mode='indrect', split='test', k=200)
+
+    results_dict['indirect'] = {}
+    results_dict['indirect']['i2t rank'] = float(sum(i2t_ranks) / len(i2t_ranks))
+    results_dict['indirect']['t2i rank'] = float(sum(t2i_ranks) / len(t2i_ranks))
+    results_dict['indirect']['ndcg score'] = float(ndcg_score)
+    print(f"Random split, mode={mode} i2t rank: ", sum(i2t_ranks) / len(i2t_ranks))
+    print(f"Random split, mode={mode} t2i rank: ", sum(t2i_ranks) / len(t2i_ranks))
+    return results_dict
+
+# %%
+def main(json_file):
+    ### Setup
+    # os.environ['ROOT'] = os.path.dirname(os.path.realpath(__file__))
+
+    # %%
+
+    ### Load data
+    # if os.path.isfile(os.path.join(os.environ['ROOT'], json_file)): #'results_pair_dict.json')):
+    #     predictions_file_path = os.path.join(os.environ['ROOT'],  json_file) #'results_pair_dict.json')
+    # else:
+    #     predictions_file_path = os.path.join(os.environ['ROOT'], json_file) #'data/results_pair_dict.json')
+
+    # with open(predictions_file_path) as f:
+    #     predictions = json.load(f)
+    predictions = json_file
+
+    with open(os.path.join(os.environ['ROOT'], 'data/key_pair.json')) as f:
+        key_pair = json.load(f)
+
+    # key_pair_reversed = {v: k for k, v in key_pair.items()}
+
+    with open(os.path.join(os.environ['ROOT'], 'data/random_sample_test_direct_ids.json')) as f:
+        random_sample_pair_test_direct = json.load(f)
+
+    with open(os.path.join(os.environ['ROOT'], 'data/random_sample_test_indirect_ids.json')) as f:
+        random_sample_pair_test_indirect = json.load(f)
+
+    with open(os.path.join(os.environ['ROOT'], 'data/chatgpt_similarity_score_test_direct.json')) as f:
+        similarity_score_test_direct = json.load(f)
+
+    with open(os.path.join(os.environ['ROOT'], 'data/chatgpt_similarity_score_test_indirect.json')) as f:
+        similarity_score_test_indirect = json.load(f)
+
+    # %%
+    ### Compute scores
+
+    result_direct = get_score_direct(random_sample_pair_test_direct, predictions, key_pair, similarity_score_test_direct, k = 200)
+    result_indirect = get_score_indirect(random_sample_pair_test_indirect, predictions, key_pair, similarity_score_test_indirect, k = 200)
+    result_dict = {**result_direct, **result_indirect}
+    return result_dict
+# %%
+if __name__ == '__main__':
+    os.environ['ROOT'] = os.path.dirname(os.path.realpath(__file__))
+    st.title("Evaluation Server for Driving Hazard Prediction and Reasoning ")
+    st.image(os.path.join(os.environ['ROOT'],'data/preview_image.jpeg'))
+    st.divider()
+    result_text = ''
+    result_dict = {}
+    uploaded_files = None
+    json_file = None
+
+    uploaded_files = st.file_uploader("Upload All Result Files Here (results_pair_dict1.json, results_pair_dict2.json)", type=["csv"], accept_multiple_files=True)
+    if st.button('Run Evaluation with no upload files (using demo files)'):
+        json_file1 = load_json(os.path.join(os.environ['ROOT'], 'results_pair_dict1.json'))
+        json_file2 = load_json(os.path.join(os.environ['ROOT'], 'results_pair_dict2.json'))
+        json_file = {**json_file1, **json_file2}
+    dataframe = pd.DataFrame([])
+    if uploaded_files is not None:
+        for i in range(len(uploaded_files)):
+            dataframe = pd.concat([dataframe, pd.read_csv(uploaded_files[i])])
+
+        result = dataframe.to_dict('tight')['data']
+        json_file = {}
+        for i in range(len(result)):
+            json_file[str(result[i][1])] = float(result[i][2])
+
+    if len(json_file) >= 1:
+        result_dict = main(json_file)
+        result_text = json.dumps(result_dict)
+
+    st.download_button('Download Results', result_text)
+    st.json(result_dict)
+
+    # !streamlit run app.py --server.fileWatcherType none
+
+    # if st.button('Load Results File1 from local instead'):
+    #     json_file1_path = os.path.join(os.environ['ROOT'], 'results_pair_dict1.json')
+    #     json_file1 = load_json(json_file1_path)
+    #     st.write(json_file1)
+
+    # if uploaded_files is not None:
+    #     with open(uploaded_file1) as jf:
+    #         json_file1 = json.load(jf)
+    #     json_file1 = load_json(uploaded_file1)
+    #     uploaded1 = True
+
+    # uploaded_file2 = st.file_uploader("Upload Results File2")
+    # if st.button('Load Results File2 from local instead'):
+    #     json_file2_path = os.path.join(os.environ['ROOT'], 'results_pair_dict2.json')
+    #     json_file2 = load_json(json_file2_path)
+    #     st.write(json_file2)
+
+    # if uploaded_file2 is not None:
+    #     with open(uploaded_file2) as jf:
+    #         json_file2 = json.load(jf)
+    #     uploaded2 = True
+
+    # # if uploaded1 and uploaded2:
+    # json_file = {**json_file1, **json_file2}
+
+
+
+# %%

config/examples.py

@@ -0,0 +1,16 @@
+import os
+import sys
+import json
+
+from omegaconf import OmegaConf
+from hydra.core.global_hydra import GlobalHydra
+from hydra import initialize, initialize_config_module, initialize_config_dir, compose
+
+os.environ['ROOT'] = "/home/quang/workspace/traffic_var"
+os.environ['DATA_ROOT'] = "/home/quang/datasets/traffic_var"
+
+# initialize hydra config
+GlobalHydra.instance().clear()
+initialize(config_path="./")
+with_decoder_config = compose(config_name='with_decoder.yaml', overrides=["clip_model=ViT-B/16", "rationale_type=0", "val_rationale_type=0"])
+original_config = compose(config_name='experiment.yaml', overrides=["clip_model=ViT-B/16", "rationale_type=0", "val_rationale_type=0"])

config/experiment.yaml

@@ -0,0 +1,41 @@
+# config.yaml
+
+exp_name: exp1-2
+wandb: 90788c79e1500570b08e5acf283e17df7e0c54b2
+root: "${oc.env:DATA_ROOT}"
+overfit: False
+batch_size: 4
+num_workers: 4
+img_size: 224
+rationale_type: 0 # 0 - only rationale, 1 - randomly add entity desc, 2 - add entity desc
+val_rationale_type: 0
+hide_true_bbox: 8 # clues and inferences selected randomly
+widescreen_processing: 1 # 0 - no widescreen, 1 - widescreen
+h_flip: False
+ema_decay: 0.9999
+
+clip_model: 'ViT-B/16' # 'RN101' 'RN50x4''RN50x16' 'RN50x64' 'ViT-L/14@336px' 'ViT-B/32'
+num_layers: 3
+dim_hidden: 512
+
+
+warmup: 1000
+init_from: ''
+lr: .00001
+n_epochs: 15
+save_every: 0
+early_stop: 5
+val_stat: 'loss'
+device: 'cuda'
+use_multi: False
+local_rank: 0
+
+hydra:
+  run:
+    dir: ./results/${exp_name}
+  output_subdir: ./ # directory for saving the yaml configs
+  job:
+    config:
+      override_dirname:
+        exclude_keys:
+          - exp.name

config/with_decoder.yaml

@@ -0,0 +1,60 @@
+# config.yaml
+
+exp_name: exp1-2
+wandb: 90788c79e1500570b08e5acf283e17df7e0c54b2
+root: "${oc.env:DATA_ROOT}"
+overfit: False
+batch_size: 4
+num_workers: 4
+img_size: 224
+rationale_type: 0 # 0 - only rationale, 1 - randomly add entity desc, 2 - add entity desc
+val_rationale_type: 0
+hide_true_bbox: 8 # clues and inferences selected randomly
+widescreen_processing: 1 # 0 - no widescreen, 1 - widescreen
+h_flip: False
+ema_decay: 0.9999
+aux_weight: 0.2
+no_hard_negative_itm: False
+
+clip_model: 'ViT-B/16' # 'RN101' 'RN50x4''RN50x16' 'RN50x64' 'ViT-L/14@336px' 'ViT-B/32'
+has_extra_txt_decoder: False
+has_extra_img_decoder: False
+has_extra_mix_decoder: False
+has_extra_gen_decoder: False
+
+extra_decoder:
+  is_decoder: True
+  vocab_size: 1000
+  d_ff: 512
+  d_kv: 64
+  d_model: 512
+  dropout_rate: 0.1
+  num_heads: 8
+  num_layers: 2
+  # eos_token_id: 1
+  # pad_token_id: 0
+  # decoder_start_token_id: 0
+  # n_positions: 512
+  relative_attention_max_distance: 128
+  relative_attention_num_buckets: 32
+
+warmup: 1000
+init_from: ''
+lr: .00001
+n_epochs: 15
+save_every: 0
+early_stop: 5
+val_stat: 'loss'
+device: 'cuda'
+use_multi: False
+local_rank: 0
+
+hydra:
+  run:
+    dir: ./results/${exp_name}
+  output_subdir: ./ # directory for saving the yaml configs
+  job:
+    config:
+      override_dirname:
+        exclude_keys:
+          - exp.name

create_result_script.py

@@ -0,0 +1,531 @@
+# %%
+
+import json
+import sys
+import pickle
+sys.path.append("../")
+import collections
+from models.fused_model import Model
+import os
+import tqdm
+import time
+import json
+import random
+from PIL import ImageFile
+from PIL import Image, ImageDraw
+import clip
+import torch
+import numpy as np
+import torchvision.transforms as T
+import torchvision.transforms.functional as F
+from pathlib import Path
+import pandas as pd
+
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
+# %%
+from types import SimpleNamespace
+# get config
+import os
+from omegaconf import OmegaConf
+from hydra.core.global_hydra import GlobalHydra
+from hydra import initialize, initialize_config_module, initialize_config_dir, compose
+
+os.environ['ROOT'] = os.path.dirname(os.path.realpath(__file__))
+os.environ['DATA_ROOT'] = os.path.join(os.environ['ROOT'], 'data')
+
+# initialize hydra config
+GlobalHydra.instance().clear()
+initialize(config_path="./config")
+
+config = compose(config_name='with_decoder.yaml',
+                 overrides=["clip_model=ViT-L/14@336px", 
+                            "rationale_type=0", "val_rationale_type=0"])
+
+class SquarePad:
+
+    def __call__(self, image):
+        max_wh = max(image.size)
+        p_left, p_top = [(max_wh - s) // 2 for s in image.size]
+        p_right, p_bottom = [max_wh - (s + pad) for s, pad in zip(image.size, [p_left, p_top])]
+        padding = (p_left, p_top, p_right, p_bottom)
+        return F.pad(image, padding, 0, 'constant')
+
+class VarDatasetForAuxEncoders:
+    def __init__(self, config, file_path, split="train", mode="combined", do_swap=False, tensorize=True, do_crop=True):
+        self.config = config
+        self.mode = mode
+        self.split = split
+        self.do_swap = do_swap
+        self.rationale_type = config.rationale_type if split == "train" else config.val_rationale_type
+        self.root_path = Path(config.root)
+        self.anno_path = file_path #self.root_path / f'annotations/13_05/anno_{split}_{mode}.json'
+        if split == "test" and mode == "combined" and config.overfit:
+            self.anno_path = self.root_path / f'annotations/13_05/anno_{split}_{mode}_overfit.json'
+
+        self.data = json.load(open(self.anno_path))
+        self.idx2name = list(self.data.keys())
+
+        if 'bounding_box' in self.data[list(self.data.keys())[0]]['details'][-1]:
+            self.one_ent_keys = [k for k, v in self.data.items() if len(v['details'][-1]["bounding_box"]) == 1]
+            self.two_ent_keys = [k for k, v in self.data.items() if len(v['details'][-1]["bounding_box"]) == 2]
+            self.three_ent_keys = [k for k, v in self.data.items() if len(v['details'][-1]["bounding_box"]) == 3]
+            self.all_ent_keys = self.one_ent_keys + self.two_ent_keys + self.three_ent_keys
+
+            self.keys = {1: self.one_ent_keys, 2: self.two_ent_keys, 3: self.three_ent_keys}
+
+        if self.config.widescreen_processing in [0, 1]:
+            self.resize_crop = self.get_transform(config.img_size, split == "train", padding=False)
+        else:
+            self.resize_crop = self.get_transform(config.img_size, split == "train", padding=True)
+
+        self.tensorize = tensorize
+        self.jitter_transform = T.ColorJitter(brightness=.5, hue=.3, saturation=.3) if split == "train" else lambda x: x
+
+        self.final_transform = T.Compose([
+            lambda image: image.convert("RGB"),
+            T.ToTensor() if tensorize else lambda x: x,
+            T.Normalize(
+                (0.48145466, 0.4578275, 0.40821073),
+                (0.26862954, 0.26130258, 0.27577711),
+            ) if tensorize else lambda x: x
+        ])
+
+    def get_transform(self, n_px, training, padding=False):
+        resize = T.Resize((n_px + 16, n_px + 16), interpolation=Image.BICUBIC)
+
+        # for traning split
+        if training and not padding:  # train
+            return T.Compose([resize, T.RandomCrop(n_px)])
+
+        if training and padding:  # train_pad
+            return T.Compose([SquarePad(), resize, T.RandomCrop(n_px)])
+
+        # for test and val split
+        if not training and not padding:  # test
+            return T.Compose([resize, T.CenterCrop(n_px)])
+
+        if not training and padding:  # test_pad
+            return T.Compose([SquarePad(), resize, T.CenterCrop(n_px)])
+
+    def key2img_path(self, key):
+        file_paths = [
+            self.root_path / f"var_images/{key}.jpg",
+            self.root_path / f"var_images/{key}.png",
+            self.root_path / f"images/{key}.jpg",
+            self.root_path / f"img/train/{key.split('_')[0]}/{key}.png",
+            self.root_path / f"img/val/{key.split('_')[0]}/{key}.png",
+            self.root_path / f"img/test/{key.split('_')[0]}/{key}.png",
+            self.root_path / f"img/{key}.png",
+            self.root_path / f"img/{key}.jpg",
+            self.root_path / f"images/{key}.png",
+            self.root_path / f"images/{key}.jpg",
+        ]
+        for file_path in file_paths:
+            if file_path.exists():
+                return file_path
+
+    def key2img(self, key):
+        file_path = self.key2img_path(key)
+        return Image.open(file_path)
+
+    def hide_region(self, image, bboxes):
+        image = image.convert('RGBA')
+
+        if self.config.hide_true_bbox == 1:  # hide mode
+            draw = ImageDraw.Draw(image, 'RGBA')
+
+        if self.config.hide_true_bbox in [2, 5, 7, 8, 9]:  #highlight mode
+            overlay = Image.new('RGBA', image.size, '#00000000')
+            draw = ImageDraw.Draw(overlay, 'RGBA')
+
+        if self.config.hide_true_bbox == 3 or self.config.hide_true_bbox == 6:  #blackout mode or position only mode
+            overlay = Image.new('RGBA', image.size, '#7B7575ff')
+            draw = ImageDraw.Draw(overlay, 'RGBA')
+
+        color_fill_list = ['#ff05cd3c', '#00F1E83c', '#F2D4003c']  # Green, Blue, Yellow?
+
+        for idx, bbox in enumerate(bboxes):
+            if bbox == None:
+                continue
+
+            color_fill = color_fill_list[idx]
+            x, y = bbox['left'], bbox['top']
+
+            if self.config.hide_true_bbox == 1:  # hide mode
+                draw.rectangle([(x, y), (x + bbox['width'], y + bbox['height'])], fill='#7B7575')
+            elif self.config.hide_true_bbox in [2, 5, 7, 8, 9]:  # highlight mode
+                draw.rectangle([(x, y), (x + bbox['width'], y + bbox['height'])], fill=color_fill, outline='#05ff37ff',
+                               width=3)  # Fill with Pink 60% ##00F1E8
+            elif self.config.hide_true_bbox == 3:  # blackout mode
+                draw.rectangle([(x, y), (x + bbox['width'], y + bbox['height'])], fill='#00000000')
+            elif self.config.hide_true_bbox == 6:  # position only mode
+                draw.rectangle([(x, y), (x + bbox['width'], y + bbox['height'])], fill=color_fill)
+
+        if self.config.hide_true_bbox in [2, 3, 5, 6, 7, 8, 9]:
+            image = Image.alpha_composite(image, overlay)
+
+        return image
+
+    def get_entity_codes(self):
+        entity_codes = [0, 1, 2]
+        if self.do_swap:
+            random.shuffle(entity_codes)
+        return entity_codes
+
+    def swap_entities(self, bboxes, text, entity_codes):
+        # text
+        for entity_idx, entity_code in enumerate(entity_codes):
+            text = text.replace(f"Entity #{entity_idx + 1}", f"Entity #{entity_code + 1}")
+
+        # bboxes: [1, 0, 2] -> [b[1], b[0], b[2]]
+        new_boxes = [bboxes[entity_code] for entity_code in entity_codes]
+        return new_boxes, text
+
+    def get_text_from_meta(self, meta):
+        n_boxes = len(meta['bounding_box'])  # key ['1', '2', '3']
+
+        # for rationale
+        text = 'Rationale: ' + str(meta['rationale'])
+
+        if self.rationale_type == 1 or self.rationale_type == 2:
+            for box_idx in range(n_boxes):
+                ent_name = f'Entity #{box_idx + 1}'
+                ent_desc = f'{ent_name}, {meta[ent_name]}'
+                # todo: replace randomly
+                text = text.replace(ent_name, ent_desc, 1)
+        return text
+
+    def get_itm_text(self, ori_file_key):
+        file_key = ori_file_key
+
+        if random.random() < 0.5:
+            n_boxes = len(self.data[file_key]['details'][-1]['bounding_box'])
+
+            file_key = random.choice(self.keys[n_boxes])
+
+            if self.config.get('no_hard_negative_itm', False):
+                file_key = random.choice(self.all_ent_keys)
+
+        itm_label = 1 if file_key == ori_file_key else 0
+        meta = self.data[file_key]['details'][-1]
+
+        itm_text = self.get_text_from_meta(meta)
+        return itm_text, itm_label
+
+    def get_bboxes_and_text(self, file_key, meta):
+        text = self.get_text_from_meta(meta)
+        bboxes = [meta['bounding_box'].get(str(box_idx + 1), None) for box_idx in range(3)]
+
+        entity_codes = self.get_entity_codes()
+        bboxes, text = self.swap_entities(bboxes, text, entity_codes)
+
+        itm_text, itm_label = self.get_itm_text(file_key)
+        _, itm_text = self.swap_entities([None, None, None], itm_text, entity_codes)
+        return {'bboxes': bboxes, 'text': text, 'itm_text': itm_text, 'itm_label': itm_label}
+
+    def get_image(self, file_key, bboxes):
+        image = self.key2img(file_key)
+        image = self.jitter_transform(image)
+        image = self.hide_region(image, bboxes)
+        image = self.final_transform(self.resize_crop(image))
+        return image
+
+    def __getitem__(self, idx):
+        file_key = self.idx2name[idx]
+
+        # Select the last version of label of the sample
+        meta = self.data[file_key]['details'][-1]
+
+        # read bboxes and rationale
+        outputs = self.get_bboxes_and_text(file_key, meta)
+        text = clip.tokenize(outputs['text'], truncate=True).squeeze()
+        itm_text = clip.tokenize(outputs['itm_text'], truncate=True).squeeze()
+        itm_label = torch.tensor(outputs['itm_label'])
+
+        image = self.get_image(file_key, outputs['bboxes'])
+
+        return {'image': image, 'caption': text, 'raw_text': text, 'file_key': file_key, 'itm_text': itm_text, 'itm_label': itm_label}
+
+    def __len__(self):
+        if self.config.overfit and not (self.split == 'test' and self.mode == 'combined'):
+            return 16
+        return len(self.data)
+
+# %%
+class VarDatasetImageOnly(VarDatasetForAuxEncoders):
+    def __init__(self, args, file_path, split="val", mode="combined", do_swap= False):
+        super().__init__(args, file_path, split=split, mode=mode, do_swap=do_swap)
+
+    def __getitem__(self, idx):
+        file_key = self.idx2name[idx]
+        meta = self.data[file_key]['details'][-1]
+        bboxes = [meta['bounding_box'].get(str(box_idx + 1), None) for box_idx in range(3)]
+        entity_codes = self.get_entity_codes()
+        bboxes = [bboxes[entity_code] for entity_code in entity_codes]
+        image = self.get_image(file_key, bboxes)
+        return {'image': image, 'file_key': file_key}
+
+# %%
+class VarDatasetTextOnly(VarDatasetForAuxEncoders):
+    def __init__(self, args, file_path, split="val", mode="combined", do_swap= False):
+        super().__init__(args, file_path, split=split, mode=mode, do_swap=do_swap)
+
+    def __getitem__(self, idx):
+        file_key = self.idx2name[idx]
+        meta = self.data[file_key]['details'][-1]
+        # text = self.get_text_from_meta(meta)
+        if 'Entity #3' in meta['hazard']:
+            n_boxes = 3 
+        elif 'Entity #2' in meta['hazard']:
+            n_boxes = 2
+        else:
+            n_boxes = 1
+
+        # for rationale
+        text = 'Rationale: ' + str(meta['hazard'])
+
+        if self.rationale_type == 1 or self.rationale_type == 2:
+            for box_idx in range(n_boxes):
+                ent_name = f'Entity #{box_idx + 1}'
+                ent_desc = f'{ent_name}, {meta[ent_name]}'
+                # todo: replace randomly
+                text = text.replace(ent_name, ent_desc, 1)
+
+        entity_codes = self.get_entity_codes()
+        for entity_idx, entity_code in enumerate(entity_codes):
+            text = text.replace(f"Entity #{entity_idx + 1}", f"Entity #{entity_code + 1}")
+        text = clip.tokenize(text, truncate=True).squeeze()
+        return {'caption': text,'file_key': file_key}
+
+# %%
+import os
+import sys
+
+sys.path.append('..')
+import json
+import fire
+import tqdm
+
+import clip
+import torch
+import sklearn
+import numpy as np
+
+from omegaconf import OmegaConf
+from models.fused_model import Model
+from torch.utils.data import DataLoader
+# from datasets import VarDatasetForAuxEncoders
+
+from scipy.stats import rankdata
+from sklearn.metrics import ndcg_score
+from sklearn.metrics import pairwise_distances
+
+
+# def get_data_loader(config, split="test", mode="combined", do_swap=False):
+#     dataset = VarDatasetForAuxEncoders(config, split=split, mode=mode, do_swap=do_swap)
+#     return DataLoader(dataset, batch_size=4, shuffle=False)
+
+def get_image_data_loader(config, file_path, split="test", mode="combined", do_swap=False):
+    dataset = VarDatasetImageOnly(config, file_path,   split=split, mode=mode, do_swap=do_swap)
+    return DataLoader(dataset, batch_size=4, shuffle=False)
+
+def get_text_data_loader(config, file_path,  split="test", mode="combined", do_swap=False):
+    dataset = VarDatasetTextOnly(config, file_path, split=split, mode=mode, do_swap=do_swap)
+    return DataLoader(dataset, batch_size=4, shuffle=False)
+
+# def get_data_loader(config, split="test", mode="combined", do_swap=False):
+#     dataset = VarDatasetForAuxEncoders(config, split=split, mode=mode, do_swap=do_swap)
+#     return DataLoader(dataset, batch_size=4, shuffle=False)
+
+def compute_rand_rank(split='test', mode='spec', img_token_dict={}, txt_token_dict={}):  # the dicts contain all 2000 test samples
+    data = json.load(open( os.path.join(os.environ['ROOT'], f"data/annotations/13_05/anno_random_{split}_{mode}_ids.json")))
+
+    i2t_ranks = []
+    t2i_ranks = []
+    i2t_rank_dict = {}
+    t2i_rank_dict = {}
+
+    for file_key in data.keys():
+        img_emb = (img_token_dict[file_key]).unsqueeze(0)
+        txt_emb = (txt_token_dict[file_key]).unsqueeze(0)
+
+        txt_embs = torch.stack([txt_token_dict[k] for k in data[file_key]])
+        img_embs = torch.stack([img_token_dict[k] for k in data[file_key]])
+        assert txt_embs.shape[0] == img_embs.shape[0] == 1000
+
+        i2t_rank = rankdata(pairwise_distances(img_emb, txt_embs, metric='cosine', n_jobs=8), axis=1)[0]
+        t2i_rank = rankdata(pairwise_distances(txt_emb, img_embs, metric='cosine', n_jobs=8), axis=1)[0]
+
+        i2t_ranks.append(i2t_rank[-1])
+        t2i_ranks.append(t2i_rank[-1])
+
+        i2t_rank_dict[file_key] = i2t_rank
+        t2i_rank_dict[file_key] = t2i_rank
+
+    assert len(i2t_ranks) == len(t2i_ranks) == 1000
+    print(f"Random split, mode={mode} i2t rank: ", sum(i2t_ranks) / len(i2t_ranks))
+    print(f"Random split, mode={mode} t2i rank: ", sum(t2i_ranks) / len(t2i_ranks))
+    # for k in i2t_rank_dict.keys():
+    #     print(k, i2t_rank_dict[k])
+    #     print('------------------')
+    #     break
+    return i2t_rank_dict  # for computing the NDCG scores
+
+
+def read_relevance_scores(anno_path="anno_random_test_obvi_ids.json", gpt_path="chatgpt_similarity_score_test_direct_combined.json"):
+    gpt_scores = json.load(open(gpt_path))
+    data = json.load(open(anno_path))
+
+    # add_missing_relevance_scores
+    for k in tqdm.tqdm(data, total=len(data)):
+        cand_keys = data[k]
+        for cand_key in cand_keys:
+            if cand_key not in gpt_scores[k]:
+                gpt_scores[k][cand_key] = 0.0
+            if cand_key == k:
+                gpt_scores[k][cand_key] = 1.0
+
+    return gpt_scores
+# %%
+
+def compute_ndcg(ranks, scores, k=3):
+    """
+    ranks = [5, 1, 4, 2, 3]
+    scores = [0.1, 0.5, 0.3, 0.95, 1.0]
+    """
+    rank_score_tuple = list(zip(ranks, scores))
+
+    top_k = sorted(rank_score_tuple, key=lambda x: x[1], reverse=True)[:k]
+
+    dcg = sum([score / np.log2(rank + 1) for rank, score in top_k])
+
+    ideal_dcg = sum([score / np.log2(idx + 2) for idx, (_, score) in enumerate(top_k)])
+
+    ndcg = dcg / ideal_dcg
+    return ndcg
+
+
+def compute_ndcg_score_per_mode(pred_rank_dict, gpt_rel_scores, mode='spec', split='test', k=200):
+    data = json.load(open(os.path.join(os.environ['ROOT'],f"data/annotations/13_05/anno_random_{split}_{mode}_ids.json")))
+
+    ndcg_scores = []
+
+    for key in tqdm.tqdm(pred_rank_dict.keys(), total=len(pred_rank_dict.keys())):
+        gpt_scores_for_key = [gpt_rel_scores[key][cand_key] for cand_key in data[key]]
+        pred_rank_for_key = pred_rank_dict[key]
+
+        ndcg_score = compute_ndcg(pred_rank_for_key, gpt_scores_for_key, k=k)
+        ndcg_scores.append(ndcg_score)
+
+    avg_ndcg_score = sum(ndcg_scores) / len(ndcg_scores)
+    print(f"Random split, mode={mode} ndcg score: ", avg_ndcg_score)
+    return avg_ndcg_score
+
+
+# %%
+def main():
+    # %%
+    ## Load Model
+    config_path= os.path.join(os.environ['ROOT'],"results/config.yaml")
+    model_path= os.path.join(os.environ['ROOT'],"results/model_epoch3.pth")
+    # %%
+    print("Loading config from:", config_path)
+    config = OmegaConf.load(config_path)
+    #print(OmegaConf.to_yaml(config))
+    # %%
+
+    # load checkpoint
+    checkpoint = torch.load(model_path, map_location=torch.device('cpu'))
+    print("Loaded model from:", model_path)
+
+    clip_model, _ = clip.load(config.clip_model, jit=False)
+    model = Model(clip_model, config)
+    model.load_state_dict(checkpoint['model_state_dict'])
+
+    model = model.to(config.device)
+    model = model.eval()
+    model = model.float()
+    logit_scale = model.clip_model.logit_scale.exp()
+
+    image_path = os.path.join(os.environ['ROOT'], "data/eval_test_image.json")
+    text_path = os.path.join(os.environ['ROOT'], "data/eval_test_text.json")
+
+    data_loader_image = get_image_data_loader(config, image_path, split='test', mode='combined' )
+    data_loader_text = get_text_data_loader(config, text_path, split='test', mode='combined' )
+
+    # %%
+    key_text_dict = {}
+    text_tensor_embedding = None
+    with torch.no_grad():
+        for i, d in tqdm.tqdm(enumerate(data_loader_text), total=len(data_loader_text)):
+            # print("d", d['file_key'])
+
+            # with torch.amp.autocast(device_type=config.device, dtype=torch.float16):
+            text_tensor_out, text_cls_out = model.var_txt_forward(d['caption'].to(config.device))
+            #print("text_tensor_out", text_tensor_out[0].shape)
+
+            if text_tensor_embedding == None:
+                text_tensor_embedding = text_cls_out
+            else:
+                text_tensor_embedding = torch.cat((text_tensor_embedding, text_cls_out), 0)
+            
+            for j,key in enumerate(d['file_key']):
+                key_text_dict[key] = int(i*len(d['file_key']) +j)
+
+    # %%
+    key_image_dict = {}
+    image_tensor_embedding = None
+    with torch.no_grad():
+        for i, d in tqdm.tqdm(enumerate(data_loader_image), total=len(data_loader_image)):
+            image_tensor_out, img_cls_out = model.var_img_forward(d['image'].to(config.device))
+
+            if image_tensor_embedding == None:
+                image_tensor_embedding = img_cls_out
+            else:
+                image_tensor_embedding = torch.cat((image_tensor_embedding, img_cls_out), 0)
+            
+            for j,key in enumerate(d['file_key']):
+                key_image_dict[key] = int(i*len(d['file_key']) +j)
+
+    idx2img = {idx: k for idx, k in enumerate(key_image_dict)}
+    idx2text = {idx: k for idx, k in enumerate(key_text_dict)}
+    # %%
+    image_tensor_embedding = image_tensor_embedding.to('cpu')
+    text_tensor_embedding = text_tensor_embedding.to('cpu')
+
+    # %% 
+    similarity_matrix = pairwise_distances(image_tensor_embedding, text_tensor_embedding, metric='cosine', n_jobs=8)
+
+    # %%
+    results_pair_dict = {}
+    ## put into matrix
+    for i in range (2000):
+        for j in range (2000):
+            results_pair_dict[str(idx2img[i])+':'+str(idx2text[j])] = float(similarity_matrix[i][j])
+
+    # %%
+    results_pair_dict1 = {}
+    results_pair_dict2 = {}
+    len_ = int(len(results_pair_dict)/2)
+    for j, key in enumerate(results_pair_dict):
+        if j <= len_:
+            results_pair_dict1[key] = results_pair_dict[key]
+        else:
+            results_pair_dict2[key] = results_pair_dict[key]
+
+    # %%
+    # with open(os.path.join(os.environ['ROOT'],'results_pair_dict1.json'), 'w', encoding='utf-8') as f:
+    #     json.dump(results_pair_dict1, f, ensure_ascii=False, indent=4)
+    # with open(os.path.join(os.environ['ROOT'],'results_pair_dict2.json'), 'w', encoding='utf-8') as f:
+    #     json.dump(results_pair_dict2, f, ensure_ascii=False, indent=4)
+    df = pd.DataFrame(results_pair_dict1.items(), columns=['key_pair','score'])
+    df.to_csv(os.path.join(os.environ['ROOT'],'results_pair_dict1.csv'))
+    df = pd.DataFrame(results_pair_dict2.items(), columns=['key_pair','score'])
+    df.to_csv(os.path.join(os.environ['ROOT'],'results_pair_dict2.csv'))
+
+# %%
+if __name__ == "__main__":
+    main()
+
+# %%

data/random_sample_test_direct_ids.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:86c25ef7d12166f6c27fe725004a8857ddfcd4dbc8cfafd14722c189c13efbf5
+size 27110039

data/random_sample_test_indirect_ids.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:74851fc1e326e6c838b868062949f2c3daa6a1c3c35f6dae87549caf86acd39d
+size 27109999

models/fused_model.py

@@ -0,0 +1,544 @@
+import copy
+import math
+from typing import Optional, Tuple, Union
+import numpy as np
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torch.nn import CrossEntropyLoss
+
+from transformers.activations import ACT2FN
+from einops import rearrange
+from transformers.models.t5.configuration_t5 import T5Config
+from transformers.modeling_utils import ModuleUtilsMixin
+
+from einops import rearrange, reduce
+
+
+class FeedForward(nn.Module):
+
+    def __init__(self, config: T5Config):
+        super().__init__()
+        self.wi = nn.Linear(config.d_model, config.d_ff, bias=False)
+        self.wo = nn.Linear(config.d_ff, config.d_model, bias=False)
+        self.dropout = nn.Dropout(config.dropout_rate)
+        self.act = ACT2FN["gelu"]
+        self.layer_norm = nn.LayerNorm(config.d_model)
+
+    def forward(self, x):
+        x_hidden = self.wo(self.dropout(self.act(self.wi(self.layer_norm(x)))))
+        return x + self.dropout(x_hidden)
+
+
+class Attention(nn.Module):
+
+    def __init__(self, config: T5Config, has_relative_attention_bias=False):
+        super().__init__()
+        self.is_decoder = config.is_decoder
+        self.has_relative_attention_bias = has_relative_attention_bias
+        self.relative_attention_num_buckets = config.relative_attention_num_buckets
+        self.relative_attention_max_distance = config.relative_attention_max_distance
+        self.d_model = config.d_model
+        self.key_value_proj_dim = config.d_kv
+        self.n_heads = config.num_heads
+        self.dropout = config.dropout_rate
+        self.inner_dim = self.n_heads * self.key_value_proj_dim
+
+        # Mesh TensorFlow initialization to avoid scaling before softmax
+        self.q = nn.Linear(self.d_model, self.inner_dim, bias=False)
+        self.k = nn.Linear(self.d_model, self.inner_dim, bias=False)
+        self.v = nn.Linear(self.d_model, self.inner_dim, bias=False)
+        self.o = nn.Linear(self.inner_dim, self.d_model, bias=False)
+
+        if self.has_relative_attention_bias:
+            self.relative_attention_bias = nn.Embedding(self.relative_attention_num_buckets, self.n_heads)
+
+    @staticmethod
+    def _relative_position_bucket(relative_position, bidirectional=True, num_buckets=32, max_distance=128):
+        """
+        memory_position - query_position -> bucket_idx. 
+        If bidirectional=False, then positive relative positions are invalid. 
+        We use smaller buckets for small absolute relative_position 
+        and larger buckets for larger absolute relative_positions.
+        * All relative positions >=max_distance map to the same bucket. 
+        * All relative positions <=-max_distance map to the same bucket.
+        This should allow for more graceful generalization to longer sequences than the model has been trained on
+        Args:
+            relative_position: an int32 Tensor
+            bidirectional: a boolean - whether the attention is bidirectional
+            num_buckets: an integer
+            max_distance: an integer
+        Returns:
+            a Tensor with the same shape as relative_position, containing int32 values in the range [0, num_buckets)
+        """
+        relative_buckets = 0
+        if bidirectional:
+            num_buckets //= 2
+            relative_buckets += (relative_position > 0).to(torch.long) * num_buckets
+            relative_position = torch.abs(relative_position)
+        else:
+            relative_position = -torch.min(relative_position, torch.zeros_like(relative_position))
+        # now relative_position is in the range [0, inf)
+
+        # half of the buckets are for exact increments in positions
+        max_exact = num_buckets // 2
+        is_small = relative_position < max_exact
+
+        # The other half of the buckets are for logarithmically bigger bins in positions up to max_distance
+        relative_position_if_large = max_exact + (torch.log(relative_position.float() / max_exact) / math.log(max_distance / max_exact) *
+                                                  (num_buckets - max_exact)).to(torch.long)
+        relative_position_if_large = torch.min(relative_position_if_large, torch.full_like(relative_position_if_large, num_buckets - 1))
+
+        relative_buckets += torch.where(is_small, relative_position, relative_position_if_large)
+        return relative_buckets
+
+    def compute_bias(self, query_length, key_length, device=None):
+        """Compute binned relative position bias"""
+        if device is None:
+            device = self.relative_attention_bias.weight.device
+        context_position = torch.arange(query_length, dtype=torch.long, device=device)[:, None]
+        memory_position = torch.arange(key_length, dtype=torch.long, device=device)[None, :]
+        relative_position = memory_position - context_position  # shape (query_length, key_length)
+        relative_position_bucket = self._relative_position_bucket(
+            relative_position,  # shape (query_length, key_length)
+            bidirectional=(not self.is_decoder),
+            num_buckets=self.relative_attention_num_buckets,
+            max_distance=self.relative_attention_max_distance,
+        )
+        values = self.relative_attention_bias(relative_position_bucket)  # shape (query_length, key_length, num_heads)
+        values = values.permute([2, 0, 1]).unsqueeze(0)  # shape (1, num_heads, query_length, key_length)
+        return values
+
+    def forward(self, x, mask=None, x_kv=None, pos_bias=None):
+        """
+        Self-attention (if key_value_states is None) or attention over source sentence (provided by key_value_states).
+        """
+        # Input is (batch_size, seq_length, dim)
+        # Mask is (batch_size, key_length) (non-causal) or (batch_size, key_length, key_length)
+        batch_size, seq_length = x.shape[:2]
+
+        real_seq_length = seq_length
+        key_length = real_seq_length if x_kv is None else x_kv.shape[1]
+
+        reshape = lambda states: rearrange(states, 'b s (h d) -> b h s d', h=self.n_heads)
+        unshape = lambda states: rearrange(states, 'b h s d -> b s (h d)')
+
+        q = reshape(self.q(x))  # (batch_size, n_heads, seq_length, dim_per_head)
+        k = reshape(self.k(x if x_kv is None else x_kv))
+        v = reshape(self.v(x if x_kv is None else x_kv))
+
+        # compute scores
+        scores = torch.matmul(q, k.transpose(3, 2))
+
+        if pos_bias is None:
+            if not self.has_relative_attention_bias:
+                pos_bias = torch.zeros((1, self.n_heads, real_seq_length, key_length), device=scores.device, dtype=scores.dtype)
+            else:
+                pos_bias = self.compute_bias(real_seq_length, key_length, device=scores.device)
+
+            if mask is not None:
+                pos_bias = pos_bias + mask  # (batch_size, n_heads, seq_length, key_length)
+
+        position_bias_masked = pos_bias
+        scores += position_bias_masked
+        attn_weights = F.softmax(scores.float(), dim=-1).type_as(scores)  # (B, H, seq_length, key_length)
+        attn_weights = F.dropout(attn_weights, p=self.dropout, training=self.training)  # (B, H, seq_length, key_length)
+
+        attn_output = unshape(torch.matmul(attn_weights, v))  # (batch_size, seq_length, dim)
+        attn_output = self.o(attn_output)
+        return (attn_output, pos_bias)
+
+
+class LayerSelfAttention(nn.Module):
+
+    def __init__(self, config, has_relative_attention_bias=False):
+        super().__init__()
+        self.SelfAttention = Attention(config, has_relative_attention_bias=has_relative_attention_bias)
+        self.layer_norm = nn.LayerNorm(config.d_model)
+        self.dropout = nn.Dropout(config.dropout_rate)
+
+    def forward(self, x, mask=None, pos_bias=None):  # x + drop(attn(ln(x)))
+        h = self.layer_norm(x)
+        outputs = self.SelfAttention(h, mask=mask, pos_bias=pos_bias)
+        x = x + self.dropout(outputs[0])
+        return (x, outputs[1])  # outputs[1] is pos_bias
+
+
+class LayerCrossAttention(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        self.EncDecAttention = Attention(config, has_relative_attention_bias=False)
+        self.layer_norm = nn.LayerNorm(config.d_model)
+        self.dropout = nn.Dropout(config.dropout_rate)
+
+    def forward(self, x, x_kv, mask=None, pos_bias=None):  # x + drop(attn(ln(x), x_kv))
+        h = self.layer_norm(x)
+
+        outputs = self.EncDecAttention(h, mask=mask, x_kv=x_kv, pos_bias=pos_bias)
+        x = x + self.dropout(outputs[0])
+        return (x, outputs[1])  # outputs[1] is pos_bias
+
+
+class Block(nn.Module):
+
+    def __init__(self, config, has_relative_attention_bias=False):
+        super().__init__()
+        self.is_decoder = config.is_decoder
+        self.layer = nn.ModuleList()
+        self.layer.append(LayerSelfAttention(config, has_relative_attention_bias=has_relative_attention_bias))
+        if self.is_decoder:
+            self.layer.append(LayerCrossAttention(config))
+
+        self.layer.append(FeedForward(config))
+
+    def forward(self, x, mask=None, pos_bias=None, context=None, context_mask=None, context_pos_bias=None):
+
+        self_attention_outputs = self.layer[0](x, mask=mask, pos_bias=pos_bias)
+        hidden_states = self_attention_outputs[0]
+
+        do_cross_attention = self.is_decoder and context is not None
+        if do_cross_attention:
+
+            cross_attention_outputs = self.layer[1](
+                hidden_states,
+                x_kv=context,
+                mask=context_mask,
+                pos_bias=context_pos_bias,
+            )
+            hidden_states = cross_attention_outputs[0]
+
+        # Apply Feed Forward layer
+        hidden_states = self.layer[-1](hidden_states)
+
+        pos_bias = self_attention_outputs[1]
+        context_pos_bias = cross_attention_outputs[1] if do_cross_attention else None
+
+        return (hidden_states, pos_bias, context_pos_bias)
+
+
+class Stack(nn.Module):
+
+    def __init__(self, config, is_decoder=True, has_embedding=False, generate_causal_mask=False):
+        super().__init__()
+        self.config = config
+        if has_embedding:
+            self.embed_tokens = nn.Embedding(config.vocab_size, config.d_model)
+
+        self.is_decoder = is_decoder
+        self.dtype = torch.float32
+
+        self.generate_causal_mask = generate_causal_mask
+
+        self.block = nn.ModuleList([Block(config, has_relative_attention_bias=bool(i == 0)) for i in range(config.num_layers)])
+        self.final_layer_norm = nn.LayerNorm(config.d_model)
+        self.dropout = nn.Dropout(config.dropout_rate)
+
+    def forward(
+        self,
+        input_ids=None,
+        dec_hidden_states=None,
+        enc_hidden_states=None,
+        dec_attention_mask=None,
+        enc_attention_mask=None,
+    ):
+        input_shape = input_ids.size() if input_ids is not None else dec_hidden_states.shape[:-1]
+        batch_size, seq_length = input_shape
+
+        if input_ids is not None:
+            input_ids = input_ids.view(-1, input_shape[-1])
+            inputs_embeds = self.embed_tokens(input_ids)
+        else:
+            inputs_embeds = dec_hidden_states
+
+        # required mask seq length can be calculated via length of past
+        mask_seq_length = seq_length
+
+        if dec_attention_mask is None:
+            dec_attention_mask = torch.ones(batch_size, mask_seq_length, device=inputs_embeds.device)
+        if self.is_decoder and enc_attention_mask is None and enc_hidden_states is not None:
+            encoder_seq_length = enc_hidden_states.shape[1]
+            enc_attention_mask = torch.ones(batch_size, encoder_seq_length, device=inputs_embeds.device, dtype=torch.long)
+
+        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
+        # ourselves in which case we just need to make it broadcastable to all heads.
+        extended_attention_mask = self.get_extended_attention_mask(dec_attention_mask, input_shape)
+
+        # If a 2D or 3D attention mask is provided for the cross-attention
+        # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
+        if self.is_decoder and enc_hidden_states is not None:
+            encoder_batch_size, encoder_sequence_length, _ = enc_hidden_states.size()
+            encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
+            if enc_attention_mask is None:
+                enc_attention_mask = torch.ones(encoder_hidden_shape, device=inputs_embeds.device)
+            encoder_extended_attention_mask = self.invert_attention_mask(enc_attention_mask)
+        else:
+            encoder_extended_attention_mask = None
+
+        pos_bias = None
+        context_pos_bias = None
+
+        hidden_states = self.dropout(inputs_embeds)
+
+        for i, layer_module in enumerate(self.block):
+
+            layer_outputs = layer_module(
+                hidden_states,
+                mask=extended_attention_mask,  # [1, 1, 1, 1 ] [B, L]
+                pos_bias=pos_bias,
+                context=enc_hidden_states,
+                context_mask=encoder_extended_attention_mask,
+                context_pos_bias=context_pos_bias,
+            )
+
+            # layer_outputs is a tuple with:
+            layer_outputs = layer_outputs[:1] + (None,) + layer_outputs[1:]
+
+            hidden_states, present_key_value_state = layer_outputs[:2]  # [B, L, D], None
+
+            # We share the position biases between the layers - the first layer store them
+            pos_bias = layer_outputs[2]  # [B, H, L, L]
+            if self.is_decoder and enc_hidden_states is not None:
+                context_pos_bias = layer_outputs[3]
+
+        hidden_states = self.final_layer_norm(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+
+        return (hidden_states,)
+
+    def invert_attention_mask(self, attention_mask):
+        """ 
+        Input:  1 for attend, 0 for masked/ignored
+        Output: 0 for attend, -1e30 for masked/ignored. 
+        Then we can add it to the attention logits.
+        [B, L]    -> [B, 1, 1, L]
+        [B, L, L] -> [B, 1, L, L]
+        """
+        if attention_mask.dim() == 3:
+            extended_attention_mask = attention_mask[:, None, :, :]
+        if attention_mask.dim() == 2:
+            extended_attention_mask = attention_mask[:, None, None, :]
+
+        extended_attention_mask = extended_attention_mask.to(dtype=self.dtype)  # fp16 compatibility
+        extended_attention_mask = (1.0 - extended_attention_mask) * torch.finfo(self.dtype).min
+
+        return extended_attention_mask
+
+    def get_extended_attention_mask(self, attention_mask, input_shape, device=None, dtype=None):
+        """
+        Makes broadcastable attention and causal masks so that future and masked tokens are ignored.
+        attention_mask: 1 for attend, 0 for masked/ignored
+        Return: The extended attention mask:  0 for attend, -1e30 for masked/ignored
+        [B, L]    -> [B, 1, 1, L]
+        [B, L, L] -> [B, 1, L, L]
+        """
+        dtype = dtype if dtype else attention_mask.dtype
+
+        # If input [B, query_length, key_length] -> [B, 1, query_length, key_length]
+        if attention_mask.dim() == 3:
+            extended_attention_mask = attention_mask[:, None, :, :]
+        elif attention_mask.dim() == 2:
+            # Provided a padding mask of dimensions [batch_size, seq_length]
+            # - if the model is a decoder, apply a causal mask in addition to the padding mask
+            # - if the model is an encoder, make the mask broadcastable to [batch_size, num_heads, seq_length, seq_length]
+            if self.config.is_decoder and self.generate_causal_mask:
+                extended_attention_mask = ModuleUtilsMixin.create_extended_attention_mask_for_decoder(input_shape, attention_mask, device)
+            else:
+                extended_attention_mask = attention_mask[:, None, None, :]
+        else:
+            raise ValueError(f"Wrong shape for input_ids (shape {input_shape}) or attention_mask (shape {attention_mask.shape})")
+
+        # Input:  valid = 1, padding = 0
+        # Output: valid = 0, padding = -1e30
+        #      => then we can add it to the attention logits
+        extended_attention_mask = extended_attention_mask.to(dtype=dtype)  # fp16 compatibility
+        extended_attention_mask = (1.0 - extended_attention_mask) * torch.finfo(dtype).min
+        return extended_attention_mask
+
+
+class Model(torch.nn.Module):
+
+    def __init__(self, clip_model, config):
+        super().__init__()
+        self.clip_model = clip_model
+        self.config = config
+
+        if self.config.has_extra_txt_decoder:
+            self.txt_decoder = Stack(config.extra_decoder)
+            self.itm_txt_head = torch.nn.Linear(config.extra_decoder.d_model, 2)
+
+        if self.config.has_extra_img_decoder:
+            self.img_decoder = Stack(config.extra_decoder)
+            self.itm_img_head = torch.nn.Linear(config.extra_decoder.d_model, 2)
+
+        if self.config.has_extra_mix_decoder:
+            self.mix_decoder = Stack(config.extra_decoder)
+            self.mix_itm_head = torch.nn.Linear(config.extra_decoder.d_model, 2)
+
+        if self.config.has_extra_gen_decoder:
+            self.gen_decoder = Stack(config.extra_decoder, has_embedding=True, generate_causal_mask=True)
+            self.gen_head = torch.nn.Linear(config.extra_decoder.d_model, config.vocab_size)
+
+        self.config = config
+
+    def img_forward(self, x: torch.Tensor):  # [N, 3, 224, 224]
+        x = self.clip_model.visual.conv1(x)  # shape = [*, width, grid, grid]
+        x = x.reshape(x.shape[0], x.shape[1], -1)  # shape = [*, width, grid ** 2]
+        x = x.permute(0, 2, 1)  # shape = [*, gri d ** 2, width]
+        x = torch.cat(
+            [self.clip_model.visual.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device), x],
+            dim=1)  # shape = [*, grid ** 2 + 1, width]
+        x = x + self.clip_model.visual.positional_embedding.to(x.dtype)
+        x = self.clip_model.visual.ln_pre(x)
+
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.clip_model.visual.transformer(x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        x = self.clip_model.visual.ln_post(x)  # [NLD]
+
+        if self.clip_model.visual.proj is not None:
+            proj = self.clip_model.visual.proj[None, :, :]
+            x = (x @ proj)
+
+        cls_token = x[:, 0, :]
+        return x, cls_token
+
+    def txt_forward(self, text):
+        dtype = self.clip_model.dtype
+        x = self.clip_model.token_embedding(text).type(dtype)  # [batch_size, n_ctx, d_model]
+
+        x = x + self.clip_model.positional_embedding.type(dtype)
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.clip_model.transformer(x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        x = self.clip_model.ln_final(x).type(dtype)
+
+        proj = self.clip_model.text_projection[None, :, :]
+        x = (x @ proj)
+
+        # take features from the eot embedding (eot_token is the highest number in each sequence)
+        eot = x[torch.arange(x.shape[0]), text.argmax(dim=-1)]
+        return x, eot  # [NLD]
+
+    def var_img_forward(self, image):
+        if len(image.shape) == 5:
+            img_features1, img_token1 = self.img_forward(image[:, 0, ...])
+            img_features2, img_token2 = self.img_forward(image[:, 1, ...])
+            img_token = (img_token1 + img_token2) / 2
+            img_features = (img_features1 + img_features2) / 2
+        else:
+            img_features, img_token = self.img_forward(image)
+        img_token = img_token / img_token.norm(dim=-1, keepdim=True)
+        return img_features, img_token
+
+    def var_txt_forward(self, text):
+        txt_features, txt_token = self.txt_forward(text)
+        txt_token = txt_token / txt_token.norm(dim=-1, keepdim=True)
+        return txt_features, txt_token
+
+    def get_device(self):
+        return next(self.parameters()).device
+
+    def get_features(self, image=None, text_ids=None):
+        outputs = {}
+        if image is not None:
+            img_features, img_token = self.var_img_forward(image)
+            outputs['img_features'] = img_features
+            outputs['img_token'] = img_token
+            outputs['img_mask'] = torch.ones_like(img_features[:, :, 0])
+        if text_ids is not None:
+            txt_features, txt_token = self.var_txt_forward(text_ids)
+            outputs['txt_features'] = txt_features
+            outputs['txt_token'] = txt_token
+            outputs['txt_mask'] = (text_ids != 0).to(txt_features.dtype)
+        return outputs
+
+    def get_prediction(self, img_features, txt_features, img_mask=None, txt_mask=None, decoder="txt_decoder", **kwargs):
+        outputs = {}
+        if decoder == 'txt_decoder':
+            hidden_states = self.txt_decoder(
+                dec_hidden_states=txt_features,
+                enc_hidden_states=img_features,
+                enc_attention_mask=img_mask,
+                dec_attention_mask=txt_mask,
+            )
+            outputs['itm_txt_logits'] = self.itm_txt_head(hidden_states[0][:, 0, :])
+            outputs['itm_txt_probs'] = torch.softmax(outputs['itm_txt_logits'], dim=-1)
+
+        if decoder == 'img_decoder':
+            hidden_states = self.img_decoder(
+                dec_hidden_states=img_features,
+                enc_hidden_states=txt_features,
+                enc_attention_mask=txt_mask,
+                dec_attention_mask=img_mask,
+            )
+            outputs['itm_img_logits'] = self.itm_img_head(hidden_states[0][:, 0, :])
+            outputs['itm_img_probs'] = torch.softmax(outputs['itm_img_logits'], dim=-1)
+        return outputs
+
+    def forward(self, image, text, itm_text=None, itm_labels=None, gen_inputs=None, gen_labels=None):  # , gen_inputs, gen_labels, **kwargs):
+        img_features, img_token = self.var_img_forward(image)
+        txt_features, txt_token = self.var_txt_forward(text)
+
+        itm_txt_features, _ = self.var_txt_forward(itm_text)
+        itm_txt_mask = (itm_text != 0).to(itm_txt_features.dtype)
+
+        outputs = dict(
+            img_token=img_token,
+            txt_token=txt_token,
+            img_features=img_features,
+            txt_features=txt_features,
+        )
+        if self.config.has_extra_txt_decoder and itm_text is not None:
+            itm_img_features = img_features
+            itm_txt_states = self.txt_decoder(
+                dec_hidden_states=itm_txt_features,
+                enc_hidden_states=itm_img_features,
+                enc_attention_mask=None,
+                dec_attention_mask=itm_txt_mask,
+            )
+            outputs['itm_txt_logits'] = self.itm_txt_head(itm_txt_states[0][:, 0])
+
+        if self.config.has_extra_img_decoder and itm_text is not None:
+            itm_img_features = img_features
+            itm_img_states = self.img_decoder(
+                dec_hidden_states=itm_img_features,
+                enc_hidden_states=itm_txt_features,
+                enc_attention_mask=itm_txt_mask,
+                dec_attention_mask=None,
+            )
+            outputs['itm_img_logits'] = self.itm_img_head(itm_img_states[0][:, 0])
+
+        if self.config.has_extra_mix_decoder:
+            pass
+
+        if self.config.has_extra_gen_decoder:
+            gen_features = self.gen_decoder(
+                input_ids=gen_inputs,
+                enc_hidden_states=img_features,
+                enc_attention_mask=None,
+                dec_attention_mask=None,
+                labels=gen_labels,
+            )
+            outputs['gen_logits'] = self.gen_head(gen_features[0])
+
+        return outputs
+
+
+if __name__ == "__main__":
+    import sys
+    from omegaconf import OmegaConf
+    sys.path.append("/home/quang/workspace/traffic_var")
+    from config.examples import with_decoder_config as config
+    config.has_extra_txt_decoder = True
+    print(OmegaConf.to_yaml(config))
+
+    import clip
+
+    def get_resolution(model):
+        return model.visual.input_resolution if hasattr(model, 'visual') else model.input_resolution
+
+    model, _ = clip.load(config.clip_model, jit=False, device="cpu")
+    config.img_size = get_resolution(model)
+    model = Model(model, config)

models/model.py

@@ -0,0 +1,117 @@
+import torch
+import torch.nn as nn
+from easydict import EasyDict as edict
+
+
+class Block(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+
+
+class Model(torch.nn.Module):
+
+    def __init__(self, clip_model, config):
+        super().__init__()
+        self.clip_model = clip_model
+        # if config.i2t_encoder_layers > 0:
+        #     self.i2t_encoder = nn.ModuleList([Block(config) for _ in range(config.i2t_encoder_layers)])
+
+        # if config.t2i_encoder_layers > 0:
+        #     self.t2i_encoder = nn.ModuleList([Block(config) for _ in range(config.i2t_encoder_layers)])
+
+        self.config = config
+
+    def img_forward(self, x: torch.Tensor):  # [N, 3, 224, 224]
+        x = self.clip_model.visual.conv1(x)  # shape = [*, width, grid, grid]
+        x = x.reshape(x.shape[0], x.shape[1], -1)  # shape = [*, width, grid ** 2]
+        x = x.permute(0, 2, 1)  # shape = [*, gri d ** 2, width]
+        x = torch.cat(
+            [self.clip_model.visual.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device), x],
+            dim=1)  # shape = [*, grid ** 2 + 1, width]
+        x = x + self.clip_model.visual.positional_embedding.to(x.dtype)
+        x = self.clip_model.visual.ln_pre(x)
+
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.clip_model.visual.transformer(x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        x = self.clip_model.visual.ln_post(x)  # [NLD]
+        cls_token = self.clip_model.visual.ln_post(x[:, 0, :])
+
+        if self.clip_model.visual.proj is not None:
+            cls_token = cls_token @ self.clip_model.visual.proj
+        return x, cls_token
+
+    def txt_forward(self, text):
+        dtype = self.clip_model.dtype
+        x = self.clip_model.token_embedding(text).type(dtype)  # [batch_size, n_ctx, d_model]
+
+        x = x + self.clip_model.positional_embedding.type(dtype)
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.clip_model.transformer(x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        x = self.clip_model.ln_final(x).type(dtype)
+
+        # take features from the eot embedding (eot_token is the highest number in each sequence)
+        eot = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.clip_model.text_projection
+        return x, eot  # [NLD]
+
+    def var_img_forward(self, image):
+        if len(image.shape) == 5:
+            img_features1, img_token1 = self.img_forward(image[:, 0, ...])
+            img_features2, img_token2 = self.img_forward(image[:, 1, ...])
+            img_token = (img_token1 + img_token2) / 2
+            img_features = (img_features1 + img_features2) / 2
+        else:
+            img_features, img_token = self.img_forward(image)
+        img_token = img_token / img_token.norm(dim=-1, keepdim=True)
+        return img_features, img_token
+
+    def var_txt_forward(self, text):
+        txt_features, txt_token = self.txt_forward(text)
+        txt_token = txt_token / txt_token.norm(dim=-1, keepdim=True)
+        return txt_features, txt_token
+
+    def forward(self, image, text, past_img_tokens=None, past_txt_tokens=None):
+        # TODO: aggregate past img and txt tokens
+        img_features, img_token = self.var_img_forward(image)
+        txt_features, txt_token = self.var_txt_forward(text)
+        logit_scale = self.clip_model.logit_scale.exp()
+
+        if past_img_tokens is not None:
+            past_img_tokens = torch.cat([past_img_tokens, img_token], dim=0)
+            past_txt_tokens = torch.cat([past_txt_tokens, txt_token], dim=0)
+
+            batch_size = past_img_tokens.shape[0]
+            ground_truth = torch.arange(batch_size, dtype=torch.long, device=img_token.device)
+
+            logits_for_imgs = logit_scale * past_img_tokens @ past_txt_tokens.t()
+            logits_for_txts = logits_for_imgs.t()
+            # print(f"past_img_tokens: {past_img_tokens.shape}, past_txt_tokens: {past_txt_tokens.shape}")
+
+            # CLIP Contrastive Learning Loss Function
+            loss_img = torch.nn.CrossEntropyLoss()
+            loss_txt = torch.nn.CrossEntropyLoss()
+            loss = (loss_img(logits_for_imgs, ground_truth[:batch_size]) + loss_txt(logits_for_txts, ground_truth[:batch_size])) / 2
+        else:
+            batch_size = img_token.shape[0]
+            ground_truth = torch.arange(batch_size, dtype=torch.long, device=img_token.device)
+
+            logits_for_imgs = logit_scale * img_token @ txt_token.t()
+            logits_for_txts = logits_for_imgs.t()
+
+            # CLIP Contrastive Learning Loss Function
+            loss_img = torch.nn.CrossEntropyLoss()
+            loss_txt = torch.nn.CrossEntropyLoss()
+            loss = (loss_img(logits_for_imgs, ground_truth[:batch_size]) + loss_txt(logits_for_txts, ground_truth[:batch_size])) / 2
+
+        return dict(
+            img_token=img_token,
+            txt_token=txt_token,
+            img_features=img_features,
+            txt_features=txt_features,
+            loss=loss,
+            past_img_tokens=past_img_tokens,
+            past_txt_tokens=past_txt_tokens,
+        )

results/config.yaml

@@ -0,0 +1,42 @@
+exp_name: expl5-3a-bs16-imgtxtdec-ema-aux1
+wandb: 90788c79e1500570b08e5acf283e17df7e0c54b2
+root: ${oc.env:DATA_ROOT}
+overfit: false
+batch_size: 16
+num_workers: 4
+img_size: 336
+rationale_type: 0
+val_rationale_type: 0
+hide_true_bbox: 8
+widescreen_processing: 1
+h_flip: false
+ema_decay: 0.9999
+aux_weight: 1.0
+no_hard_negative_itm: false
+clip_model: ViT-L/14@336px
+has_extra_txt_decoder: true
+has_extra_img_decoder: true
+has_extra_mix_decoder: false
+has_extra_gen_decoder: false
+extra_decoder:
+  is_decoder: true
+  vocab_size: 1000
+  d_ff: 768
+  d_kv: 64
+  d_model: 768
+  dropout_rate: 0.1
+  num_heads: 12
+  num_layers: 2
+  relative_attention_max_distance: 128
+  relative_attention_num_buckets: 32
+warmup: 1000
+init_from: ''
+lr: 1.0e-05
+n_epochs: 15
+save_every: 0
+early_stop: 5
+val_stat: loss
+device: cuda
+use_multi: false
+local_rank: 0
+run_file: /home/acd13872jh/workspace/traffic_var/results/.runinfo/.sh

results/model_epoch3.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a45011f16c569075fcfab80c3f0e6df273a825dbad2ebeca65a438a57a95977f
+size 251753353

results/results_pair_dict.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1d3fa8165cf47a2562a3aa534880df625caf3705a59b5f5b221ba40f37a524fb
+size 303815155

results_pair_dict1.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7a469b8b91ebec66d55cc3389731e3e9f7e11111f112f13a06ea2e668f573309
+size 151976036

results_pair_dict2.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:59e06a5e308a71fd97b3f2f87439fde8fbac2c2834896a3ac5c54868ec9c480f
+size 151839121