add tag process in dataloader

AlchemistZoro · AlchemistZoro · commit 3692444b24b8 · 2023-02-01T00:58:24.000+08:00
diff --git a/data_utils.py b/data_utils.py
diff --git a/greaselm.py b/greaselm.py
@@ -61,7 +61,11 @@ def load_data(args, devices, kg):
         model_name=args.encoder,
         max_node_num=args.max_node_num, max_seq_length=args.max_seq_len,
         is_inhouse=args.inhouse, inhouse_train_qids_path=args.inhouse_train_qids,
-        subsample=args.subsample, n_train=args.n_train, debug=args.debug, cxt_node_connects_all=args.cxt_node_connects_all, kg=kg)
+        subsample=args.subsample, n_train=args.n_train, debug=args.debug, cxt_node_connects_all=args.cxt_node_connects_all, kg=kg,emp=args.emp,
+        train_tagged_path = args.train_tagged,
+        dev_tagged_path=args.dev_tagged,
+        test_tagged_path=args.test_tagged,
+        )
 
     return dataset
 
@@ -560,7 +564,8 @@ def main(args):
 
     parser = parser_utils.get_parser()
     args, _ = parser.parse_known_args()
-
+    # print(args.train_statements,args.train_tagged)
+    # input()
     # General
     parser.add_argument('--mode', default='train', choices=['train', 'eval'], help='run training or evaluation')
     parser.add_argument('--save_dir', default=f'./saved_models/greaselm/', help='model output directory')
@@ -609,5 +614,10 @@ def main(args):
     parser.add_argument('--refreeze_epoch', default=10000, type=int)
     parser.add_argument('--init_range', default=0.02, type=float, help='stddev when initializing with normal distribution')
 
+    # MyGLM
+    parser.add_argument('--emp',default=True,type=bool)
+    
     args = parser.parse_args()
+    # print(args.train_statements,args.train_tagged)
+ 
     main(args)
diff --git a/preprocess_utils/tagging.py b/preprocess_utils/tagging.py
@@ -173,8 +173,10 @@ def tag(statement_path, cpnet_vocab_path, pattern_path, output_path, num_process
 
     # check_path(output_path)
     with open(output_path, 'w') as fout:
-        for dic in res:
-            fout.write(json.dumps(dic) + '\n')
+        # change write file to json format
+        fout.write(json.dumps(res))
+        # for dic in res:
+        #     fout.write(json.dumps(dic) + '\n')
 
     print(f'grounded concepts saved to {output_path}')
     print()
@@ -185,7 +187,7 @@ def tag(statement_path, cpnet_vocab_path, pattern_path, output_path, num_process
     statement_path = '../data/csqa/statement/test.statement.jsonl'
     cpnet_vocab_path = '../data/cpnet/concept.txt'
     pattern_path = '../data/cpnet/matcher_patterns.json'
-    output_path = '../data/obqa/tagged/test.tagged.jsonl'
+    output_path = '../data/obqa/tagged/test.jsonl'
     num_processes = 1
     debug=True
     tag(statement_path, cpnet_vocab_path, pattern_path, output_path, num_processes, debug)
diff --git a/run_greaselm.sh b/run_greaselm.sh
@@ -2,7 +2,6 @@
 export TOKENIZERS_PARALLELISM=true
 dt=`date '+%Y%m%d_%H%M%S'`
 
-
 dataset=$1
 shift
 encoder='roberta-large'
diff --git a/utils/data_utils.py b/utils/data_utils.py
@@ -3,12 +3,13 @@
 import pickle
 import os
 
+
 import numpy as np
 import torch
 from tqdm import tqdm
 from transformers import (OPENAI_GPT_PRETRAINED_CONFIG_ARCHIVE_MAP, BERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
                           XLNET_PRETRAINED_CONFIG_ARCHIVE_MAP, ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP)
-from transformers import (OpenAIGPTTokenizer, BertTokenizer, XLNetTokenizer, RobertaTokenizer)
+from transformers import (OpenAIGPTTokenizer, BertTokenizer, XLNetTokenizer, RobertaTokenizer,AutoTokenizer)
 try:
     from transformers import ALBERT_PRETRAINED_CONFIG_ARCHIVE_MAP
     from transformers import AlbertTokenizer
@@ -18,6 +19,8 @@
 from preprocess_utils import conceptnet
 from utils import utils
 
+
+
 MODEL_CLASS_TO_NAME = {
     'gpt': list(OPENAI_GPT_PRETRAINED_CONFIG_ARCHIVE_MAP.keys()),
     'bert': list(BERT_PRETRAINED_CONFIG_ARCHIVE_MAP.keys()),
@@ -94,7 +97,7 @@ def __init__(self, train_statement_path, train_adj_path,
                  test_statement_path, test_adj_path,
                  batch_size, eval_batch_size, device, model_name, max_node_num=200, max_seq_length=128,
                  is_inhouse=False, inhouse_train_qids_path=None,
-                 subsample=1.0, n_train=-1, debug=False, cxt_node_connects_all=False, kg="cpnet"):
+                 subsample=1.0, n_train=-1, debug=False, cxt_node_connects_all=False, kg="cpnet",emp=False,train_tagged_path='',dev_tagged_path = '',test_tagged_path = ''):
         super().__init__()
         self.batch_size = batch_size
         self.eval_batch_size = eval_batch_size
@@ -105,13 +108,16 @@ def __init__(self, train_statement_path, train_adj_path,
         self.max_node_num = max_node_num
         self.debug_sample_size = 32
         self.cxt_node_connects_all = cxt_node_connects_all
+        
+        # emp control the embedding pooling process
+        self.emp = emp
 
         self.model_type = MODEL_NAME_TO_CLASS[model_name]
         self.load_resources(kg)
 
         # Load training data
         print ('train_statement_path', train_statement_path)
-        self.train_qids, self.train_labels, self.train_encoder_data, train_concepts_by_sents_list = self.load_input_tensors(train_statement_path, max_seq_length)
+        self.train_qids, self.train_labels, self.train_encoder_data, train_concepts_by_sents_list = self.load_input_tensors(train_statement_path, max_seq_length,train_tagged_path,emp)
 
         num_choice = self.train_encoder_data[0].size(1)
         self.num_choice = num_choice
@@ -123,7 +129,7 @@ def __init__(self, train_statement_path, train_adj_path,
         print("Finish loading training data.")
         
         # Load dev data
-        self.dev_qids, self.dev_labels, self.dev_encoder_data, dev_concepts_by_sents_list = self.load_input_tensors(dev_statement_path, max_seq_length)
+        self.dev_qids, self.dev_labels, self.dev_encoder_data, dev_concepts_by_sents_list = self.load_input_tensors(dev_statement_path, max_seq_length,dev_tagged_path,emp)
         *self.dev_decoder_data, self.dev_adj_data = self.load_sparse_adj_data_with_contextnode(dev_adj_path, max_node_num, dev_concepts_by_sents_list)
         if not debug:
             assert all(len(self.dev_qids) == len(self.dev_adj_data[0]) == x.size(0) for x in [self.dev_labels] + self.dev_encoder_data + self.dev_decoder_data)
@@ -132,7 +138,7 @@ def __init__(self, train_statement_path, train_adj_path,
 
         # Load test data
         if test_statement_path is not None:
-            self.test_qids, self.test_labels, self.test_encoder_data, test_concepts_by_sents_list = self.load_input_tensors(test_statement_path, max_seq_length)
+            self.test_qids, self.test_labels,self.test_encoder_data, test_concepts_by_sents_list = self.load_input_tensors(test_statement_path, max_seq_length,test_tagged_path,emp)
             *self.test_decoder_data, self.test_adj_data = self.load_sparse_adj_data_with_contextnode(test_adj_path, max_node_num, test_concepts_by_sents_list)
             if not debug:
                 assert all(len(self.test_qids) == len(self.test_adj_data[0]) == x.size(0) for x in [self.test_labels] + self.test_encoder_data + self.test_decoder_data)
@@ -210,11 +216,13 @@ def test(self):
 
     def load_resources(self, kg):
         # Load the tokenizer
-        try:
-            tokenizer_class = {'bert': BertTokenizer, 'xlnet': XLNetTokenizer, 'roberta': RobertaTokenizer, 'albert': AlbertTokenizer}.get(self.model_type)
-        except:
-            tokenizer_class = {'bert': BertTokenizer, 'xlnet': XLNetTokenizer, 'roberta': RobertaTokenizer}.get(self.model_type)
-        tokenizer = tokenizer_class.from_pretrained(self.model_name)
+        # try:
+        #     tokenizer_class = {'bert': BertTokenizer, 'xlnet': XLNetTokenizer, 'roberta': RobertaTokenizer, 'albert': AlbertTokenizer}.get(self.model_type)
+        # except:
+        #     tokenizer_class = {'bert': BertTokenizer, 'xlnet': XLNetTokenizer, 'roberta': RobertaTokenizer}.get(self.model_type)
+        # use autotokenizer to use tagging
+        tokenizer = AutoTokenizer.from_pretrained(self.model_name,use_fast =True)
+        #tokenizer = tokenizer_class.from_pretrained(self.model_name)
         self.tokenizer = tokenizer
 
         if kg == "cpnet":
@@ -249,14 +257,21 @@ def load_resources(self, kg):
         else:
             raise ValueError("Invalid value for kg.")
 
-    def load_input_tensors(self, input_jsonl_path, max_seq_length):
+    def load_input_tensors(self, input_jsonl_path, max_seq_length,tagged_jsonl_path,emp):
         """Construct input tensors for the LM component of the model."""
-        cache_path = input_jsonl_path + "-sl{}".format(max_seq_length) + (("-" + self.model_type) if self.model_type != "roberta" else "") + '.loaded_cache'
+        if emp:
+            cache_path = input_jsonl_path + "-sl{}".format(max_seq_length) + (("-" + self.model_type) if self.model_type != "roberta" else "") + '-tag' + '.loaded_cache'
+        else:
+            cache_path = input_jsonl_path + "-sl{}".format(max_seq_length) + (("-" + self.model_type) if self.model_type != "roberta" else "") + '.loaded_cache'
+            
         use_cache = True
         
         if use_cache and not os.path.exists(cache_path):
             use_cache = False
 
+        #debug
+        # use_cache = False
+        
         if use_cache:
             with open(cache_path, 'rb') as f:
                 input_tensors = utils.CPU_Unpickler(f).load()
@@ -266,7 +281,7 @@ def load_input_tensors(self, input_jsonl_path, max_seq_length):
             elif self.model_type in ('gpt',):
                 input_tensors = load_gpt_input_tensors(input_jsonl_path, max_seq_length)
             elif self.model_type in ('bert', 'xlnet', 'roberta', 'albert'):
-                input_tensors = load_bert_xlnet_roberta_input_tensors(input_jsonl_path, max_seq_length, self.debug, self.tokenizer, self.debug_sample_size)
+                input_tensors = load_bert_xlnet_roberta_input_tensors(input_jsonl_path, max_seq_length, self.debug, self.tokenizer, self.debug_sample_size,tagged_jsonl_path,emp)
                 
             if not self.debug:
                 utils.save_pickle(input_tensors, cache_path)
@@ -508,7 +523,7 @@ def tokenize_and_encode(tokenizer, obj):
     return examples_ids, mc_labels, input_ids, mc_token_ids, lm_labels
 
 
-def load_bert_xlnet_roberta_input_tensors(statement_jsonl_path, max_seq_length, debug, tokenizer, debug_sample_size):
+def load_bert_xlnet_roberta_input_tensors(statement_jsonl_path, max_seq_length, debug, tokenizer, debug_sample_size,tagged_jsonl_path,emp):
     class InputExample(object):
 
         def __init__(self, example_id, question, contexts, endings, label=None):
@@ -528,8 +543,9 @@ def __init__(self, example_id, choices_features, label):
                     'input_mask': input_mask,
                     'segment_ids': segment_ids,
                     'output_mask': output_mask,
+                    'pool_mask':pool_mask
                 }
-                for input_ids, input_mask, segment_ids, output_mask in choices_features
+                for input_ids, input_mask, segment_ids, output_mask,pool_mask in choices_features
             ]
             self.label = label
 
@@ -554,8 +570,52 @@ def read_examples(input_file):
                         label=label
                     ))
         return examples
+    
+    def read_tagged_file(input_file):
+        # with open(input_file,"r",encoding="utf-8") as f:
+        #     examples = []
+        #     for line in f.readlines():
+        #         json_dic = json.loads(line)
+        #         statement = json_dic['statements']
+        #         answers = json_dic['answers']
+        #         stem = json_dic['stem']
+        #         examples.append((statement,answers,stem))
+        f = open(input_file)
+        examples = json.load(f)
+        return examples
+                
+                
+    def get_pool_mask(encoded_input_words,context,ending,tagged_ending,tagged_context,pool_mask):
+        # print(encoded_input_words,'\n\n',context,'\n\n',ending,'\n\n',tagged_ending,'\n\n',tagged_context)
+        def get_pool_set(items):
+            pool_set = set()
+            # item [11, 12, 'afford']
+            for item in items:
+                start = item[0]
+                end = item[1]
+                if end-start>1:
+                    for i in range(start,end):
+                        pool_set.add(i)
+            return pool_set
+        
+        end_pool_set = get_pool_set(tagged_ending)
+        context_pool_set = get_pool_set(tagged_context)
+        flag = True
+        for i in range(len(pool_mask)):
+            if flag:
+                if encoded_input_words[i] in context_pool_set:
+                    pool_mask[i] = 1
+                if encoded_input_words[i] == None and encoded_input_words[i+1]==None:
+                    flag = False
+            else:
+                if encoded_input_words[i] in end_pool_set:
+                    pool_mask[i] = 1
+                if encoded_input_words[i] == None and encoded_input_words[i+1]==None:
+                    break
+        # print(pool_mask)
+        return pool_mask             
 
-    def simple_convert_examples_to_features(examples, label_list, max_seq_length, tokenizer):
+    def simple_convert_examples_to_features(examples, label_list, max_seq_length, tokenizer,emp,tagged):
         """ Loads a data file into a list of `InputBatch`s
             `cls_token_at_end` define the location of the CLS token:
                 - False (Default, BERT/XLM pattern): [CLS] + A + [SEP] + B + [SEP]
@@ -572,20 +632,35 @@ def simple_convert_examples_to_features(examples, label_list, max_seq_length, to
             choices_features = []
             for ending_idx, (context, ending) in enumerate(zip(example.contexts, example.endings)):
                 ans = example.question + " " + ending
-
+                if emp: 
+                    tagged_ending = tagged[ex_index]['answers'][ending_idx]
+                    tagged_context = tagged[ex_index]['stem']
+                
                 encoded_input = tokenizer(context, ans, padding="max_length", truncation=True, max_length=max_seq_length, return_token_type_ids=True, return_special_tokens_mask=True)
+                # print(encoded_input.words())
                 input_ids = encoded_input["input_ids"]
                 output_mask = encoded_input["special_tokens_mask"]
                 input_mask = encoded_input["attention_mask"]
                 segment_ids = encoded_input["token_type_ids"]
-                # print(context,'\n',ans,'\n',encoded_input["input_ids"])
-                # input()
+                pool_mask = [0]*max_seq_length
+                if emp:
+                    pool_mask = get_pool_mask(encoded_input.words(),context,ending,tagged_ending,tagged_context,pool_mask)
+
+
                 assert len(input_ids) == max_seq_length
                 assert len(output_mask) == max_seq_length
                 assert len(input_mask) == max_seq_length
                 assert len(segment_ids) == max_seq_length
 
-                choices_features.append((input_ids, input_mask, segment_ids, output_mask))
+                choices_features.append((input_ids, input_mask, segment_ids, output_mask,pool_mask))
+                
+                # for i in range(max_seq_length):
+                #     if input_ids[i] == 1:     
+                #         print(len(context.split())+len(ans.split()),i)
+                #         break
+                
+                # input()
+                
             label = label_map[example.label]
             features.append(InputFeatures(example_id=example.example_id, choices_features=choices_features, label=label))
 
@@ -604,8 +679,34 @@ def convert_features_to_tensors(features):
         return all_input_ids, all_input_mask, all_segment_ids, all_output_mask, all_label
 
     examples = read_examples(statement_jsonl_path)
-    features, concepts_by_sents_list = simple_convert_examples_to_features(examples, list(range(len(examples[0].endings))), max_seq_length, tokenizer)
+    tagged = []
+    if emp:
+        tagged = read_tagged_file(tagged_jsonl_path)
+
+    
+    features, concepts_by_sents_list = simple_convert_examples_to_features(examples, list(range(len(examples[0].endings))), max_seq_length, tokenizer,emp,tagged)
     
     example_ids = [f.example_id for f in features]
     *data_tensors, all_label = convert_features_to_tensors(features)
     return example_ids, all_label, data_tensors, concepts_by_sents_list
+
+if __name__ == "__main__":
+    
+
+ 
+    model_name = 'roberta-large'
+    tokenizer = AutoTokenizer.from_pretrained(model_name,use_fast=True)   
+    # model_name = 'bert-base-uncased'
+    # tokenizer = AutoTokenizer.from_pretrained(model_name,use_fast=True)
+    
+    
+    statement_jsonl_path = './data/obqa/statement/test.statement.jsonl'
+    max_seq_length = 128
+    debug = False
+    emp = False
+    
+    debug_sample_size = 32
+    tagged_jsonl_path = './data/obqa/tagged/test.tagged.jsonl'
+    
+    example_ids, all_label, data_tensors, concepts_by_sents_list=load_bert_xlnet_roberta_input_tensors(statement_jsonl_path, max_seq_length, debug, tokenizer, debug_sample_size,tagged_jsonl_path,emp)
+    
diff --git a/utils/parser_utils.py b/utils/parser_utils.py
@@ -53,6 +53,13 @@ def add_data_arguments(parser):
     parser.add_argument('--train_statements', default='{data_dir}/{dataset}/statement/train.statement.jsonl')
     parser.add_argument('--dev_statements', default='{data_dir}/{dataset}/statement/dev.statement.jsonl')
     parser.add_argument('--test_statements', default='{data_dir}/{dataset}/statement/test.statement.jsonl')
+    
+    # tagged 
+    parser.add_argument('--train_tagged', default='{data_dir}/{dataset}/tagged/train.tagged.jsonl')
+    parser.add_argument('--dev_tagged', default='{data_dir}/{dataset}/tagged/dev.tagged.jsonl')
+    parser.add_argument('--test_tagged', default='{data_dir}/{dataset}/tagged/test.tagged.jsonl')
+    
+
     # preprocessing options
     parser.add_argument('-sl', '--max_seq_len', default=100, type=int)
     # set dataset defaults
@@ -62,11 +69,16 @@ def add_data_arguments(parser):
                         inhouse_train_qids=args.inhouse_train_qids.format(dataset=args.dataset))
     data_splits = ('train', 'dev') if args.dataset in DATASET_NO_TEST else ('train', 'dev', 'test')
     for split in data_splits:
-        for attribute in ('statements',):
+        # &#36825;&#37324;&#19981;&#21152;&#19978;tagged&#65292;tagged&#30340;&#25991;&#20214;&#36335;&#24452;&#30340;{data_dir}&#19981;&#20250;&#34987;&#26367;&#20195;&#20250;&#22914;&#19979;&#25152;&#31034;
+        #{data_dir}/{dataset}/tagged/train.tagged.jsonl
+        for attribute in ('statements','tagged'):
             attr_name = f'{split}_{attribute}'
             parser.set_defaults(**{attr_name: getattr(args, attr_name).format(dataset=args.dataset, data_dir=args.data_dir)})
     if 'test' not in data_splits:
         parser.set_defaults(test_statements=None)
+    # args, _ = parser.parse_known_args()
+    # print(args.train_statements,args.train_tagged)
+    # input()
 
 
 def add_encoder_arguments(parser):
@@ -106,4 +118,5 @@ def get_parser():
     add_encoder_arguments(parser)
     add_optimization_arguments(parser)
     add_additional_arguments(parser)
+    
     return parser
