import pickle import random from collections import defaultdict from typing import List import torch import numpy as np from os import listdir from os.path import isfile, join import pandas as pd def set_seed(seed): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed(seed) torch.cuda.manual_seed_all(seed) def get_ckpt(args): certain_epoch, given_ckpt, ckpt_saved_folder = args.certain_epoch, args.given_ckpt, args.logs_dir ckpt_files = [f for f in listdir(ckpt_saved_folder) if isfile(join(ckpt_saved_folder, f))] assert len(ckpt_files) >= 1 desired_ckpt_name = '' desired_ckpt_epoch = 0 if given_ckpt is not None: desired_ckpt_epoch = int(given_ckpt.split('=')[1].split('-')[0]) desired_ckpt_name = ckpt_saved_folder + given_ckpt else: for ckpt_name in ckpt_files: if ckpt_name[-4:] != 'ckpt': continue if ckpt_name.split('_epoch')[0] != args.tag_info: continue try: ckpt_epoch = int(ckpt_name.split('epoch=')[1].split('-')[0]) except: continue if certain_epoch is not None: if certain_epoch == ckpt_epoch: desired_ckpt_epoch, desired_ckpt_name = ckpt_epoch, ckpt_name else: if ckpt_epoch > desired_ckpt_epoch: desired_ckpt_epoch, desired_ckpt_name = ckpt_epoch, ckpt_name print('=' * 20) print('Loading: ' + desired_ckpt_name) print('=' * 20) assert desired_ckpt_name in ckpt_files return ckpt_saved_folder + desired_ckpt_name, desired_ckpt_epoch def grad_status(model): return (par.requires_grad for par in model.parameters()) def lmap(f, x): """list(map(f, x))""" return list(map(f, x)) def assert_all_frozen(model): model_grads: List[bool] = list(grad_status(model)) n_require_grad = sum(lmap(int, model_grads)) npars = len(model_grads) assert not any(model_grads), f"{n_require_grad / npars:.1%} of {npars} weights require grad" def dec_2d(dec, size): res = [] i = 0 while i < len(dec): res.append(dec[i: i + size]) i = i + size return res ###### decoder helper def numerical_decoder(args, cuda_ids, output): np_ids = cuda_ids.cpu().numpy() begin_and_end_token = np.where(np_ids == 1) if output: if len(begin_and_end_token) != 1 or begin_and_end_token[0].size < 1: print("Invalid Case") return "0" if args.hierarchic_decode: np_ids = np_ids[1:begin_and_end_token[0][0]] - 2 else: np_ids = (np_ids[1:begin_and_end_token[0][0]] - 2) % args.output_vocab_size else: if args.hierarchic_decode: np_ids = np_ids[:begin_and_end_token[0][0]] - 2 else: np_ids = (np_ids[:begin_and_end_token[0][0]] - 2) % args.output_vocab_size bits = int(np.log10(args.output_vocab_size)) num_list = list(map(str, np_ids)) str_ids = ''.join([c.zfill(bits) if (i != len(num_list) - 1) else c for i, c in enumerate(num_list)]) return str_ids def random_shuffle(doc_id): new_doc_id = "" for index in range(0, len(doc_id)): while True: rand_digit = np.random.randint(0, 9) if not rand_digit == int(doc_id[index]): new_doc_id += str(rand_digit) break return new_doc_id def augment(query): if len(query) < 20*10: start_pos = np.random.randint(0, int(len(query)+1/2)) end_pos = np.random.randint(start_pos, len(query)) span_length = max(start_pos-end_pos, 10*10) new_query = str(query[start_pos:start_pos+span_length]) else: start_pos = np.random.randint(0, len(query)-10*10) end_pos = np.random.randint(start_pos+5*10, len(query)) span_length = min(start_pos-end_pos, 20*10) new_query = str(query[start_pos:start_pos+span_length]) #print(new_query) return new_query def load_data(args): def process_func(doc_to_query_list, index, query, newid, rank=1): if args.kary: docid = '-'.join([c for c in str(newid).split('-')]) softmax_index_list = str(newid).split('-')[:3] softmax_index = 0 for id, num in enumerate(softmax_index_list[::-1]): softmax_index += int(num) * (args.kary ** id) else: docid = ''.join([c for c in str(newid)]) softmax_index = int(str(newid)[:3]) ##import hard negative here neg_docid_list = [] if args.hard_negative: for i in range(1): neg_docid_list.append(random_shuffle(docid)) aug_query_list = [] if args.aug_query: if args.aug_query_type == 'aug_query': if newid in doc_to_query_list: aug_query_list = doc_to_query_list[newid] else: for i in range(10): aug_query_list.append(augment(query)) return query, docid, rank, softmax_index, neg_docid_list, aug_query_list result = None q_emb, query_id_dict_train = None, None prefix_embedding = None prefix_mask = None prefix2idx_dict = None doc_to_query_list = None assert args.contrastive_variant == '' ## load pre-defined id_class in train.sh if 'gtq' in args.query_type: if args.trivia: # train_file = '../Data_process/Trivia_dataset/train.tsv' train_file = args.trivia_train_path df = pd.read_csv( train_file, encoding='utf-8', names=["query", "queryid", "oldid", "bert_k30_c30_1", "bert_k30_c30_2", "bert_k30_c30_3", "bert_k30_c30_4", "bert_k30_c30_5"], header=None, sep='\t', dtype={'query': str, 'queryid': str, 'oldid': str, args.id_class: str}).loc[:, ["query", args.id_class]] assert not df.isnull().values.any() doc_to_query_list = defaultdict(set) for [query, docid] in df.values.tolist(): doc_to_query_list[docid].add(query) elif args.nq: # train_filename = '../Data_process/NQ_dataset/nq_train_doc_newid.tsv' train_filename = args.nq_train_doc_newid_path print(train_filename) df = pd.read_csv( train_filename, names=["query", "queryid", "oldid", "bert_k30_c30_1", "bert_k30_c30_2", "bert_k30_c30_3", "bert_k30_c30_4", "bert_k30_c30_5"], encoding='utf-8', header=None, sep='\t', dtype={'query': str, 'queryid': str, 'oldid': str, args.id_class: str}).loc[:, ["query", args.id_class]] assert not df.isnull().values.any() doc_to_query_list = defaultdict(set) for [query, docid] in df.values.tolist(): doc_to_query_list[docid].add(query) ## Query Generation Data if 'qg' in args.query_type: if args.trivia: # qg_file = '../Data_process/Trivia_dataset/trivia_512_qg.tsv' qg_file = args.trivia_qg_path gq_df1 = pd.read_csv( qg_file, names=["query", "oldid", "bert_k30_c30_1", "bert_k30_c30_2", "bert_k30_c30_3", "bert_k30_c30_4", "bert_k30_c30_5"], encoding='utf-8', quoting=3, header=None, sep='\t', dtype={'query': str, 'queryid': str, "oldid": str, args.id_class: str}).loc[:, ["query", args.id_class]] gq_df1 = gq_df1.dropna(axis=0) elif args.nq: # qg_file = '../Data_process/NQ_dataset/NQ_512_qg.tsv' qg_file = args.nq_qg_path gq_df1 = pd.read_csv( qg_file, names=["query", "queryid", "oldid", "bert_k30_c30_1", "bert_k30_c30_2", "bert_k30_c30_3", "bert_k30_c30_4", "bert_k30_c30_5"], encoding='utf-8', header=None, sep='\t', dtype={'query': str, 'queryid': str, 'oldid': str, args.id_class: str}).loc[:, ["query", args.id_class]] gq_df1 = gq_df1.dropna(axis=0) print(len(gq_df1)) for [query, docid] in gq_df1.values.tolist(): doc_to_query_list[docid].add(query) temp = defaultdict(list) for k,v in doc_to_query_list.items(): temp[k] = list(v) doc_to_query_list = temp result = tuple( process_func(doc_to_query_list, index, *row) for index, row in enumerate(zip(df["query"], df[args.id_class])) ) result_add1 = tuple( process_func(doc_to_query_list, index, *row) for index, row in enumerate(zip(gq_df1["query"], gq_df1[args.id_class])) ) result = result + result_add1 else: result = tuple( process_func(doc_to_query_list, index, *row) for index, row in enumerate(zip(df["query"], df[args.id_class])) ) path_list = [] if 'doc' in args.query_type: if args.trivia: # filename = '../Data_process/Trivia_dataset/trivia_title_cont.tsv' filename = args.trivia_title_cont_path path_list.append(filename) elif args.nq: # filename = '../Data_process/NQ_dataset/nq_title_abs.tsv' filename = args.nq_title_abs_path path_list.append(filename) if 'doc_aug' in args.query_type: if args.trivia: # filename = '../Data_process/Trivia_dataset/trivia_doc_aug.tsv' filename = args.trivia_doc_aug_path path_list.append(filename) elif args.nq: # filename = '../Data_process/NQ_dataset/NQ_doc_aug.tsv' filename = args.nq_doc_aug_path path_list.append(filename) for file_path in path_list: print(file_path) if args.trivia: df1 = pd.read_csv( file_path, names=["query", "queryid", "oldid", "bert_k30_c30_1", "bert_k30_c30_2", "bert_k30_c30_3", "bert_k30_c30_4", "bert_k30_c30_5"], header=None, sep='\t', dtype={'query': str, 'queryid': str, 'oldid': str, args.id_class: str}).loc[:, ["query", "queryid", args.id_class]] elif args.nq: df1 = pd.read_csv( file_path, names=["query", "queryid", "oldid", "bert_k30_c30_1", "bert_k30_c30_2", "bert_k30_c30_3", "bert_k30_c30_4", "bert_k30_c30_5"], encoding='utf-8', header=None, sep='\t', dtype={'query': str, 'queryid': str, 'oldid': str, args.id_class: str}).loc[:, ["query", "queryid", args.id_class]] df1.dropna(axis=0, inplace=True) assert not df1.isnull().values.any() result_add1 = tuple( filter(None, (process_func(doc_to_query_list, index, *row) for index, row in enumerate(zip(df1["query"], df1[args.id_class]))) ) ) result = result_add1 if result is None else result + result_add1 print('&' * 20) print(result[0]) print('&' * 20) return result, q_emb, query_id_dict_train, prefix_embedding, prefix_mask, prefix2idx_dict def load_data_infer(args): df = None if args.test_set == 'dev': if args.trivia: # dev_file = '../Data_process/Trivia_dataset/dev.tsv' # test_file = '../Data_process/Trivia_dataset/test.tsv' dev_file = args.trivia_dev_path test_file = args.trivia_test_path df = pd.read_csv( dev_file, encoding='utf-8', names=["query", "queryid", "oldid", "bert_k30_c30_1", "bert_k30_c30_2", "bert_k30_c30_3", "bert_k30_c30_4", "bert_k30_c30_5"], header=None, sep='\t', dtype={'query': str, 'queryid': str, 'oldid': str, args.id_class: str}).loc[:, ["query", 'queryid', args.id_class]] elif args.nq: # dev_filename = '../Data_process/NQ_dataset/nq_dev_doc_newid.tsv' dev_filename = args.nq_dev_doc_newid_path df = pd.read_csv( dev_filename, names=["query", "queryid", "oldid", "bert_k30_c30_1", "bert_k30_c30_2", "bert_k30_c30_3", "bert_k30_c30_4", "bert_k30_c30_5"], encoding='utf-8', header=None, sep='\t', dtype={'query': str, 'queryid': str, 'oldid': str, args.id_class: str}).loc[:, ["query", "queryid", args.id_class]] assert not df.isnull().values.any() result = [] softmax_index = -1 for index, row in df.iterrows(): query = row["query"] if args.trivia: if args.kary: rank1 = row[args.id_class] rank1 = rank1.split(",") rank1 = rank1[0] rank1 = '-'.join([c for c in str(rank1).split('-')]) else: rank1 = row[args.id_class] rank1 = rank1.split(",") rank1 = rank1[0] elif args.nq: rank1 = ''.join([c for c in str(row[args.id_class])]) list_sum = [] if args.kary: docid = '-'.join([c for c in str(row[args.id_class]).split('-')]) softmax_index_list = str(row[args.id_class]).split('-')[:3] else: docid = ''.join([c for c in str(row[args.id_class])]) softmax_index = int(str(row[args.id_class])[:3]) rank = 1 list_sum.append((docid, rank)) neg_docid_list = [] if args.hard_negative: for i in range(1): neg_docid_list.append(random_shuffle(rank1)) aug_query_list = [] #do not need for inference if args.aug_query: for i in range(20): aug_query_list.append(augment(query)) result.append((query, rank1, list_sum, softmax_index, neg_docid_list, aug_query_list)) return tuple(result)