Major Refactoring

Author: danparizher

evaluations.py

@@ -3,10 +3,12 @@
 '''
 
 import re
-from rouge import Rouge
+
 from nltk.translate.bleu_score import sentence_bleu
+from rouge import Rouge
 from sentence_transformers import util
 
+
 ########################
 ## BLEU
 ########################
@@ -43,9 +45,7 @@ def caculate_bleu(results, data, gram):
         bleu = bleu_score(target, prediction, gram)
         bleus.append(bleu)
 
-    avg_bleu = sum(bleus) / len(bleus)
-
-    return avg_bleu
+    return sum(bleus) / len(bleus)
 
 
 ########################
@@ -54,8 +54,7 @@ def caculate_bleu(results, data, gram):
 def score_rouge(str1, str2):
     rouge = Rouge(metrics=["rouge-l"])
     scores = rouge.get_scores(str1, str2, avg=True)
-    rouge_l = scores['rouge-l']['f']
-    return rouge_l
+    return scores['rouge-l']['f']
 
 
 def caculate_rouge(results, data):
@@ -71,8 +70,7 @@ def caculate_rouge(results, data):
         rouge = score_rouge(target, prediction)
         rouges.append(rouge)
 
-    avg_rouge = sum(rouges) / len(rouges)
-    return avg_rouge
+    return sum(rouges) / len(rouges)
 
 
 ########################
@@ -82,8 +80,7 @@ def similariry_score(str1, str2, model):
     # compute embedding for both lists
     embedding_1 = model.encode(str1, convert_to_tensor=True)
     embedding_2 = model.encode(str2, convert_to_tensor=True)
-    score = util.pytorch_cos_sim(embedding_1, embedding_2).item()
-    return score
+    return util.pytorch_cos_sim(embedding_1, embedding_2).item()
 
 
 def caculate_similariry(results, data, model):
@@ -96,5 +93,4 @@ def caculate_similariry(results, data, model):
         score = similariry_score(target, prediction, model)
         scores.append(score)
 
-    avg_score = sum(scores) / len(scores)
-    return avg_score
+    return sum(scores) / len(scores)

main.py

@@ -1,23 +1,25 @@
+import argparse
+import json
 import os
+import random
+import re
+
 import numpy as np
 import torch
-import os
-import re
-import json
-import argparse
-import random
-from transformers import T5Tokenizer, DataCollatorForSeq2Seq, Seq2SeqTrainingArguments, Seq2SeqTrainer, T5ForConditionalGeneration
-from model import T5ForConditionalGeneration, T5ForMultimodalGeneration
-from utils_data import img_shape, load_data_std, load_data_img, ScienceQADatasetStd, ScienceQADatasetImg
-from utils_prompt import *
-from utils_evaluate import get_scores
-from rich.table import Column, Table
 from rich import box
 from rich.console import Console
+from rich.table import Column, Table
+from transformers import (DataCollatorForSeq2Seq, Seq2SeqTrainer, Seq2SeqTrainingArguments, T5ForConditionalGeneration, T5Tokenizer)
+
+from model import T5ForConditionalGeneration, T5ForMultimodalGeneration
+from utils_data import (ScienceQADatasetImg, ScienceQADatasetStd, img_shape, load_data_img, load_data_std)
+from utils_evaluate import get_scores
+from utils_prompt import *
+
 console = Console(record=True)
-from torch import cuda
-import nltk
 import evaluate
+import nltk
+from torch import cuda
 
 
 def parse_args():
@@ -36,7 +38,7 @@ def parse_args():
     parser.add_argument('--train_split', type=str, default='train', choices=['train', 'trainval', 'minitrain'])
     parser.add_argument('--val_split', type=str, default='val', choices=['test', 'val', 'minival'])
     parser.add_argument('--test_split', type=str, default='test', choices=['test', 'minitest'])
-    
+
     parser.add_argument('--use_generate', action='store_true', help='only for baseline to improve inference speed')
     parser.add_argument('--final_eval', action='store_true', help='only evaluate the model at the final epoch')
     parser.add_argument('--user_msg', type=str, default="baseline", help='experiment type in the save_dir')
@@ -50,16 +52,15 @@ def parse_args():
                         choices=['QCM-A', 'QCM-LE', 'QCMG-A', 'QCM-LEA', 'QCM-ALE'])
     parser.add_argument('--seed', type=int, default=42, help='random seed')
 
-    args = parser.parse_args()
-    return args
+    return parser.parse_args()
 
 def T5Trainer(
     dataframe, args,
 ):
     torch.manual_seed(args.seed)  # pytorch random seed
     np.random.seed(args.seed)  # numpy random seed
     torch.backends.cudnn.deterministic = True
-    
+
     if args.evaluate_dir is not None:
         args.model = args.evaluate_dir
 
@@ -72,7 +73,7 @@ def T5Trainer(
     train_qids = qids['train']
     test_qids = qids['test']
     val_qids = qids['val']
-    
+
     if args.evaluate_dir is not None:
         save_dir = args.evaluate_dir
     else:
@@ -139,7 +140,7 @@ def T5Trainer(
             args,
             args.eval_le,
         )
-        
+
         test_set = ScienceQADatasetStd(
             problems,
             test_qids,
@@ -155,11 +156,8 @@ def T5Trainer(
     def extract_ans(ans):
         pattern = re.compile(r'The answer is \(([A-Z])\)')
         res = pattern.findall(ans)
-        
-        if len(res) == 1:
-            answer = res[0]  # 'A', 'B', ...
-        else:
-            answer = "FAILED" 
+
+        answer = res[0] if len(res) == 1 else "FAILED"
         return answer  
 
     # accuracy for answer inference
@@ -184,7 +182,7 @@ def compute_metrics_acc(eval_preds):
             if reference == best_option:
                 correct +=1 
         return {'accuracy': 1.0*correct/len(targets)}
-    
+
     # rougel for rationale generation
     metric = evaluate.load("rouge")
     def postprocess_text(preds, labels):
@@ -218,13 +216,13 @@ def compute_metrics_rougel(eval_preds):
     if args.final_eval:
         training_args = Seq2SeqTrainingArguments(
             save_dir,
-            do_train=True if args.evaluate_dir is None else False,
+            do_train=args.evaluate_dir is None,
             do_eval=False,
             evaluation_strategy="no",
             logging_strategy="steps",
             save_strategy="epoch",
-            save_total_limit = 2,
-            learning_rate= args.lr,
+            save_total_limit=2,
+            learning_rate=args.lr,
             eval_accumulation_steps=args.eval_acc,
             per_device_train_batch_size=args.bs,
             per_device_eval_batch_size=args.eval_bs,
@@ -233,23 +231,24 @@ def compute_metrics_rougel(eval_preds):
             predict_with_generate=args.use_generate,
             report_to="none",
         )
-    # evaluate at each epoch
     else:
         training_args = Seq2SeqTrainingArguments(
             save_dir,
-            do_train=True if args.evaluate_dir is None else False,
+            do_train=args.evaluate_dir is None,
             do_eval=True,
             evaluation_strategy="epoch",
             logging_strategy="steps",
             save_strategy="epoch",
-            save_total_limit = 2,
-            learning_rate= args.lr,
+            save_total_limit=2,
+            learning_rate=args.lr,
             eval_accumulation_steps=args.eval_acc,
             per_device_train_batch_size=args.bs,
             per_device_eval_batch_size=args.eval_bs,
             weight_decay=0.01,
             num_train_epochs=args.epoch,
-            metric_for_best_model="accuracy" if args.prompt_format != "QCM-LE" else "rougeL",
+            metric_for_best_model="accuracy"
+            if args.prompt_format != "QCM-LE"
+            else "rougeL",
             predict_with_generate=args.use_generate,
             load_best_model_at_end=True,
             report_to="none",
@@ -268,12 +267,12 @@ def compute_metrics_rougel(eval_preds):
     if args.evaluate_dir is None:
         trainer.train()
         trainer.save_model(save_dir)
-        
+
     metrics = trainer.evaluate(eval_dataset = test_set)
     trainer.log_metrics("test", metrics)
     trainer.save_metrics("test", metrics)
 
-    predict_results = trainer.predict(test_dataset=test_set, max_length=args.output_len) 
+    predict_results = trainer.predict(test_dataset=test_set, max_length=args.output_len)
     if trainer.is_world_process_zero():
         if args.use_generate:
             preds, targets = predict_results.predictions, predict_results.label_ids
@@ -292,7 +291,7 @@ def compute_metrics_rougel(eval_preds):
         results_ans = {}
         results_rationale = {}
         results_reference = {}
-        
+
         num_fail = 0
         for idx, qid in enumerate(test_qids):
             pred = preds[int(idx)]
@@ -302,7 +301,7 @@ def compute_metrics_rougel(eval_preds):
                 if extract_pred in args.options:
                     extract_pred = args.options.index(extract_pred)
                 else:
-                    extract_pred = random.choice(range(0,len(args.options)))
+                    extract_pred = random.choice(range(len(args.options)))
             else:
                 num_fail += 1
                 extract_pred = random.choice(range(len(args.options))) # random choose one option
@@ -320,7 +319,7 @@ def compute_metrics_rougel(eval_preds):
         output_prediction_file = os.path.join(save_dir,"predictions_ans_test.json")
         with open(output_prediction_file, "w") as writer:
             writer.write(json.dumps(output_data, indent=4))
-    
+
     # generate the rationale for the eval set
     if args.prompt_format == "QCM-LE":
         torch.cuda.empty_cache()

model.py

@@ -2,20 +2,19 @@
 Adapted from https://github.com/huggingface/transformers
 '''
 
-from transformers import T5Config, T5ForConditionalGeneration
-from transformers.models.t5.modeling_t5 import T5Stack, __HEAD_MASK_WARNING_MSG, T5EncoderModel
 import copy
 import math
 import os
 import warnings
 from typing import Optional, Tuple, Union
+
 import torch
 from torch import nn
 from torch.nn import CrossEntropyLoss
-from transformers.modeling_outputs import (
-    BaseModelOutput,
-    Seq2SeqLMOutput,
-)
+from transformers import T5Config, T5ForConditionalGeneration
+from transformers.modeling_outputs import BaseModelOutput, Seq2SeqLMOutput
+from transformers.models.t5.modeling_t5 import (__HEAD_MASK_WARNING_MSG, T5EncoderModel, T5Stack)
+
 
 class T5ForMultimodalGeneration(T5ForConditionalGeneration):
     _keys_to_ignore_on_load_missing = [
@@ -87,10 +86,13 @@ def forward(
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
         # FutureWarning: head_mask was separated into two input args - head_mask, decoder_head_mask
-        if head_mask is not None and decoder_head_mask is None:
-            if self.config.num_layers == self.config.num_decoder_layers:
-                warnings.warn(__HEAD_MASK_WARNING_MSG, FutureWarning)
-                decoder_head_mask = head_mask
+        if (
+            head_mask is not None
+            and decoder_head_mask is None
+            and self.config.num_layers == self.config.num_decoder_layers
+        ):
+            warnings.warn(__HEAD_MASK_WARNING_MSG, FutureWarning)
+            decoder_head_mask = head_mask
 
         # Encode if needed (training, first prediction pass)
         if encoder_outputs is None:
@@ -114,7 +116,7 @@ def forward(
 
 
         hidden_states = encoder_outputs[0]
-        
+
         image_embedding = self.image_dense(image_ids)
         image_att, _ = self.mha_layer(hidden_states, image_embedding, image_embedding)
 

utils_data.py

@@ -1,9 +1,10 @@
-import os
-from torch.utils.data import Dataset
-import os
 import json
+import os
+
 import numpy as np
 import torch
+from torch.utils.data import Dataset
+
 from utils_prompt import *
 
 img_shape = {
@@ -20,9 +21,9 @@ def load_data_std(args):
     for qid in problems:
         problems[qid]['caption'] = captions[qid] if qid in captions else ""
 
-    train_qids = pid_splits['%s' % (args.train_split)]
-    val_qids = pid_splits['%s' % (args.val_split)]
-    test_qids = pid_splits['%s' % (args.test_split)]
+    train_qids = pid_splits[f'{args.train_split}']
+    val_qids = pid_splits[f'{args.val_split}']
+    test_qids = pid_splits[f'{args.test_split}']
     print(f"number of train problems: {len(train_qids)}\n")
     print(f"number of val problems: {len(val_qids)}\n")
     print(f"number of test problems: {len(test_qids)}\n")
@@ -43,18 +44,16 @@ def load_data_img(args):
         image_features = image_features.repeat(512, axis=1)
     elif args.img_type == "clip":
         image_features = np.load('vision_features/clip.npy')
-    elif args.img_type == "detr":
-        image_features = np.load('vision_features/detr.npy')
     else:
         image_features = np.load('vision_features/detr.npy')
     print("img_features size: ", image_features.shape)
 
     for qid in problems:
         problems[qid]['caption'] = captions[qid] if qid in captions else ""
 
-    train_qids = pid_splits['%s' % (args.train_split)]
-    val_qids = pid_splits['%s' % (args.val_split)]
-    test_qids = pid_splits['%s' % (args.test_split)]
+    train_qids = pid_splits[f'{args.train_split}']
+    val_qids = pid_splits[f'{args.val_split}']
+    test_qids = pid_splits[f'{args.test_split}']
     print(f"number of train problems: {len(train_qids)}\n")
     print(f"number of val problems: {len(val_qids)}\n")
     print(f"number of test problems: {len(test_qids)}\n")
@@ -79,10 +78,7 @@ def __init__(
         self.summ_len = target_len
         self.target_text = []
         self.source_text = []
-        if test_le is not None:
-            test_le_data =json.load(open(test_le))["preds"]
-        else:
-            test_le_data = None
+        test_le_data = None if test_le is None else json.load(open(test_le))["preds"]
         idx = 0
         for qid in self.data:
             if test_le_data is not None:
@@ -161,10 +157,7 @@ def __init__(
         self.target_text = []
         self.source_text = []
         self.image_ids = []
-        if test_le is not None:
-            test_le_data =json.load(open(test_le))["preds"]
-        else:
-            test_le_data = None
+        test_le_data = None if test_le is None else json.load(open(test_le))["preds"]
         idx = 0
         for qid in self.data:
             if test_le_data is not None: