even cleaner code

Author: Ye Shaokai

action/ek_eval.py

@@ -18,7 +18,7 @@
 import json
 import logging
 from llava.utils import rank0_print
-from action.utils import generate_label_map, MultiChoiceGenerator, match_answer
+from action.utils import generate_label_map, MultiChoiceGenerator, match_answer, parse_avion_predictions
 
 def datetime2sec(str):
     hh, mm, ss = str.split(':')
@@ -377,17 +377,10 @@ def get_topk_predictions(data, idx,  k):
     options = list(range(26))[:k]
 
     predictions = data[str(idx)]['predictions'][:k]
-    new_predictions = []
-    for pred in predictions:
-        # the prediction looks like verb:noun1:noun2..
-        # we want to it look like verb noun1:noun2 
-        first_sep = pred.index(':')
-        prediction = pred[:first_sep] + ' ' + pred[first_sep+1:]
-        new_predictions.append(prediction)
-
-    predictions = new_predictions    
-    for i in range(len(options)):
-              
+
+    predictions = parse_avion_predictions(predictions)    
+
+    for i in range(len(options)):              
         options[i] = f'{letters[i]}. {predictions[i]}'
 
     mc_data = {

action/generate_description.py

@@ -2,7 +2,7 @@
 import csv
 import os
 import argparse
-from action.utils import generate_label_map, MultiChoiceGenerator
+from action.utils import generate_label_map, MultiChoiceGenerator, AvionMultiChoiceGenerator
 from pathlib import Path
 
 
@@ -12,7 +12,7 @@ def datetime2sec(str):
     hh, mm, ss = str.split(':')
     return int(hh) * 3600 + int(mm) * 60 + float(ss)
 
-def generate_train_ann(ann_file, verb_ids, noun_ids, gen_type = 'naive'):
+def generate_train_ann(ann_file, verb_ids, noun_ids, gen_type = 'naive', avion_prediction_path = ''):
     assert gen_type in GEN_TYPES
     # epic kitchen uses csv
     csv_reader = csv.reader(open(ann_file))
@@ -21,8 +21,12 @@ def generate_train_ann(ann_file, verb_ids, noun_ids, gen_type = 'naive'):
     ann_root = Path(ann_file).parent
     if gen_type == "random_mc":
         mc_generator = MultiChoiceGenerator(ann_root)
+    elif gen_type == 'avion_mc':
+        mc_generator = AvionMultiChoiceGenerator(ann_root)
+        with open(avion_prediction_path, 'r') as f:
+            avion_train_predictions = json.load(f)
 
-    for row in csv_reader:
+    for idx, row in enumerate(csv_reader):
         start_timestamp, end_timestamp = datetime2sec(row[4]), datetime2sec(row[5])
 
         pid, vid = row[1:3]
@@ -40,6 +44,14 @@ def generate_train_ann(ann_file, verb_ids, noun_ids, gen_type = 'naive'):
             gt_answer_letter = mc_data['gt_answer_letter'][0]
             gt_answer_name = mc_data['gt_answer_name'][0]
             conversation = generate_random_mc_conversation(options, gt_answer_letter, gt_answer_name )
+        elif gen_type == "avion_mc":
+            vn_str = f'{row[10]}:{row[12]}'
+            avion_preds = avion_train_predictions[str(idx)]['predictions']
+            mc_data = mc_generator.generate_multi_choice(vn_str, avion_preds, 5)
+            options = mc_data['option'][0]
+            gt_answer_letter = mc_data['gt_answer_letter'][0]
+            gt_answer_name = mc_data['gt_answer_name'][0]
+            conversation = generate_random_mc_conversation(options, gt_answer_letter, gt_answer_name )
 
         data = {'video': vid_path,
                 'conversations': conversation,
@@ -67,24 +79,38 @@ def generate_random_mc_conversation(options:list[str], gt_answer_letter, gt_answ
         {"from": "gpt", "value": f"{gt_answer_letter}. {gt_answer_name}"} 
     ]
 
+def generate_avion_mc_conversation():
+    pass
+
 
 def get_args():
     parser = argparse.ArgumentParser(description="For generating VQA for EPIC-KITCHEN")
     parser.add_argument('--train_metadata', default='/data/shaokai/epic-kitchens-100-annotations/EPIC_100_train.csv', type=str)
     parser.add_argument('--out_folder', default = '/data/shaokai/EK100_in_LLAVA/', type = str)
+    parser.add_argument('--avion_train_predictions', default = '/data/shaokai/avion_predictions_train.json', type = str)
+    parser.add_argument('--gen_type', default = 'avion_mc', type = str, choices = GEN_TYPES)
     return parser.parse_args()
 
-def main():    
+def main(): 
     args = get_args()    
     ann_file = args.train_metadata
-    inst_train_folder = args.out_folder
-    print (ann_file)
-    anno_path = Path(ann_file).parent
-    labels, mapping_vn2act, verb_ids, noun_ids = generate_label_map(anno_path)
-    conv_lst = generate_train_ann(ann_file, verb_ids, noun_ids, gen_type = 'random_mc')
-    
+    inst_train_folder = os.path.join(args.out_folder, args.gen_type)
+
+    print ('train_metadata', args.train_metadata)
+    print ('out_folder', args.out_folder)
+    print ('loading predictions from ', args.avion_train_predictions)
+    print ('gen_type is ', args.gen_type)
+
     os.makedirs(inst_train_folder, exist_ok=True)
 
+    anno_path = Path(ann_file).parent
+    _, _, verb_ids, noun_ids = generate_label_map(anno_path)
+    conv_lst = generate_train_ann(ann_file, 
+                                  verb_ids, 
+                                  noun_ids, 
+                                  gen_type = args.gen_type, 
+                                  avion_prediction_path = args.avion_train_predictions)
+        
     # save it to a jsonl
     with open(os.path.join(inst_train_folder,'train_convs_narration.jsonl'), 'w') as f:
         for conv in conv_lst:

action/utils.py

@@ -48,13 +48,25 @@ def match_answer(pred, gt):
     return pred.intersection(gt) == gt
 
 
+def parse_avion_predictions(predictions):
+    new_predictions = []
+    for pred in predictions:
+        # the prediction looks like verb:noun1:noun2..
+        # we want to it look like verb noun1:noun2 
+        first_sep = pred.index(':')
+        prediction = pred[:first_sep] + ' ' + pred[first_sep+1:]
+        new_predictions.append(prediction)
+    return new_predictions
+
+
 class MultiChoiceGenerator:
     """
     Generating multi choice
     """
     def __init__(self, ann_root):
         self.ann_root = ann_root
         _, self.mapping_vn2act, self.verb_maps, self.noun_maps = generate_label_map(ann_root)
+    
 
     def generate_multi_choice(self, gt_vn, k):
         """
@@ -98,12 +110,73 @@ def generate_multi_choice(self, gt_vn, k):
 
         return data
 
+class AvionMultiChoiceGenerator(MultiChoiceGenerator):
+    """
+    Generate multichoice using avion predictions
+    """
+    def __init__(self, ann_root):
+        super().__init__(ann_root)
+    
+    def generate_multi_choice(self, gt_vn, avion_predictions, k):
+        """
+        Generate k multiple choices from gt_vn pairs
+
+        randomly pick 1 letter for gt_vn
+        randomly pick k-1 letters from vn_list that is not gt_vn (this is important as avion_predictions can contain correct prediction)        
+
+        """    
+        gt_v_id, gt_n_id = gt_vn.split(':')
+        gt_v_name, gt_n_name = self.verb_maps[gt_v_id], self.noun_maps[gt_n_id]
+        gt_answer = f'{gt_v_name} {gt_n_name}'
+
+        letters = [chr(65+i) for i in range(26)][:k]
+        options = list(range(26))[:k]
+
+        # we should have plenty of predictions to select, so let's not always pick the hardest
+        assert len(avion_predictions) > 2*k
+        avion_predictions = avion_predictions[:k*2]
+        avion_predictions = parse_avion_predictions(avion_predictions)
+        if gt_answer in avion_predictions:
+            avion_predictions.remove(gt_answer)
+        # just so that it's not strictly desending with confidence
+        random.shuffle(avion_predictions)
+        avion_predictions = avion_predictions[:k-1]
+
+        answers = [gt_answer] + avion_predictions
+        random.shuffle(answers)
+
+        options = []
+        for answer, letter in zip(answers, letters):
+            options.append(f'{letter}. {answer}')
+
+        gt_letter = letters[answers.index(gt_answer)]
+        
+        data = {
+                'question': {0: 'the video is an egocentric view of a person. What is the person doing? Pick the the letter that has the correct answer'},
+                'option': {0: options},
+                # the correct letter in mc
+                # for inspecting
+                'gt_answer_letter': {0: gt_letter},
+                'gt_answer_name': {0: gt_answer},
+                'valid_letters': letters
+            }        
+        return data
+
 
 if __name__ == '__main__':
 
     anno_root = "/storage-rcp-pure/upmwmathis_scratch/shaokai/epic-kitchens-100-annotations/"
-    generator = MultiChoiceGenerator(anno_root)
+    #generator = MultiChoiceGenerator(anno_root)
+    generator = AvionMultiChoiceGenerator(anno_root)
+    import json
+
+    with open('/storage-rcp-pure/upmwmathis_scratch/shaokai/avion_predictions_train.json') as f:
+        predictions = json.load(f)
+
+    print (len(predictions))
+    print (predictions['0'])
+    print (len(predictions['0']['predictions']))
 
-    print (generator.generate_multi_choice('3:3',5))
+    print (generator.generate_multi_choice('3:3',  predictions['0']['predictions'],  5))
 
     pass

scripts/train/EK100.yaml

@@ -0,0 +1,3 @@
+datasets:
+  - json_path: /data/shaokai/EK100_inst_train/avion_mc/train_convs_narration.jsonl
+    sampling_strategy: all

scripts/train/EK100_avion_mc.yaml

@@ -0,0 +1,3 @@
+datasets:
+  - json_path: /data/shaokai/EK100_inst_train/random_mc/train_convs_narration.jsonl
+    sampling_strategy: all

scripts/train/EK100_random_mc.yaml

@@ -5,4 +5,7 @@
 
 python3 action/generate_description.py \
    --train_metadata /storage-rcp-pure/upmwmathis_scratch/shaokai/epic-kitchens-100-annotations/EPIC_100_train.csv \
-   --out_folder /storage-rcp-pure/upmwmathis_scratch/shaokai/EK100_avion_mc #  > train_gen.out 2>&1
+   --out_folder /storage-rcp-pure/upmwmathis_scratch/shaokai/EK100_inst_train \
+   --avion_train_predictions /storage-rcp-pure/upmwmathis_scratch/shaokai/avion_predictions_train.json \
+   --gen_type avion_mc \
+

shaokai_generate_train.sh

@@ -0,0 +1,68 @@
+#!/bin/bash
+
+# Export environment variables
+export CUDA_VISIBLE_DEVICES="0,1,2,3"
+export OMP_NUM_THREADS="8"
+export NCCL_IB_DISABLE="0"
+export NCCL_IB_GID_INDEX="3"
+export NCCL_SOCKET_IFNAME="eth0"
+export NCCL_DEBUG="INFO"
+export ACCELERATE_CPU_AFFINITY="1"
+export WANDB_API_KEY="4474ec79de023b0c3ffb43588ab6163264f875db"
+export HF_HOME=/data/shaokai
+
+
+# Run the command using torchrun
+torchrun --nproc_per_node=4 \
+         --nnodes=1 \
+         --node_rank=0 \
+         --master_addr=127.0.0.1 \
+         --master_port=29500 \
+         llava/train/train_mem.py \
+         --deepspeed scripts/zero3.json \
+         --model_name_or_path lmms-lab/llava-onevision-qwen2-0.5b-ov \
+         --version qwen_1_5 \
+         --data_path scripts/train/EK100_avion_mc.yaml \
+         --video_folder /data/shaokai/\
+         --mm_tunable_parts mm_vision_tower,mm_mlp_adapter,mm_language_model \
+         --mm_vision_tower_lr 2e-6 \
+         --vision_tower google/siglip-so400m-patch14-384 \
+         --mm_projector_type mlp2x_gelu \
+         --mm_vision_select_layer -2 \
+         --mm_use_im_start_end False \
+         --mm_use_im_patch_token False \
+         --group_by_modality_length True \
+         --image_aspect_ratio anyres_max_9 \
+         --image_grid_pinpoints "(1x1),...,(6x6)" \
+         --mm_patch_merge_type spatial_unpad \
+         --bf16 True \
+         --run_name shaokai_llama_ov_0.5b_avion_mc \
+         --output_dir experiments/shaokai_llama_ov_0.5b_avion_mc \
+         --num_train_epochs 1 \
+         --per_device_train_batch_size 1 \
+         --per_device_eval_batch_size 4 \
+         --gradient_accumulation_steps 2 \
+         --evaluation_strategy steps \
+         --eval_steps 500\
+         --save_strategy steps \
+         --save_steps 1000 \
+         --learning_rate 1e-5 \
+         --weight_decay 0. \
+         --warmup_ratio 0.03 \
+         --lr_scheduler_type cosine \
+         --logging_steps 1 \
+         --tf32 True \
+         --model_max_length 32768 \
+         --gradient_checkpointing True \
+         --dataloader_num_workers 4 \
+         --lazy_preprocess True \
+         --report_to wandb \
+         --torch_compile True \
+         --torch_compile_backend inductor \
+         --dataloader_drop_last True \
+         --frames_upbound 32 \
+         --root /data/shaokai/EK100 \
+         --action_predictions /data/shaokai/avaion_predictions_test.json \
+         --val_metadata /data/shaokai/epic-kitchens-100-annotations/EPIC_100_validation.csv \
+         --llava_num_frames 16 \
+         --topk_predictions 5 > train_kitchen_0.5b_avion_mc.out 2>&1

shaokai_train_avion_mc.sh

@@ -8,9 +8,7 @@ export NCCL_IB_GID_INDEX="3"
 export NCCL_SOCKET_IFNAME="eth0"
 export NCCL_DEBUG="INFO"
 export ACCELERATE_CPU_AFFINITY="1"
-# export LD_PRELOAD="/usr/lib/x86_64-linux-gnu/libffi.so.7"
 export WANDB_API_KEY="4474ec79de023b0c3ffb43588ab6163264f875db"
-experiment_name="shaokai_llama_ov_0.5b_debug"
 export HF_HOME=/data/shaokai
 
 
@@ -24,7 +22,7 @@ torchrun --nproc_per_node=4 \
          --deepspeed scripts/zero3.json \
          --model_name_or_path lmms-lab/llava-onevision-qwen2-0.5b-ov \
          --version qwen_1_5 \
-         --data_path scripts/train/EK100.yaml \
+         --data_path scripts/train/EK100_random_mc.yaml \
          --video_folder /data/shaokai/\
          --mm_tunable_parts mm_vision_tower,mm_mlp_adapter,mm_language_model \
          --mm_vision_tower_lr 2e-6 \
@@ -64,7 +62,7 @@ torchrun --nproc_per_node=4 \
          --dataloader_drop_last True \
          --frames_upbound 32 \
          --root /data/shaokai/EK100 \
-         --action_predictions /data/shaokai/avaion_predictions.json \
+         --action_predictions /data/shaokai/avaion_predictions_test.json \
          --val_metadata /data/shaokai/epic-kitchens-100-annotations/EPIC_100_validation.csv \
          --llava_num_frames 16 \
-         --topk_predictions 5 > train_kitchen_0.5b.out 2>&1
+         --topk_predictions 5 > train_kitchen_0.5b_random_mc.out 2>&1