import os
import sys
import json
import torch
import random
import pickle
import itertools
import numpy as np

from tqdm import tqdm
from typing import Any, Callable, Dict, List, NewType, Optional, Tuple, Union
from transformers import (OPENAI_GPT_PRETRAINED_CONFIG_ARCHIVE_MAP, BERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
                          XLNET_PRETRAINED_CONFIG_ARCHIVE_MAP, ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP, T5_PRETRAINED_CONFIG_ARCHIVE_MAP)
from transformers import (OpenAIGPTTokenizer, BertTokenizer, BertTokenizerFast, XLNetTokenizer, RobertaTokenizer, RobertaTokenizerFast, AutoTokenizer, PreTrainedTokenizerFast)
# T5Tokenizer

try:
    from transformers import ALBERT_PRETRAINED_CONFIG_ARCHIVE_MAP
    from transformers import AlbertTokenizer
except:
    pass

from preprocess_utils import conceptnet
from utils import utils

# added
from torch import Tensor
from torch_geometric.utils.num_nodes import maybe_num_nodes

from math import ceil
from sentence_transformers.util import dot_score
from sentence_transformers import SentenceTransformer

MODEL_CLASS_TO_NAME = {
    'gpt': list(OPENAI_GPT_PRETRAINED_CONFIG_ARCHIVE_MAP.keys()),
    'bert': list(BERT_PRETRAINED_CONFIG_ARCHIVE_MAP.keys()),
    'xlnet': list(XLNET_PRETRAINED_CONFIG_ARCHIVE_MAP.keys()),
    'roberta': list(ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP.keys()),
    'lstm': ['lstm'],
    't5': list(T5_PRETRAINED_CONFIG_ARCHIVE_MAP.keys()),
}
try:
    MODEL_CLASS_TO_NAME['albert'] = list(ALBERT_PRETRAINED_CONFIG_ARCHIVE_MAP.keys())
except:
    pass

MODEL_NAME_TO_CLASS = {model_name: model_class for model_class, model_name_list in MODEL_CLASS_TO_NAME.items() for model_name in model_name_list}

# Add SapBERT, PubMedBERT configuration
model_name = 'cambridgeltl/SapBERT-from-PubMedBERT-fulltext'
MODEL_NAME_TO_CLASS[model_name] = 'bert'
model_name = 'microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract'
MODEL_NAME_TO_CLASS[model_name] = 'bert'
model_name = 'michiyasunaga/BioLinkBERT-large'
MODEL_NAME_TO_CLASS[model_name] = 'bert'

# model_name = 'google/flan-t5-small'
# model_name = 'google/flan-t5-xl'
# model_name = 'google/flan-t5-xxl'
MODEL_NAME_TO_CLASS['google/flan-t5-small'] = 't5'
MODEL_NAME_TO_CLASS['google/flan-t5-base'] = 't5'
MODEL_NAME_TO_CLASS['google/flan-t5-large'] = 't5'
MODEL_NAME_TO_CLASS['google/flan-t5-xl'] = 't5'
MODEL_NAME_TO_CLASS['google/flan-t5-xxl'] = 't5'
MODEL_NAME_TO_CLASS['bigscience/bloomz-560m'] = 'bloomz'
MODEL_NAME_TO_CLASS['bigscience/bloomz-7b1'] = 'bloomz'
MODEL_NAME_TO_CLASS['bigscience/bloomz-3b'] = 'bloomz'

# print('MODEL_NAME_TO_CLASS: ', MODEL_NAME_TO_CLASS)


GPT_SPECIAL_TOKENS = ['_start_', '_delimiter_', '_classify_']


def bfs_based_on_k_hop_subgraph(
    node_idx: Union[int, List[int], Tensor],
    num_hops: int,
    edge_index: Tensor,
    relabel_nodes: bool = False,
    num_nodes: Optional[int] = None,
    flow: str = 'source_to_target',
    directed: bool = False,
) -> Tuple[Tensor, Tensor, Tensor, Tensor]:

    num_nodes = maybe_num_nodes(edge_index, num_nodes)

    assert flow in ['source_to_target', 'target_to_source']
    if flow == 'target_to_source':
        row, col = edge_index
    else:
        col, row = edge_index

    node_mask = row.new_empty(num_nodes, dtype=torch.bool)
    edge_mask = row.new_empty(row.size(0), dtype=torch.bool)

    if isinstance(node_idx, (int, list, tuple)):
        node_idx = torch.tensor([node_idx], device=row.device).flatten()
    else:
        node_idx = node_idx.to(row.device)

    subsets = [node_idx]

    for _ in range(num_hops):
        node_mask.fill_(False)
        try:
            node_mask[subsets[-1]] = True
            torch.index_select(node_mask, 0, row, out=edge_mask)
            subsets.append(col[edge_mask])
            # print('col[edge_mask]: ', col[edge_mask])
        except:
            break

    # print('subsets: ', subsets)
    subsets = torch.cat(subsets)

    # origin
    # subset = torch.cat(subsets).unique(sorted=False)  # return_inverse=True

    x_unique = torch.unique(subsets, sorted=True)
    x_idx = (torch.cat([(subsets == x_u).nonzero()[0] for x_u in x_unique])).sort()[0]

    return subsets[x_idx]
    # return subset
    # return subset, edge_index, inv, edge_mask


class MultiGPUSparseAdjDataBatchGenerator(object):
    """A data generator that batches the data and moves them to the corresponding devices."""

    def __init__(self, args, device0, device1, batch_size, indexes, qids, labels,
                 tensors0=[], lists0=[], tensors1=[], lists1=[], adj_data=None, tokenizer=None, id2concept=None):
        self.args = args
        self.device0 = device0
        self.device1 = device1
        self.batch_size = batch_size
        self.indexes = indexes
        self.qids = qids
        self.labels = labels
        self.tensors0 = tensors0
        self.lists0 = lists0
        self.tensors1 = tensors1
        self.lists1 = lists1
        self.adj_data = adj_data
        self.tokenizer = tokenizer
        self.id2concept = id2concept

        # self.mlm_probability = args.mlm_probability
        if args.span_mask:
            print('span_mask', args.span_mask, file=sys.stderr)
        self.geo_p = 0.2
        self.span_len_upper = 10
        self.span_len_lower = 1
        self.span_lens = list(range(self.span_len_lower, self.span_len_upper + 1))
        self.span_len_dist = [self.geo_p * (1-self.geo_p)**(i - self.span_len_lower) for i in range(self.span_len_lower, self.span_len_upper + 1)]
        self.span_len_dist = [x / (sum(self.span_len_dist)) for x in self.span_len_dist]

        self.eval_end_task_mode = False  # if True, use the non-modified text and KG inputs

        if self.args.baseline_flattern and 'triplet' in self.args.baseline_flattern:
            if self.args.dataset in ['medqa', 'pubmedqa', 'bioasq']:
                cpnet_vocab_path = self.args.kg_vocab_path  # "data/umls/concepts.txt"
                with open(cpnet_vocab_path, "r", encoding="utf8") as fin:
                    self.id2concept = [w.strip() for w in fin]
                self.concept2id = {w: i for i, w in enumerate(self.id2concept)}
                self.id2relation_text = [r.strip() for r in open(f"{os.path.dirname(self.args.kg_vocab_path)}/relations.txt")]
            else:
                # Load cpnet
                cpnet_vocab_path = self.args.kg_vocab_path  # "data/cpnet/concept.txt"
                with open(cpnet_vocab_path, "r", encoding="utf8") as fin:
                    self.id2concept = [w.strip() for w in fin]
                self.concept2id = {w: i for i, w in enumerate(self.id2concept)}
                # self.id2relation = conceptnet.merged_relations
                self.id2relation_text = conceptnet.relation_text

            if self.args.baseline_flattern_similarity:
                # TEXT_EMB_MODEL_TABLES = {
                #     "mpnet": "sentence-transformers/all-mpnet-base-v2",
                #     "msmarco": "sebastian-hofstaetter/distilbert-dot-tas_b-b256-msmarco"
                # }
                self.text_emb_model = SentenceTransformer('sentence-transformers/all-mpnet-base-v2')

    def __len__(self):
        return (self.indexes.size(0) - 1) // self.batch_size + 1

    def __iter__(self):
        bs = self.batch_size
        n = self.indexes.size(0)
        # print('bs: ', bs)
        for a in range(0, n, bs):
            b = min(n, a + bs)
            batch_indexes = self.indexes[a:b]
            batch_qids = [self.qids[idx] for idx in batch_indexes]
            batch_labels = self._to_device(self.labels[batch_indexes], self.device1)
            # print('batch_indexes: ', batch_indexes, batch_indexes.shape)

            batch_tensors0 = [self._to_device(x[batch_indexes], self.device1) for x in self.tensors0]
            assert len(batch_tensors0) == 4  # tensors0: all_input_ids, all_input_mask, all_segment_ids, all_output_mask

            # batch_lm_labels are masked word prediction labels
            batch_lm_inputs, batch_input_masks = batch_tensors0[0], batch_tensors0[1]
            # batch_lm_inputs, batch_lm_labels = batch_tensors0[0], batch_tensors0[1]
            # origin
            # batch_lm_inputs, batch_lm_labels = self.process_lm_data(batch_tensors0)

            # print('batch_lm_inputs: ', batch_lm_inputs.shape)
            # print('batch_lm_labels: ', batch_lm_labels)
            # print('batch_lm_labels: ', batch_lm_labels.shape)

            batch_tensors1 = [self._to_device(x[batch_indexes], self.device1) for x in self.tensors1]
            assert len(batch_tensors1) == 5  # tensors1: concept_ids, node_type_ids, node_scores, adj_lengths, special_nodes_mask
            batch_tensors1[0] = batch_tensors1[0].to(self.device0)
            # batch_lists0 = [self._to_device([x[i] for i in batch_indexes], self.device0) for x in self.lists0]
            # batch_lists1 = [self._to_device([x[i] for i in batch_indexes], self.device1) for x in self.lists1]

            edge_index_all, edge_type_all = self.adj_data
            # edge_index_all: nested list of shape (n_samples, num_choice), where each entry is tensor[2, E]
            # edge_type_all:  nested list of shape (n_samples, num_choice), where each entry is tensor[E, ]
            edge_index = self._to_device([edge_index_all[i] for i in batch_indexes], self.device1)
            edge_type = self._to_device([edge_type_all[i] for i in batch_indexes], self.device1)
            node_type_ids = batch_tensors1[1]  # [bs, nc, n_nodes]
            assert node_type_ids.dim() == 3
            edge_index, edge_type, pos_triples, neg_nodes = self.process_graph_data(edge_index, edge_type, node_type_ids)

            # validate the llm inputs
            # print('batch_qids: ', len(batch_qids))
            # print('batch_labels: ', batch_labels.shape)
            # print('batch_lm_inputs: ', batch_lm_inputs.shape)
            # # print('batch_lm_labels: ', batch_lm_labels.shape)  # masked word prediction labels

            # validate the gnn inputs
            # print('concept_ids: ', batch_tensors1[0].shape, batch_tensors1[0]) # 1*5*200
            # print('node_type_ids: ', batch_tensors1[1].shape, batch_tensors1[1])
            # print('node_scores: ', batch_tensors1[2].shape, batch_tensors1[2])
            # print('adj_lengths: ', batch_tensors1[3].shape, batch_tensors1[3])
            # print('special_nodes_mask: ', batch_tensors1[4].shape, batch_tensors1[4])
            # print('edge_index: ', len(edge_index[0][0][0]), edge_index)  # 1*5*2*(e.g., 614)
            # print('edge_type: ', len(edge_type[0]), edge_type)  # 1*4*xxx
            # print('pos_triples: ', len(pos_triples[0]), pos_triples) # 1*4*3*100
            # print('neg_nodes: ', len(neg_nodes[0]), neg_nodes)
            concept_ids = batch_tensors1[0]
            node_scores = batch_tensors1[2]
            adj_lengths = batch_tensors1[3]

            if not self.args.baseline_flattern:
                batch_graph_inputs = edge_index, edge_type, pos_triples, neg_nodes, concept_ids, node_type_ids, node_scores, adj_lengths
                yield tuple([batch_qids, batch_labels, batch_lm_inputs, batch_input_masks, batch_graph_inputs])

            else:
                # ----------------- START: baseline flattern_graph_str_list -----------------
                # prepare for the graph flattern prompt
                flattern_graph_str_list = []

                # Verify graph data loader by printing the concepts extracted for each question-answer pair
                concept_ids = batch_tensors1[0]

                if self.args.dataset in ['medqa', 'pubmedqa', 'bioasq']:
                    concept2concept_name_dict = {}
                    with open(self.args.data_dir + '/umls/concept_names.txt', "r", encoding="utf8") as f:
                        for line in f:
                            concept = line.split('\t')[0].strip()
                            concept_name = line.split('\t')[1].strip()
                            concept2concept_name_dict[concept] = concept_name
                    # print('concept2concept_name_dict: ', concept2concept_name_dict)

                    id2concept_for_baseline = []
                    for nid in self.id2concept:
                        id2concept_for_baseline.append(concept2concept_name_dict[nid])

                else:
                    id2concept_for_baseline = self.id2concept

                for i in range(len(concept_ids[0])):
                    # print('i: ', i)

                    # edge_type: 2-18 corresponds to 0-16 in self.id2relation. 0, 1 are edges from context node to q and a. 21-36 are inverse edges of 0-16. 19, 20 are inverse edges of context node to q and a.

                    # releverance score based methods
                    if self.args.baseline_flattern == 'rs':
                        concept_list_of_a_choice_rel_score = [id2concept_for_baseline[concept_ids[0][i][j]-1] if (concept_ids[0][i][j] != 0 and concept_ids[0][i][j] != 1 and node_type_ids[0][i][j] != 3 and node_type_ids[0][i][j] != 0 and node_type_ids[0][i][j] != 1) else '' for j in range(len(concept_ids[0][i]))]
                        # concept_list_of_a_choice_rel_score = [id2concept_for_baseline[concept_ids[0][i][j]-1] if (concept_ids[0][i][j] != 0 and concept_ids[0][i][j] != 1 and node_type_ids[0][i][j] != 3) else '' for j in range(len(concept_ids[0][i]))]  # low performance. need to filter q and a nodes.
                        # concept_list_of_a_choice_rel_score = concept_list_of_a_choice_rel_score[:30]
                        concept_of_a_choice_rel_score = ', '.join(t for t in concept_list_of_a_choice_rel_score if t)
                        concept_of_a_choice_rel_score = concept_of_a_choice_rel_score.strip()
                        # concept_of_a_choice_rel_score = concept_of_a_choice_rel_score.replace(' ', ', ')
                        concept_of_a_choice_rel_score = concept_of_a_choice_rel_score.replace('_', ' ')
                        concept_of_a_choice_rel_score = concept_of_a_choice_rel_score + '.'
                        flattern_graph_str_list.append(concept_of_a_choice_rel_score)
                        # print('concept_of_a_choice_rel_score: ', concept_of_a_choice_rel_score)
                        # print('---'*10)

                    # bfs method
                    elif self.args.baseline_flattern == 'bfs':
                        bfs_neighbors_node_idx = bfs_based_on_k_hop_subgraph(0, 10, edge_index[0][i])

                        # remove context, q, and a nodes. only include
                        new_bfs_neighbors_node_idx = []
                        for temp_idx in bfs_neighbors_node_idx:
                            if node_type_ids[0][i][temp_idx] != 3 and node_type_ids[0][i][temp_idx] != 0 and node_type_ids[0][i][temp_idx] != 1:
                                new_bfs_neighbors_node_idx.append(int(temp_idx))
                        bfs_neighbors_node_idx = new_bfs_neighbors_node_idx

                        bfs_neighbors = concept_ids[0][i][bfs_neighbors_node_idx]
                        # print('bfs_neighbors: ', bfs_neighbors)
                        # concept_of_a_choice_bfs = ' '.join([id2concept_for_baseline[temp-1] if (temp != 0 and temp != 1 and  else '' for temp in bfs_neighbors])
                        concept_of_a_choice_bfs = [id2concept_for_baseline[temp-1] if (temp != 0 and temp != 1 and temp != '') else '' for temp in bfs_neighbors]
                        concept_of_a_choice_bfs = ', '.join(t for t in concept_of_a_choice_bfs if t)
                        concept_of_a_choice_bfs = concept_of_a_choice_bfs.strip()
                        # concept_of_a_choice_bfs = concept_of_a_choice_bfs.replace(' ', ', ')
                        concept_of_a_choice_bfs = concept_of_a_choice_bfs.replace('_', ' ')
                        concept_of_a_choice_bfs = concept_of_a_choice_bfs + '.'
                        flattern_graph_str_list.append(concept_of_a_choice_bfs)
                        # print('concept_of_a_choice_bfs: ', concept_of_a_choice_bfs)

                    # elif self.args.baseline_flattern == 'triplet':
                    elif self.args.baseline_flattern and 'triplet' in self.args.baseline_flattern:
                        # print('edge_index: ', len(edge_index[0][0][0]), edge_index)  # 1*4*2*(e.g., 2498)
                        # print('edge_type: ', len(edge_type[0]), edge_type)  # 1*4*(e.g., 2498)
                        # print('pos_triples: ', len(pos_triples[0]), pos_triples) # 1*4*3*100
                        # print('neg_nodes: ', len(neg_nodes[0]), neg_nodes) # 1*4*100*64
                        # breakpoint()
                        all_triplets_list = []
                        all_triplets_list_for_case_study = []
                        node_type_in_triplets_for_case_study = {}
                        
                        if self.args.dataset in ['medqa', 'pubmedqa', 'bioasq']:
                            valid_edge_type_id_list = [val for val in range(2, 100)]
                            # breakpoint()
                        else:
                            valid_edge_type_id_list = [val for val in range(2, 19)]
                        for aa in range(len(edge_index[0])):
                            src_node_list = edge_index[0][aa][0].tolist()
                            dst_node_list = edge_index[0][aa][1].tolist()
                            edge_type_for_each_subgraph_list = edge_type[0][aa].tolist()
                            assert len(src_node_list) == len(dst_node_list) == len(edge_type_for_each_subgraph_list)
                            # print('len(src_node_list): ', len(src_node_list))
                            for bb in range(len(src_node_list)):
                                cur_src_node_id = src_node_list[bb]
                                cur_dst_node_id = dst_node_list[bb]
                                cur_edge_type_id = edge_type_for_each_subgraph_list[bb]
                                # filter triplets based on cur_edge_type_id
                                if cur_edge_type_id in valid_edge_type_id_list:
                                    if self.args.baseline_flattern == 'triplet_a':
                                        # only focus on the a nodes
                                        # if int(node_type_ids[0][i][cur_src_node_id]) == 1 or int(node_type_ids[0][i][cur_dst_node_id]) == 1:
                                        if not (int(node_type_ids[0][i][cur_src_node_id]) == 1 or int(node_type_ids[0][i][cur_dst_node_id]) == 1):
                                            continue
                                    elif self.args.baseline_flattern == 'triplet_q':
                                        # only focus on the q nodes
                                        if not (int(node_type_ids[0][i][cur_src_node_id]) == 0 or int(node_type_ids[0][i][cur_dst_node_id]) == 0):
                                            continue
                                    elif self.args.baseline_flattern == 'triplet_qa':
                                        # focus on both q and a nodes
                                        if not (int(node_type_ids[0][i][cur_src_node_id]) == 0 or int(node_type_ids[0][i][cur_dst_node_id]) == 0 or int(node_type_ids[0][i][cur_src_node_id]) == 1 or int(node_type_ids[0][i][cur_dst_node_id]) == 1):
                                            continue
                                    elif self.args.baseline_flattern == 'every_triplet':
                                        # focus on both q and a nodes
                                        if not (int(node_type_ids[0][i][cur_src_node_id]) == 0 or int(node_type_ids[0][i][cur_dst_node_id]) == 0 or int(node_type_ids[0][i][cur_src_node_id]) == 1 or int(node_type_ids[0][i][cur_dst_node_id]) == 1 or int(node_type_ids[0][i][cur_src_node_id]) == 2 or int(node_type_ids[0][i][cur_dst_node_id]) == 2):
                                            continue

                                    cur_src_node_concept_id = int(concept_ids[0][i][cur_src_node_id])
                                    cur_dst_node_concept_id = int(concept_ids[0][i][cur_dst_node_id])
                                    if cur_src_node_concept_id == 0 or cur_src_node_concept_id == 1:
                                        continue
                                    if cur_dst_node_concept_id == 0 or cur_dst_node_concept_id == 1:
                                        continue
                                    cur_src_node_concept = id2concept_for_baseline[cur_src_node_concept_id-1]
                                    cur_dst_node_concept = id2concept_for_baseline[cur_dst_node_concept_id-1]
                                    cur_edge_type_str = self.id2relation_text[cur_edge_type_id-2]
                                    # cur_triplet = cur_src_node_concept + ', ' + cur_edge_type_str + ', ' + cur_dst_node_concept
                                    cur_triplet = cur_src_node_concept + ' ' + cur_edge_type_str + ' ' + cur_dst_node_concept
                                    cur_triplet = cur_triplet.strip()
                                    cur_triplet = cur_triplet.replace('_', ' ')
                                    cur_triplet = '(' + cur_triplet + ')'
                                    # print('cur_triplet: ', cur_triplet)
                                    all_triplets_list.append(cur_triplet)
                                    if self.args.case_study:
                                        all_triplets_list_for_case_study.append((cur_src_node_concept.replace('_', ' '), cur_edge_type_str.replace('_', ' '), cur_dst_node_concept.replace('_', ' ')))
                                        # check_for_two_hop = True
                                        # src q
                                        if int(node_type_ids[0][i][cur_src_node_id]) == 0:
                                            node_type_in_triplets_for_case_study[cur_src_node_concept] = 'q'
                                            # if int(node_type_ids[0][i][cur_dst_node_id]) != 0 and int(node_type_ids[0][i][cur_dst_node_id]) != 1:
                                            #     node_type_in_triplets_for_case_study[cur_dst_node_concept] = 'one-hop'
                                            # check_for_two_hop = False
                                        # src a
                                        elif int(node_type_ids[0][i][cur_src_node_id]) == 1:
                                            node_type_in_triplets_for_case_study[cur_src_node_concept] = 'a'
                                            # if int(node_type_ids[0][i][cur_dst_node_id]) != 0 and int(node_type_ids[0][i][cur_dst_node_id]) != 1:
                                            #     node_type_in_triplets_for_case_study[cur_dst_node_concept] = 'one-hop'
                                            # check_for_two_hop = False
                                                
                                        # dst q
                                        if int(node_type_ids[0][i][cur_dst_node_id]) == 0:
                                            node_type_in_triplets_for_case_study[cur_dst_node_concept] = 'q'
                                            # if int(node_type_ids[0][i][cur_src_node_id]) != 0 and int(node_type_ids[0][i][cur_src_node_id]) != 1:
                                            #     node_type_in_triplets_for_case_study[cur_src_node_concept] = 'one-hop'
                                            # check_for_two_hop = False
                                        elif int(node_type_ids[0][i][cur_dst_node_id]) == 1:
                                            node_type_in_triplets_for_case_study[cur_dst_node_concept] = 'a'
                                            # if int(node_type_ids[0][i][cur_src_node_id]) != 0 and int(node_type_ids[0][i][cur_src_node_id]) != 1:
                                            #     node_type_in_triplets_for_case_study[cur_src_node_concept] = 'one-hop'
                                            # check_for_two_hop = False
                                        
                                        # if check_for_two_hop:
                                        #     node_type_in_triplets_for_case_study[cur_src_node_concept] = 'two-hop'
                                        #     node_type_in_triplets_for_case_study[cur_dst_node_concept] = 'two-hop'

                        # remove duplicates
                        all_triplets_list = list(set(all_triplets_list))
                        # print('len(all_triplets_list): ', len(all_triplets_list))

                        # calculate similarity between graph concepts and text question and answer.
                        if self.args.baseline_flattern_similarity:
                            num_batch = 20
                            # print('torch.LongTensor(batch_lm_inputs): ', torch.LongTensor(batch_lm_inputs))
                            text_input = self.tokenizer.batch_decode(batch_lm_inputs[0], skip_special_tokens=True)
                            text_input = text_input[0]
                            # print('text_input: ', text_input)

                            similarities = []
                            query_emb = self.text_emb_model.encode(text_input)
                            # pool = self.text_emb_model.start_multi_process_pool()

                            for batch_index in range(0, ceil(len(all_triplets_list) / num_batch)):
                                start_index = batch_index * num_batch
                                end_index = (batch_index + 1) * num_batch

                                # triplets_embs = self.text_emb_model.encode_multi_process(all_triplets_list[start_index : end_index], pool, batch_size=num_batch)
                                triplets_embs = self.text_emb_model.encode(all_triplets_list[start_index: end_index], batch_size=num_batch)
                                similarities.extend(dot_score(query_emb, triplets_embs)[0].tolist())
                                # similarities.extend(sentence_transformers.util.dot_score(query_emb, triplets_embs)[0].tolist())

                            similarities = np.array(similarities)
                            # print('similarities: ', similarities)
                            sorted_indices = np.argsort(-similarities)
                            sorted_triplets_list = [all_triplets_list[index] for index in sorted_indices]
                            all_triplets_list = sorted_triplets_list[:10]  # top 10 teiplets
                                

                        if self.args.case_study:
                            all_triplets_list = all_triplets_list_for_case_study
                            flattern_graph_str_list.extend(all_triplets_list)
                        else:
                            all_triplets_list = ', '.join(all_triplets_list)
                            # breakpoint()
                            flattern_graph_str_list.append(all_triplets_list)
                
                if self.args.case_study:
                    flattern_graph_str_list = [flattern_graph_str_list, node_type_in_triplets_for_case_study]
                    
                # # print('length of tuple output: ', len(tuple([batch_qids, batch_labels, batch_lm_inputs, batch_lm_labels, *batch_tensors0, *batch_lists0, *batch_tensors1, *batch_lists1, edge_index, edge_type, pos_triples, neg_nodes])))
                # # yield tuple([batch_qids, batch_labels, batch_lm_inputs, batch_lm_labels])
                # yield tuple([batch_qids, batch_labels, batch_lm_inputs, flattern_graph_str_list])
                batch_graph_inputs = edge_index, edge_type, pos_triples, neg_nodes, concept_ids, node_type_ids, node_scores, adj_lengths, flattern_graph_str_list
                yield tuple([batch_qids, batch_labels, batch_lm_inputs, batch_input_masks, batch_graph_inputs])
                # ----------------- END: baseline flattern_graph_str_list -----------------

            # origin tuple output, length 17
            # yield tuple([batch_qids, batch_labels, batch_lm_inputs, batch_lm_labels, *batch_tensors0, *batch_lists0, *batch_tensors1, *batch_lists1, edge_index, edge_type, pos_triples, neg_nodes])

    def _to_device(self, obj, device):
        if isinstance(obj, (tuple, list)):
            return [self._to_device(item, device) for item in obj]
        else:
            return obj.to(device)

    def set_eval_end_task_mode(self, flag: bool):
        self.eval_end_task_mode = flag

    def process_lm_data(self, batch_tensors0):
        input_ids, special_tokens_mask = batch_tensors0[0], batch_tensors0[3]
        assert input_ids.dim() == 3 and special_tokens_mask.dim() == 3
        _bs, _nc, _seqlen = input_ids.size()

        _inputs = input_ids.clone().view(-1, _seqlen)  # remember to clone input_ids
        _mask_labels = []
        for ex in _inputs:
            if self.args.span_mask:
                _mask_label = self._span_mask(self.tokenizer.convert_ids_to_tokens(ex))
            else:
                _mask_label = self._word_mask(self.tokenizer.convert_ids_to_tokens(ex))
            _mask_labels.append(_mask_label)
        _mask_labels = torch.tensor(_mask_labels, device=_inputs.device)

        batch_lm_inputs, batch_lm_labels = self.mask_tokens(inputs=_inputs, mask_labels=_mask_labels, special_tokens_mask=special_tokens_mask.view(-1, _seqlen))

        batch_lm_inputs = batch_lm_inputs.view(_bs, _nc, _seqlen)  # this is masked
        batch_lm_labels = batch_lm_labels.view(_bs, _nc, _seqlen)

        if self.eval_end_task_mode or (self.args.mlm_task == 0):
            batch_lm_inputs = input_ids  # non-modified input

        return batch_lm_inputs, batch_lm_labels

    def mask_tokens(self, inputs, mask_labels, special_tokens_mask=None):
        """
        Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original.
        """
        assert inputs.size() == mask_labels.size()

        if self.tokenizer.mask_token is None:
            raise ValueError(
                "This tokenizer does not have a mask token which is necessary for masked language modeling. Remove the --mlm flag if you want to use this tokenizer."
            )
        labels = inputs.clone()
        # We sample a few tokens in each sequence for masked-LM training (with probability args.mlm_probability defaults to 0.15 in Bert/RoBERTa)

        probability_matrix = mask_labels

        if special_tokens_mask is None:
            special_tokens_mask = [
                self.tokenizer.get_special_tokens_mask(val, already_has_special_tokens=True) for val in labels.tolist()
            ]
            special_tokens_mask = torch.tensor(special_tokens_mask, dtype=torch.bool)
        else:
            special_tokens_mask = special_tokens_mask.bool()

        probability_matrix.masked_fill_(special_tokens_mask, value=0.0)
        # if self.tokenizer._pad_token is not None: #should be handled already
        #     padding_mask = labels.eq(self.tokenizer.pad_token_id)
        #     probability_matrix.masked_fill_(padding_mask, value=0.0)

        masked_indices = probability_matrix.bool()
        labels[~masked_indices] = -100  # We only compute loss on masked tokens

        # 80% of the time, we replace masked input tokens with tokenizer.mask_token ([MASK])
        indices_replaced = torch.bernoulli(torch.full(labels.shape, 0.8, device=labels.device)).bool() & masked_indices
        inputs[indices_replaced] = self.tokenizer.convert_tokens_to_ids(self.tokenizer.mask_token)

        # 10% of the time, we replace masked input tokens with random word
        indices_random = torch.bernoulli(torch.full(labels.shape, 0.5, device=labels.device)).bool() & masked_indices & ~indices_replaced
        random_words = torch.randint(len(self.tokenizer), labels.shape, dtype=torch.long, device=labels.device)
        inputs[indices_random] = random_words[indices_random]

        # The rest of the time (10% of the time) we keep the masked input tokens unchanged
        return inputs, labels

    def _word_mask(self, input_tokens: List[str], max_predictions=512):
        """
        Get 0/1 labels for masking tokens at word level
        """
        effective_num_toks = 0
        cand_indexes = []
        if isinstance(self.tokenizer, (BertTokenizer, BertTokenizerFast)):
            after_special_tok = False
            for (i, token) in enumerate(input_tokens):
                if token in ["[CLS]", "[SEP]", "[PAD]"]:
                    after_special_tok = True
                    continue
                if len(cand_indexes) >= 1 and (not after_special_tok) and token.startswith("##"):
                    cand_indexes[-1].append(i)
                else:
                    cand_indexes.append([i])
                after_special_tok = False
                effective_num_toks += 1
        elif isinstance(self.tokenizer, (RobertaTokenizer, RobertaTokenizerFast, AutoTokenizer, PreTrainedTokenizerFast)):
            after_special_tok = False
            for (i, token) in enumerate(input_tokens):
                if token in ["<s>",  "</s>", "<pad>"]:
                    after_special_tok = True
                    continue
                if len(cand_indexes) >= 1 and (not after_special_tok) and (not token.startswith("Ġ")):
                    cand_indexes[-1].append(i)
                else:
                    cand_indexes.append([i])
                after_special_tok = False
                effective_num_toks += 1
        else:
            raise NotImplementedError

        random.shuffle(cand_indexes)
        num_to_predict = min(max_predictions, max(1, int(round(effective_num_toks * self.mlm_probability))))
        masked_lms = []
        covered_indexes = set()
        for index_set in cand_indexes:
            if len(masked_lms) >= num_to_predict:
                break
            # If adding a whole-word mask would exceed the maximum number of
            # predictions, then just skip this candidate.
            if len(masked_lms) + len(index_set) > num_to_predict:
                continue
            is_any_index_covered = False
            for index in index_set:
                if index in covered_indexes:
                    is_any_index_covered = True
                    break
            if is_any_index_covered:
                continue
            for index in index_set:
                covered_indexes.add(index)
                masked_lms.append(index)

        assert len(covered_indexes) == len(masked_lms)
        mask_labels = [1 if i in covered_indexes else 0 for i in range(len(input_tokens))]
        return mask_labels

    def _span_mask(self, input_tokens: List[str], max_predictions=512):
        """
        Get 0/1 labels for masking tokens at word level
        """
        effective_num_toks = 0
        cand_indexes = []
        if isinstance(self.tokenizer, (BertTokenizer, BertTokenizerFast)):
            after_special_tok = False
            for (i, token) in enumerate(input_tokens):
                if token in ["[CLS]", "[SEP]", "[PAD]"]:
                    after_special_tok = True
                    continue
                if len(cand_indexes) >= 1 and (not after_special_tok) and token.startswith("##"):
                    cand_indexes[-1].append(i)
                else:
                    cand_indexes.append([i])
                after_special_tok = False
                effective_num_toks += 1
        elif isinstance(self.tokenizer, (RobertaTokenizer, RobertaTokenizerFast)):
            after_special_tok = False
            for (i, token) in enumerate(input_tokens):
                if token in ["<s>",  "</s>", "<pad>"]:
                    after_special_tok = True
                    continue
                if len(cand_indexes) >= 1 and (not after_special_tok) and (not token.startswith("Ġ")):
                    cand_indexes[-1].append(i)
                else:
                    cand_indexes.append([i])
                after_special_tok = False
                effective_num_toks += 1
        else:
            raise NotImplementedError
        cand_indexes_args = list(range(len(cand_indexes)))

        random.shuffle(cand_indexes_args)
        num_to_predict = min(max_predictions, max(1, int(round(effective_num_toks * self.mlm_probability))))
        masked_lms = []
        covered_indexes = set()
        for wid in cand_indexes_args:
            if len(masked_lms) >= num_to_predict:
                break
            span_len = np.random.choice(self.span_lens, p=self.span_len_dist)
            # If adding a whole-word mask would exceed the maximum number of
            # predictions, then just skip this candidate.
            # if len(masked_lms) + span_len > num_to_predict:
            #     continue
            index_set = []
            is_any_index_covered = False
            for _wid in range(wid, len(cand_indexes)):  # iterate over word
                if len(index_set) + len(cand_indexes[_wid]) > span_len:
                    break
                for _index in cand_indexes[_wid]:  # iterate over subword
                    if _index in covered_indexes:
                        is_any_index_covered = True
                        break
                    index_set.append(_index)
                if is_any_index_covered:
                    break
            if is_any_index_covered:
                continue
            for _index in index_set:
                covered_indexes.add(_index)
                masked_lms.append(_index)

        assert len(covered_indexes) == len(masked_lms)
        mask_labels = [1 if i in covered_indexes else 0 for i in range(len(input_tokens))]
        return mask_labels

    def process_graph_data(self, edge_index, edge_type, node_type_ids):
        # edge_index: nested list of shape (n_samples, num_choice), where each entry is tensor[2, E]
        # edge_type:  nested list of shape (n_samples, num_choice), where each entry is tensor[E, ]
        # node_type_ids: tensor[n_samples, num_choice, num_nodes]
        bs, nc = len(edge_index), len(edge_index[0])
        input_edge_index, input_edge_type, pos_triples, neg_nodes = [], [], [], []
        for bid in range(bs):
            for cid in range(nc):
                _edge_index = edge_index[bid][cid]  # .clone()
                _edge_type = edge_type[bid][cid]  # .clone()
                _node_type_ids = node_type_ids[bid][cid]  # .clone()
                _edge_index, _edge_type, _pos_triples, _neg_nodes = self._process_one_graph(_edge_index, _edge_type, _node_type_ids)
                input_edge_index.append(_edge_index)
                input_edge_type.append(_edge_type)
                pos_triples.append(_pos_triples)
                neg_nodes.append(_neg_nodes)
        input_edge_index = list(map(list, zip(*(iter(input_edge_index),) * nc)))  # nested list of shape (n_samples, num_choice)
        input_edge_type = list(map(list, zip(*(iter(input_edge_type),) * nc)))
        pos_triples = list(map(list, zip(*(iter(pos_triples),) * nc)))
        neg_nodes = list(map(list, zip(*(iter(neg_nodes),) * nc)))

        if self.eval_end_task_mode or (self.args.link_task == 0):
            input_edge_index = edge_index  # non-modified input
            input_edge_type = edge_type  # non-modified input

        return input_edge_index, input_edge_type, pos_triples, neg_nodes

    def _process_one_graph(self, _edge_index, _edge_type, _node_type_ids):
        # _edge_index: tensor[2, E]
        # _edge_type:  tensor[E, ]
        # _node_type_ids: tensor[n_nodes, ]
        E = len(_edge_type)
        if E == 0:
            # print ('KG with 0 node', file=sys.stderr)
            effective_num_nodes = 1
        else:
            effective_num_nodes = int(_edge_index.max()) + 1
        device = _edge_type.device

        tmp = _node_type_ids.max().item()
        assert isinstance(tmp, int) and 0 <= tmp <= 5
        _edge_index_node_type = _node_type_ids[_edge_index]  # [2, E]
        _is_special = (_edge_index_node_type == 3)  # [2, E]
        is_special = _is_special[0] | _is_special[1]  # [E,]
        is_special = is_special.to('cpu')  # for advanced torch version

        positions = torch.arange(E)
        # print('positions: ', positions, flush=True)
        # print('is_special: ', is_special, flush=True)
        positions = positions[~is_special]  # [some_E, ]
        drop_count = min(self.args.link_drop_max_count, int(len(positions) * self.args.link_drop_probability))
        if len(positions) > 0 and drop_count > 0:
            drop_idxs = torch.multinomial(torch.full((len(positions),), 1.0), drop_count, replacement=False)  # [drop_count, ]
        else:
            drop_idxs = torch.tensor([]).long()
        drop_positions = positions[drop_idxs]  # [drop_count, ]

        mask = torch.zeros((E,)).long()  # [E, ]
        mask = mask.index_fill_(dim=0, index=drop_positions, value=1).bool().to(device)  # [E, ]

        real_drop_count = int(drop_count * (1-self.args.link_drop_probability_in_which_keep))
        real_drop_positions = positions[drop_idxs[:real_drop_count]]  # [real_drop_count, ]
        real_mask = torch.zeros((E,)).long()  # [E, ]
        real_mask = real_mask.index_fill_(dim=0, index=real_drop_positions, value=1).bool().to(device)  # [E, ]

        assert int(mask.long().sum()) == drop_count
        # print (f'drop_E / total_E = {drop_count} / {E} = {drop_count / E}', ) #E is typically 1000-3000
        input_edge_index = _edge_index[:, ~real_mask]
        input_edge_type = _edge_type[~real_mask]
        assert input_edge_index.size(1) == E - real_drop_count

        pos_edge_index = _edge_index[:, mask]
        pos_edge_type = _edge_type[mask]
        pos_triples = [pos_edge_index[0], pos_edge_type, pos_edge_index[1]]
        # pos_triples: list[h, r, t], where each of h, r, t is [n_triple, ]
        assert pos_edge_index.size(1) == drop_count

        num_edges = len(pos_edge_type)
        num_corruption = self.args.link_negative_sample_size
        neg_nodes = torch.randint(0, effective_num_nodes, (num_edges, num_corruption), device=device)  # [n_triple, n_neg]
        return input_edge_index, input_edge_type, pos_triples, neg_nodes


class DataLoader(object):

    def __init__(self, args, train_statement_path, train_adj_path,
                 dev_statement_path, dev_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"):
        super().__init__()
        self.args = args
        self.batch_size = batch_size
        self.eval_batch_size = eval_batch_size
        self.device0, self.device1 = device
        self.is_inhouse = is_inhouse
        self.debug = debug
        self.model_name = model_name
        self.max_node_num = max_node_num
        self.debug_sample_size = 32
        self.cxt_node_connects_all = cxt_node_connects_all

        self.model_type = MODEL_NAME_TO_CLASS[model_name]

        # Load the tokenizer
        try:
            tokenizer_class = {'bert': BertTokenizer, 'xlnet': XLNetTokenizer, 'roberta': RobertaTokenizer, 'albert': AlbertTokenizer, 't5': AutoTokenizer, 'bloomz': AutoTokenizer}.get(self.model_type)
        except:
            tokenizer_class = {'bert': BertTokenizer, 'xlnet': XLNetTokenizer, 'roberta': RobertaTokenizer, 't5': AutoTokenizer, 'bloomz': AutoTokenizer}.get(self.model_type)
        tokenizer = tokenizer_class.from_pretrained(self.model_name)
        # print('tokenizer: ', tokenizer)

        # gpt
        if self.model_type == 'bloomz':
            tokenizer.padding_side = "left"
            tokenizer.pad_token = tokenizer.eos_token

        self.tokenizer = tokenizer

        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, mode='train')
        print('(len(self.train_qids): ', len(self.train_qids))
        print('(len(self.train_labels): ', len(self.train_labels))
        print('len(self.train_encoder_data): ', len(self.train_encoder_data))
        print('x.size(0): ', self.train_encoder_data[0].size(0))
        # print('train_concepts_by_sents_list: ', len(train_concepts_by_sents_list))  # 0
        # print('train_concepts_by_sents_list: ', train_concepts_by_sents_list[0])

        # print('self.train_encoder_data[0]: ', self.train_encoder_data[0].shape)
        # num_choice = self.train_encoder_data[0].size(1)
        # self.num_choice = num_choice
        # print ('num_choice', num_choice)
        if args.dataset == 'csqa':
            self.num_choice = 5
        elif args.dataset == 'obqa':
            self.num_choice = 4
        elif args.dataset == 'riddle':
            self.num_choice = 5
        elif args.dataset == 'medqa':
            self.num_choice = 4
        elif args.dataset == 'arc':
            self.num_choice = 4
        elif args.dataset == 'piqa':
            self.num_choice = 2
        elif args.dataset == 'pubmedqa':
            self.num_choice = 3
        elif args.dataset == 'bioasq':
            self.num_choice = 2
        elif args.dataset == 'cosmosqa':
            self.num_choice = 4
        else:
            raise ValueError("Please indicate how many answer choices this dataset has.")
        print('self.num_choice', self.num_choice)
        *self.train_decoder_data, self.train_adj_data = self.load_sparse_adj_data_with_contextnode(train_adj_path, max_node_num, train_concepts_by_sents_list, mode='train')
        # print('len(self.train_adj_data[0][0]): ', self.train_adj_data[0][0])
        # print('len(self.train_adj_data[0]): ', len(self.train_adj_data[0]))
        # print('len(self.train_adj_data[1][0]): ', self.train_adj_data[1][0])
        # print('len(self.train_adj_data[1]): ', len(self.train_adj_data[1]))
        if not debug:
            assert all(len(self.train_qids) == len(self.train_adj_data[0]) == x.size(0) for x in [self.train_labels] + self.train_encoder_data + self.train_decoder_data)

        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_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)

        print("Finish loading dev data.")

        # 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_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)

            print("Finish loading test data.")

        # If using inhouse split, we split the original training set into an inhouse training set and an inhouse test set.
        if self.is_inhouse:
            with open(inhouse_train_qids_path, 'r') as fin:
                inhouse_qids = set(line.strip() for line in fin)
            self.inhouse_train_indexes = torch.tensor([i for i, qid in enumerate(self.train_qids) if qid in inhouse_qids])
            self.inhouse_test_indexes = torch.tensor([i for i, qid in enumerate(self.train_qids) if qid not in inhouse_qids])

        # Optionally we can subsample the training set.
        assert 0. < subsample <= 1.
        if subsample < 1. or n_train >= 0:
            # n_train will override subsample if the former is not None
            if n_train == -1:
                n_train = int(self.train_size() * subsample)
            assert n_train > 0
            if self.is_inhouse:
                self.inhouse_train_indexes = self.inhouse_train_indexes[:n_train]
            else:
                self.train_qids = self.train_qids[:n_train]
                self.train_labels = self.train_labels[:n_train]
                self.train_encoder_data = [x[:n_train] for x in self.train_encoder_data]
                self.train_decoder_data = [x[:n_train] for x in self.train_decoder_data]
                self.train_adj_data = tuple([x[:n_train] for x in self.train_adj_data])
                assert all(len(self.train_qids) == len(self.train_adj_data[0]) == x.size(0) for x in [self.train_labels] + self.train_encoder_data + self.train_decoder_data)
            assert self.train_size() == n_train

    def train_size(self):
        return self.inhouse_train_indexes.size(0) if self.is_inhouse else len(self.train_qids)

    def dev_size(self):
        return len(self.dev_qids)

    def test_size(self):
        if self.is_inhouse:
            return self.inhouse_test_indexes.size(0)
        else:
            return len(self.test_qids) if hasattr(self, 'test_qids') else 0

    def train(self, steps=-1, local_rank=-1):
        if self.debug:
            train_indexes = torch.arange(self.debug_sample_size)
        elif self.is_inhouse:
            n_train = self.inhouse_train_indexes.size(0)
            train_indexes = self.inhouse_train_indexes[torch.randperm(n_train)]
        else:
            train_indexes = torch.randperm(len(self.train_qids))
        if steps != -1:
            train_indexes = train_indexes[: self.args.batch_size * steps]

        # print('local_rank', local_rank, 'len(train_indexes)', len(train_indexes), 'train_indexes[:10]', train_indexes[:10].tolist())
        # print('local_rank', local_rank, 'len(train_indexes)', len(train_indexes), 'train_indexes[:10]', train_indexes[:10].tolist(), file=sys.stderr)
        return MultiGPUSparseAdjDataBatchGenerator(self.args, self.device0, self.device1, self.batch_size, train_indexes, self.train_qids, self.train_labels, tensors0=self.train_encoder_data, tensors1=self.train_decoder_data, adj_data=self.train_adj_data, tokenizer=self.tokenizer, id2concept=self.id2concept)

    def train_eval(self):
        # origin
        # return MultiGPUSparseAdjDataBatchGenerator(self.args, self.device0, self.device1, self.eval_batch_size, torch.arange(len(self.train_qids)), self.train_qids, self.train_labels, tensors0=self.train_encoder_data, tensors1=self.train_decoder_data, adj_data=self.train_adj_data, tokenizer=self.tokenizer, id2concept=self.id2concept)

        # added
        # same_length_with_dev = len(self.dev_qids)
        same_length_with_dev = 128
        return MultiGPUSparseAdjDataBatchGenerator(self.args, self.device0, self.device1, self.eval_batch_size, torch.arange(len(self.train_qids[:same_length_with_dev])), self.train_qids[:same_length_with_dev], self.train_labels[:same_length_with_dev], tensors0=self.train_encoder_data, tensors1=self.train_decoder_data, adj_data=self.train_adj_data, tokenizer=self.tokenizer, id2concept=self.id2concept)

    def dev(self):
        if self.debug:
            dev_indexes = torch.arange(self.debug_sample_size)
        else:
            dev_indexes = torch.arange(len(self.dev_qids))
        return MultiGPUSparseAdjDataBatchGenerator(self.args, self.device0, self.device1, self.eval_batch_size, dev_indexes, self.dev_qids, self.dev_labels, tensors0=self.dev_encoder_data, tensors1=self.dev_decoder_data, adj_data=self.dev_adj_data, tokenizer=self.tokenizer, id2concept=self.id2concept)

    def test(self):
        if self.debug:
            test_indexes = torch.arange(self.debug_sample_size)
        elif self.is_inhouse:
            test_indexes = self.inhouse_test_indexes
        else:
            test_indexes = torch.arange(len(self.test_qids))
        if self.is_inhouse:
            return MultiGPUSparseAdjDataBatchGenerator(self.args, self.device0, self.device1, self.eval_batch_size, test_indexes, self.train_qids, self.train_labels, tensors0=self.train_encoder_data, tensors1=self.train_decoder_data, adj_data=self.train_adj_data, tokenizer=self.tokenizer, id2concept=self.id2concept)
        else:
            return MultiGPUSparseAdjDataBatchGenerator(self.args, self.device0, self.device1, self.eval_batch_size, test_indexes, self.test_qids, self.test_labels, tensors0=self.test_encoder_data, tensors1=self.test_decoder_data, adj_data=self.test_adj_data, tokenizer=self.tokenizer, id2concept=self.id2concept)

    def load_resources(self, kg):
        if kg == "cpnet":
            # Load cpnet
            cpnet_vocab_path = self.args.kg_vocab_path  # "data/cpnet/concept.txt"
            with open(cpnet_vocab_path, "r", encoding="utf8") as fin:
                self.id2concept = [w.strip() for w in fin]
            self.concept2id = {w: i for i, w in enumerate(self.id2concept)}
            self.id2relation = conceptnet.merged_relations
        elif kg == "ddb":
            cpnet_vocab_path = self.args.kg_vocab_path  # "data/ddb/vocab.txt"
            with open(cpnet_vocab_path, "r", encoding="utf8") as fin:
                self.id2concept = [w.strip() for w in fin]
            self.concept2id = {w: i for i, w in enumerate(self.id2concept)}
            self.id2relation = [
                'belongstothecategoryof',
                'isacategory',
                'maycause',
                'isasubtypeof',
                'isariskfactorof',
                'isassociatedwith',
                'maycontraindicate',
                'interactswith',
                'belongstothedrugfamilyof',
                'child-parent',
                'isavectorfor',
                'mabeallelicwith',
                'seealso',
                'isaningradientof',
                'mabeindicatedby'
            ]
        elif kg == "umls":
            cpnet_vocab_path = self.args.kg_vocab_path  # "data/umls/concepts.txt"
            with open(cpnet_vocab_path, "r", encoding="utf8") as fin:
                self.id2concept = [w.strip() for w in fin]
            self.concept2id = {w: i for i, w in enumerate(self.id2concept)}
            self.id2relation = [r.strip() for r in open(f"{os.path.dirname(self.args.kg_vocab_path)}/relations.txt")]
        else:
            raise ValueError("Invalid value for kg.")

    def load_input_tensors(self, input_jsonl_path, max_seq_length, mode='eval'):
        """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 "") + '_prompt_design_' + str(self.args.prompt_design) + '.loaded_cache'
        # use_cache = True
        # use_cache = False
        use_cache = self.args.load_llm_cache

        if use_cache and not os.path.exists(cache_path):
            use_cache = False

        if use_cache:
            print(f'Loading cache {cache_path}')
            print(f'Loading cache {cache_path}', file=sys.stderr)
            # with open(cache_path, 'rb') as f:
            #     input_tensors = utils.CPU_Unpickler(f).load()
            input_tensors = ()
            with open(cache_path, "rb") as in_file:
                try:
                    while True:
                        obj = pickle.load(in_file)
                        if type(obj) == dict:
                            assert len(obj) == 1
                            key = list(obj.keys())[0]
                            input_tensors = input_tensors + (obj[key], )
                        elif type(obj) == tuple:
                            assert len(obj) == 4  # example_ids, all_label, data_tensors, concepts_by_sents_list
                            input_tensors = obj
                        else:
                            raise TypeError("Invalid type for obj.")
                except EOFError:
                    pass
            print(f'Loaded cache {cache_path}', file=sys.stderr)
        else:
            if self.model_type in ('lstm',):
                raise NotImplementedError
            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', 't5', 'bloomz'):
                input_tensors = load_bert_xlnet_roberta_input_tensors(input_jsonl_path, max_seq_length, self.debug, self.tokenizer, self.debug_sample_size, self.args)
            if not self.debug and not self.args.prompt_design:
                if self.args.local_rank in [-1, 0]:
                    print('saving cache...', file=sys.stderr)
                    # utils.save_pickle(input_tensors, cache_path)
                    with open(cache_path, 'wb') as f:
                        for _i_, obj in enumerate(tqdm(input_tensors)):
                            pickle.dump({f'obj{_i_}': obj}, f, protocol=4)
                    print('saved cache.', file=sys.stderr)

        if mode == 'train' and self.args.local_rank != -1:
            example_ids, all_label, data_tensors, concepts_by_sents_list = input_tensors  # concepts_by_sents_list is always []
            assert len(example_ids) == len(all_label) == len(data_tensors[0])
            total_num = len(data_tensors[0])
            rem = total_num % self.args.world_size
            if rem != 0:
                example_ids = example_ids + example_ids[:self.args.world_size - rem]
                all_label = torch.cat([all_label, all_label[:self.args.world_size - rem]], dim=0)
                data_tensors = [torch.cat([t, t[:self.args.world_size - rem]], dim=0) for t in data_tensors]
                total_num_aim = total_num + self.args.world_size - rem
            else:
                total_num_aim = total_num
            assert total_num_aim % self.args.world_size == 0
            assert total_num_aim == len(data_tensors[0])
            _select = (torch.arange(total_num_aim) % self.args.world_size) == self.args.local_rank  # bool tensor
            example_ids = np.array(example_ids)[_select].tolist()
            all_label = all_label[_select]
            data_tensors = [t[_select] for t in data_tensors]
            input_tensors = (example_ids, all_label, data_tensors, [])
        example_ids = input_tensors[0]
        print('local_rank', self.args.local_rank, 'len(example_ids)', len(example_ids), file=sys.stderr)
        return input_tensors

    def load_sparse_adj_data_with_contextnode(self, adj_pk_path, max_node_num, concepts_by_sents_list, mode='eval'):
        """Construct input tensors for the GNN component of the model."""
        print("Loading sparse adj data...")
        cache_path = adj_pk_path + "-nodenum{}".format(max_node_num) + ("-cntsall" if self.cxt_node_connects_all else "") + '_prompt_design_' + str(self.args.prompt_design) + '.loaded_cache'
        # use_cache = self.args.dump_graph_cache
        use_cache = self.args.load_graph_cache

        if use_cache and not os.path.exists(cache_path):
            use_cache = False

        if use_cache:
            print(f'Loading cache {cache_path}')
            print(f'Loading cache {cache_path}', file=sys.stderr)
            # with open(cache_path, 'rb') as f:
            #     adj_lengths_ori, concept_ids, node_type_ids, node_scores, adj_lengths, edge_index, edge_type, half_n_rel, special_nodes_mask = utils.CPU_Unpickler(f).load()
            loaded_data = []
            with open(cache_path, "rb") as in_file:
                try:
                    while True:
                        obj = pickle.load(in_file)
                        if type(obj) == dict:
                            assert len(obj) == 1
                            key = list(obj.keys())[0]
                            loaded_data.append(obj[key])
                        elif type(obj) == list:
                            loaded_data.extend(obj)
                        else:
                            raise TypeError("Invalid type for obj.")
                except EOFError:
                    pass
            adj_lengths_ori, concept_ids, node_type_ids, node_scores, adj_lengths, edge_index, edge_type, half_n_rel, special_nodes_mask = loaded_data
            self.final_num_relation = half_n_rel
            print(f'Loaded cache {cache_path}', file=sys.stderr)

            ori_adj_mean = adj_lengths_ori.float().mean().item()
            ori_adj_sigma = np.sqrt(((adj_lengths_ori.float() - ori_adj_mean)**2).mean().item())
            print('| ori_adj_len: mu {:.2f} sigma {:.2f} | adj_len: {:.2f} |'.format(ori_adj_mean, ori_adj_sigma, adj_lengths.float().mean().item()) +
                  ' prune_rate： {:.2f} |'.format((adj_lengths_ori > adj_lengths).float().mean().item()) +
                  ' qc_num: {:.2f} | ac_num: {:.2f} |'.format((node_type_ids == 0).float().sum(1).mean().item(),
                                                              (node_type_ids == 1).float().sum(1).mean().item()))

            edge_index = list(map(list, zip(*(iter(edge_index),) * self.num_choice)))  # list of size (n_questions, n_choices), where each entry is tensor[2, E] #this operation corresponds to .view(n_questions, n_choices)
            edge_type = list(map(list, zip(*(iter(edge_type),) * self.num_choice)))  # list of size (n_questions, n_choices), where each entry is tensor[E, ]

            adj_lengths_ori, concept_ids, node_type_ids, node_scores, adj_lengths, special_nodes_mask = [x.view(-1, self.num_choice, *x.size()[1:]) for x in (adj_lengths_ori, concept_ids, node_type_ids, node_scores, adj_lengths, special_nodes_mask)]
            # concept_ids: (n_questions, num_choice, max_node_num)
            # node_type_ids: (n_questions, num_choice, max_node_num)
            # node_scores: (n_questions, num_choice, max_node_num)
            # adj_lengths: (n_questions,　num_choice)

            if mode == 'train' and self.args.local_rank != -1:
                assert len(adj_lengths_ori) == len(concept_ids) == len(node_type_ids) == len(node_scores) == len(adj_lengths) == len(edge_index) == len(edge_type) == len(special_nodes_mask)  # they equal to n_questions * num_choice
                _ts = [adj_lengths_ori, concept_ids, node_type_ids, node_scores, adj_lengths, special_nodes_mask]
                total_num = len(edge_index)
                rem = total_num % self.args.world_size
                if rem != 0:
                    edge_index = edge_index + edge_index[:self.args.world_size - rem]
                    edge_type = edge_type + edge_type[:self.args.world_size - rem]
                    _ts = [torch.cat([t, t[:self.args.world_size - rem]], dim=0) for t in _ts]
                    total_num_aim = total_num + self.args.world_size - rem
                else:
                    total_num_aim = total_num
                assert total_num_aim % self.args.world_size == 0
                assert total_num_aim == len(_ts[0]) == len(edge_index)
                _select = (torch.arange(total_num_aim) % self.args.world_size) == self.args.local_rank  # bool tensor
                edge_index = [e for e, TF in zip(edge_index, _select) if TF]
                edge_type = [e for e, TF in zip(edge_type, _select) if TF]
                _ts = [t[_select] for t in _ts]
                adj_lengths_ori, concept_ids, node_type_ids, node_scores, adj_lengths, special_nodes_mask = _ts
                assert len(adj_lengths_ori) == len(concept_ids) == len(node_type_ids) == len(node_scores) == len(adj_lengths) == len(edge_index) == len(edge_type) == len(special_nodes_mask)
            print('local_rank', self.args.local_rank, 'len(edge_index)', len(edge_index), file=sys.stderr)
        else:
            # Set special nodes and links
            context_node = 0
            n_special_nodes = 1
            cxt2qlinked_rel = 0
            cxt2alinked_rel = 1
            half_n_rel = len(self.id2relation) + 2
            if self.cxt_node_connects_all:
                cxt2other_rel = half_n_rel
                half_n_rel += 1

            adj_concept_pairs = []
            print(f'Loading {adj_pk_path}...', file=sys.stderr)
            with open(adj_pk_path, "rb") as in_file:
                try:
                    while True:
                        ex = pickle.load(in_file)
                        if type(ex) == dict:
                            adj_concept_pairs.append(ex)
                        elif type(ex) == tuple:
                            adj_concept_pairs.append(ex)
                        elif type(ex) == list:
                            assert len(ex) > 10
                            adj_concept_pairs.extend(ex)
                        else:
                            raise TypeError("Invalid type for ex.")
                except EOFError:
                    pass
            print(f'Loaded {adj_pk_path}...', file=sys.stderr)

            # print('adj_concept_pairs: ', len(adj_concept_pairs))
            # print('adj_concept_pairs[0]: ', adj_concept_pairs[0])
            # print('adj_concept_pairs[1]: ', adj_concept_pairs[1])

            n_samples = len(adj_concept_pairs)  # this is actually n_questions x n_choices
            assert n_samples % self.num_choice == 0
            # print('self.num_choice in load_sparse_adj_data_with_contextnode: ', self.num_choice)
            n_questions = n_samples // self.num_choice
            # print('n_questions: ', n_questions)
            if mode == 'train' and self.args.local_rank != -1:
                rem = n_questions % self.args.world_size
                if rem != 0:
                    adj_concept_pairs = adj_concept_pairs + adj_concept_pairs[: (self.args.world_size - rem) * self.num_choice]
                    n_questions_aim = n_questions + self.args.world_size - rem
                else:
                    n_questions_aim = n_questions
                assert n_questions_aim % self.args.world_size == 0
                n_samples = n_questions_aim // self.args.world_size * self.num_choice

            edge_index, edge_type = [], []
            adj_lengths = torch.zeros((n_samples,), dtype=torch.long)
            concept_ids = torch.full((n_samples, max_node_num), 1, dtype=torch.long)
            node_type_ids = torch.full((n_samples, max_node_num), 2, dtype=torch.long)  # default 2: "other node"
            node_scores = torch.zeros((n_samples, max_node_num, 1), dtype=torch.float)
            special_nodes_mask = torch.zeros(n_samples, max_node_num, dtype=torch.bool)

            adj_lengths_ori = adj_lengths.clone()
            if not concepts_by_sents_list:
                concepts_by_sents_list = itertools.repeat(None)
            idx = -1
            for _idx, (_data, cpts_by_sents) in tqdm(enumerate(zip(adj_concept_pairs, concepts_by_sents_list)), total=n_questions * self.num_choice, desc='loading adj matrices'):
                if self.debug and _idx >= self.debug_sample_size * self.num_choice:
                    break

                if mode == 'train' and self.args.local_rank != -1:
                    qidx = _idx // self.num_choice
                    if qidx % self.args.world_size != self.args.local_rank:
                        continue
                idx += 1

                if isinstance(_data, dict):
                    adj, concepts, qm, am, cid2score = _data['adj'], _data['concepts'], _data['qmask'], _data['amask'], _data['cid2score']
                else:
                    adj, concepts, qm, am = _data
                    cid2score = None
                # adj: e.g. <4233x249 (n_nodes*half_n_rels x n_nodes) sparse matrix of type '<class 'numpy.bool'>' with 2905 stored elements in COOrdinate format>
                # concepts: np.array(num_nodes, ), where entry is concept id
                # qm: np.array(num_nodes, ), where entry is True/False
                # am: np.array(num_nodes, ), where entry is True/False
                assert len(concepts) == len(set(concepts))
                qam = qm | am
                # sanity check: should be T,..,T,F,F,..F
                if len(concepts) == 0:
                    # print ("KG with 0 node", file=sys.stderr)
                    pass
                else:
                    assert qam[0] == True
                F_start = False
                for TF in qam:
                    if TF == False:
                        F_start = True
                    else:
                        assert F_start == False

                assert n_special_nodes <= max_node_num
                special_nodes_mask[idx, :n_special_nodes] = 1

                if self.args.kg_only_use_qa_nodes:
                    actual_max_node_num = torch.tensor(qam).long().sum().item()
                else:
                    actual_max_node_num = max_node_num
                num_concept = min(len(concepts) + n_special_nodes, actual_max_node_num)  # this is the final number of nodes including contextnode but excluding PAD
                adj_lengths_ori[idx] = len(concepts)
                adj_lengths[idx] = num_concept

                # Prepare nodes
                concepts = concepts[:num_concept - n_special_nodes]
                concept_ids[idx, n_special_nodes:num_concept] = torch.tensor(concepts + 1)  # To accomodate contextnode, original concept_ids incremented by 1
                concept_ids[idx, 0] = context_node  # this is the "concept_id" for contextnode

                # Prepare node scores
                if cid2score is not None:
                    if -1 not in cid2score:
                        cid2score[-1] = 0
                    for _j_ in range(num_concept):
                        _cid = int(concept_ids[idx, _j_]) - 1  # Now context node is -1
                        node_scores[idx, _j_, 0] = torch.tensor(cid2score[_cid])

                # Prepare node types
                node_type_ids[idx, 0] = 3  # context node
                node_type_ids[idx, 1:n_special_nodes] = 4  # sent nodes
                node_type_ids[idx, n_special_nodes:num_concept][torch.tensor(qm, dtype=torch.bool)[:num_concept - n_special_nodes]] = 0
                node_type_ids[idx, n_special_nodes:num_concept][torch.tensor(am, dtype=torch.bool)[:num_concept - n_special_nodes]] = 1

                # Load adj
                ij = torch.tensor(adj.row, dtype=torch.int64)  # (num_matrix_entries, ), where each entry is coordinate
                k = torch.tensor(adj.col, dtype=torch.int64)  # (num_matrix_entries, ), where each entry is coordinate
                n_node = adj.shape[1]
                if n_node > 0:
                    assert len(self.id2relation) == adj.shape[0] // n_node
                    i, j = ij // n_node, ij % n_node
                else:
                    i, j = ij, ij

                # Prepare edges
                i += 2
                j += 1
                k += 1  # **** increment coordinate by 1, rel_id by 2 ****
                extra_i, extra_j, extra_k = [], [], []
                for _coord, q_tf in enumerate(qm):
                    _new_coord = _coord + n_special_nodes
                    if _new_coord > num_concept:
                        break
                    if q_tf:
                        extra_i.append(cxt2qlinked_rel)  # rel from contextnode to question concept
                        extra_j.append(0)  # contextnode coordinate
                        extra_k.append(_new_coord)  # question concept coordinate
                    elif self.cxt_node_connects_all:
                        extra_i.append(cxt2other_rel)  # rel from contextnode to other concept
                        extra_j.append(0)  # contextnode coordinate
                        extra_k.append(_new_coord)  # other concept coordinate
                for _coord, a_tf in enumerate(am):
                    _new_coord = _coord + n_special_nodes
                    if _new_coord > num_concept:
                        break
                    if a_tf:
                        extra_i.append(cxt2alinked_rel)  # rel from contextnode to answer concept
                        extra_j.append(0)  # contextnode coordinate
                        extra_k.append(_new_coord)  # answer concept coordinate
                    elif self.cxt_node_connects_all:
                        extra_i.append(cxt2other_rel)  # rel from contextnode to other concept
                        extra_j.append(0)  # contextnode coordinate
                        extra_k.append(_new_coord)  # other concept coordinate

                # half_n_rel += 2 #should be 19 now
                if len(extra_i) > 0:
                    i = torch.cat([i, torch.tensor(extra_i)], dim=0)
                    j = torch.cat([j, torch.tensor(extra_j)], dim=0)
                    k = torch.cat([k, torch.tensor(extra_k)], dim=0)

                if self.args.max_num_relation > 0:
                    _keep = (i < self.args.max_num_relation).bool()
                    assert _keep.dim() == 1 and _keep.size(0) == i.size(0) == j.size(0) == k.size(0)
                    i = i[_keep]
                    j = j[_keep]
                    k = k[_keep]
                    half_n_rel = min(half_n_rel, self.args.max_num_relation)
                self.final_num_relation = half_n_rel
                ########################

                mask = (j < actual_max_node_num) & (k < actual_max_node_num)
                i, j, k = i[mask], j[mask], k[mask]
                i, j, k = torch.cat((i, i + half_n_rel), 0), torch.cat((j, k), 0), torch.cat((k, j), 0)  # add inverse relations
                edge_index.append(torch.stack([j, k], dim=0))  # each entry is [2, E]
                edge_type.append(i)  # each entry is [E, ]

            if not self.debug and self.args.dump_graph_cache:
                print('saving cache...', file=sys.stderr)
                with open(cache_path, 'wb') as f:
                    # pickle.dump([adj_lengths_ori, concept_ids, node_type_ids, node_scores, adj_lengths, edge_index, edge_type, half_n_rel, special_nodes_mask], f)
                    for _i_, obj in enumerate(tqdm([adj_lengths_ori, concept_ids, node_type_ids, node_scores, adj_lengths, edge_index, edge_type, half_n_rel, special_nodes_mask])):
                        pickle.dump({f'obj{_i_}': obj}, f, protocol=4)
                print('saved cache.', file=sys.stderr)

            assert n_samples == idx+1
            print('local_rank', self.args.local_rank, 'graph loading final idx+1', idx+1, file=sys.stderr)
            del adj_concept_pairs

            ori_adj_mean = adj_lengths_ori.float().mean().item()
            ori_adj_sigma = np.sqrt(((adj_lengths_ori.float() - ori_adj_mean)**2).mean().item())
            print('| ori_adj_len: mu {:.2f} sigma {:.2f} | adj_len: {:.2f} |'.format(ori_adj_mean, ori_adj_sigma, adj_lengths.float().mean().item()) +
                  ' prune_rate： {:.2f} |'.format((adj_lengths_ori > adj_lengths).float().mean().item()) +
                  ' qc_num: {:.2f} | ac_num: {:.2f} |'.format((node_type_ids == 0).float().sum(1).mean().item(),
                                                              (node_type_ids == 1).float().sum(1).mean().item()))

            edge_index = list(map(list, zip(*(iter(edge_index),) * self.num_choice)))  # list of size (n_questions, n_choices), where each entry is tensor[2, E] #this operation corresponds to .view(n_questions, n_choices)
            edge_type = list(map(list, zip(*(iter(edge_type),) * self.num_choice)))  # list of size (n_questions, n_choices), where each entry is tensor[E, ]

            adj_lengths_ori, concept_ids, node_type_ids, node_scores, adj_lengths, special_nodes_mask = [x.view(-1, self.num_choice, *x.size()[1:]) for x in (adj_lengths_ori, concept_ids, node_type_ids, node_scores, adj_lengths, special_nodes_mask)]
            # concept_ids: (n_questions, num_choice, max_node_num)
            # node_type_ids: (n_questions, num_choice, max_node_num)
            # node_scores: (n_questions, num_choice, max_node_num)
            # adj_lengths: (n_questions,　num_choice)
        return concept_ids, node_type_ids, node_scores, adj_lengths, special_nodes_mask, (edge_index, edge_type)  # , half_n_rel * 2 + 1


######################### GPT loader utils #########################
def load_gpt_input_tensors(statement_jsonl_path, max_seq_length):
    def _truncate_seq_pair(tokens_a, tokens_b, max_length):
        """Truncates a sequence pair in place to the maximum length."""
        while True:
            total_length = len(tokens_a) + len(tokens_b)
            if total_length <= max_length:
                break
            if len(tokens_a) > len(tokens_b):
                tokens_a.pop()
            else:
                tokens_b.pop()

    def load_qa_dataset(dataset_path):
        """ Output a list of tuples(story, 1st continuation, 2nd continuation, label) """
        with open(dataset_path, "r", encoding="utf-8") as fin:
            output = []
            for line in fin:
                input_json = json.loads(line)
                label = ord(input_json.get("answerKey", "A")) - ord("A")
                output.append((input_json['id'], input_json["question"]["stem"], *[ending["text"] for ending in input_json["question"]["choices"]], label))
        return output

    def pre_process_datasets(encoded_datasets, num_choices, max_seq_length, start_token, delimiter_token, clf_token):
        """ Pre-process datasets containing lists of tuples(story, 1st continuation, 2nd continuation, label)

            To Transformer inputs of shape (n_batch, n_alternative, length) comprising for each batch, continuation:
            input_ids[batch, alternative, :] = [start_token] + story[:cap_length] + [delimiter_token] + cont1[:cap_length] + [clf_token]
        """
        tensor_datasets = []
        for dataset in encoded_datasets:
            n_batch = len(dataset)
            input_ids = np.zeros((n_batch, num_choices, max_seq_length), dtype=np.int64)
            mc_token_ids = np.zeros((n_batch, num_choices), dtype=np.int64)
            lm_labels = np.full((n_batch, num_choices, max_seq_length), fill_value=-1, dtype=np.int64)
            mc_labels = np.zeros((n_batch,), dtype=np.int64)
            for i, data, in enumerate(dataset):
                q, mc_label = data[0], data[-1]
                choices = data[1:-1]
                for j in range(len(choices)):
                    _truncate_seq_pair(q, choices[j], max_seq_length - 3)
                    qa = [start_token] + q + [delimiter_token] + choices[j] + [clf_token]
                    input_ids[i, j, :len(qa)] = qa
                    mc_token_ids[i, j] = len(qa) - 1
                    lm_labels[i, j, :len(qa) - 1] = qa[1:]
                mc_labels[i] = mc_label
            all_inputs = (input_ids, mc_token_ids, lm_labels, mc_labels)
            tensor_datasets.append(tuple(torch.tensor(t) for t in all_inputs))
        return tensor_datasets

    def tokenize_and_encode(tokenizer, obj):
        """ Tokenize and encode a nested object """
        if isinstance(obj, str):
            return tokenizer.convert_tokens_to_ids(tokenizer.tokenize(obj))
        elif isinstance(obj, int):
            return obj
        else:
            return list(tokenize_and_encode(tokenizer, o) for o in obj)

    tokenizer = OpenAIGPTTokenizer.from_pretrained('openai-gpt')
    tokenizer.add_tokens(GPT_SPECIAL_TOKENS)
    special_tokens_ids = tokenizer.convert_tokens_to_ids(GPT_SPECIAL_TOKENS)

    dataset = load_qa_dataset(statement_jsonl_path)
    examples_ids = [data[0] for data in dataset]
    dataset = [data[1:] for data in dataset]  # discard example ids
    num_choices = len(dataset[0]) - 2

    encoded_dataset = tokenize_and_encode(tokenizer, dataset)

    (input_ids, mc_token_ids, lm_labels, mc_labels), = pre_process_datasets([encoded_dataset], num_choices, max_seq_length, *special_tokens_ids)
    return examples_ids, mc_labels, input_ids, mc_token_ids, lm_labels


######################### BERT/XLNet/Roberta loader utils #########################
class InputExample(object):

    def __init__(self, example_id, question, contexts, endings, label=None):
        self.example_id = example_id
        self.question = question
        self.contexts = contexts
        self.endings = endings
        self.label = label


class InputFeatures(object):

    def __init__(self, example_id, choices_features, label):
        self.example_id = example_id
        self.choices_features = [
            {
                'input_ids': input_ids,
                'input_mask': input_mask,
                'segment_ids': segment_ids,
                'output_mask': output_mask,
            }
            for input_ids, input_mask, segment_ids, output_mask in choices_features
        ]
        self.label = label


def read_examples(input_file):
    with open(input_file, "r", encoding="utf-8") as f:
        examples = []
        for line in tqdm(f.readlines()):
            json_dic = json.loads(line)
            label = ord(json_dic["answerKey"]) - ord("A") if 'answerKey' in json_dic else -100
            contexts = json_dic["question"]["stem"]
            if "para" in json_dic:
                contexts = json_dic["para"] + " " + contexts
            if "fact1" in json_dic:
                contexts = json_dic["fact1"] + " " + contexts
            engings_list = [ending["text"] for ending in json_dic["question"]["choices"]]
            examples.append(
                InputExample(
                    example_id=json_dic["id"],
                    contexts=contexts,
                    # contexts=[contexts] * len(json_dic["question"]["choices"]),
                    question="",
                    endings=engings_list,
                    label=engings_list[label] if label != -100 else -100  # return real string instead of integer.
                    # label=label
                ))
    return examples


def llm_simple_convert_examples_to_features(examples, label_list, max_seq_length, tokenizer, args, debug=False):
    """ 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]
            - True (XLNet/GPT pattern): A + [SEP] + B + [SEP] + [CLS]
        `cls_token_segment_id` define the segment id associated to the CLS token (0 for BERT, 2 for XLNet)
    """
    # label_map = {label: i for i, label in enumerate(label_list)}

    features = []
    concepts_by_sents_list = []
    for ex_index, example in tqdm(enumerate(examples), total=len(examples), desc="Converting examples to features"):
        if debug and ex_index >= debug_sample_size:
            break
        choices_features = []
        # print('example.contexts: ', example.contexts)
        # print('example.endings: ', example.endings)
        # print('example.label: ', example.label)

        # tokenize labels
        label = example.label
        # print('label: ', label)
        if label != -100:
            encoded_label = tokenizer(label, padding="max_length", truncation=True, max_length=max_seq_length, return_token_type_ids=True, return_special_tokens_mask=True)
            label_input_ids = encoded_label["input_ids"]
        else:
            label_input_ids = label
        # print('label_input_ids: ', label_input_ids)

        # tokenize input: it seems that the input should be tokenized in here instead of in training loop. Otherwise extra efforts are needed.

        # format: "Q+A1, ..., A5"
        # input_to_tokenizer = example.contexts + ' ' + ', '.join([ending for ending in example.endings])
        # format: "Question: Q\nOptions: A1, ..., A5"
        if not args.prompt_design:
            input_to_tokenizer = 'Question: '+example.contexts + '\n' + 'Options: ' + ', '.join([ending for ending in example.endings])
        elif args.prompt_design == 'prompt_A':
            # format: "Choose the best option to answer the question. Question: Q\nOptions: A1, ..., A5"
            input_to_tokenizer = 'Choose the best option to answer the question. Question: '+example.contexts + '\n' + 'Options: ' + ', '.join([ending for ending in example.endings])
        elif args.prompt_design == 'prompt_B':
            # format: "Which is the correct answer? Question: Q\nOptions: A1, ..., A5"
            input_to_tokenizer = 'Which is the correct answer? Question: '+example.contexts + '\n' + 'Options: ' + ', '.join([ending for ending in example.endings])
        elif args.prompt_design == 'prompt_C':
            # format: "Which option can answer this question? Question: Q\nOptions: A1, ..., A5"
            input_to_tokenizer = 'Which option can answer this question? Question: '+example.contexts + '\n' + 'Options: ' + ', '.join([ending for ending in example.endings])
        else:
            raise ValueError("Please indicate whether and which prompt_design to use.")

        input_to_tokenizer = ''.join(input_to_tokenizer)
        # print('input_to_tokenizer: ', input_to_tokenizer)

        encoded_input = tokenizer(input_to_tokenizer, padding="max_length", truncation=True, max_length=max_seq_length, return_token_type_ids=True, return_special_tokens_mask=True)
        input_ids = encoded_input["input_ids"]
        # print('input_ids: ', input_ids, len(input_ids))
        output_mask = encoded_input["special_tokens_mask"]
        input_mask = encoded_input["attention_mask"]
        segment_ids = encoded_input["token_type_ids"]
        choices_features.append((input_ids, input_mask, segment_ids, output_mask))

        features.append(InputFeatures(example_id=example.example_id, choices_features=choices_features, label=label_input_ids))

    return features, concepts_by_sents_list

# def simple_convert_examples_to_features(examples, label_list, max_seq_length, tokenizer, debug=False):
#     """ 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]
#             - True (XLNet/GPT pattern): A + [SEP] + B + [SEP] + [CLS]
#         `cls_token_segment_id` define the segment id associated to the CLS token (0 for BERT, 2 for XLNet)
#     """
#     label_map = {label: i for i, label in enumerate(label_list)}

#     features = []
#     concepts_by_sents_list = []
#     for ex_index, example in tqdm(enumerate(examples), total=len(examples), desc="Converting examples to features"):
#         if debug and ex_index >= debug_sample_size:
#             break
#         choices_features = []
#         # print('example.contexts: ', example.contexts)
#         # print('example.endings: ', example.endings)
#         # print('example.label: ', example.label)
#         for ending_idx, (context, ending) in enumerate(zip(example.contexts, example.endings)):
#             ans = example.question + " " + ending

#             # print('ending_idx: ', ending_idx)
#             # print('context: ', context)
#             # print('ans: ', ans)

#             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)
#             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('input_ids: ', len(input_ids))
#             # print('input_mask: ', len(input_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))
#         label = label_map.get(example.label, -100)
#         # print('label: ', label)
#         features.append(InputFeatures(example_id=example.example_id, choices_features=choices_features, label=label))
#         # print('features: ', features)

#     return features, concepts_by_sents_list


def load_bert_xlnet_roberta_input_tensors(statement_jsonl_path, max_seq_length, debug, tokenizer, debug_sample_size, args):

    def select_field(features, field):
        return [[choice[field] for choice in feature.choices_features] for feature in features]

    def convert_features_to_tensors(features):
        all_input_ids = torch.tensor(select_field(features, 'input_ids'), dtype=torch.long)
        all_input_mask = torch.tensor(select_field(features, 'input_mask'), dtype=torch.long)
        all_segment_ids = torch.tensor(select_field(features, 'segment_ids'), dtype=torch.long)
        all_output_mask = torch.tensor(select_field(features, 'output_mask'), dtype=torch.bool)
        all_label = torch.tensor([f.label for f in features], dtype=torch.long)
        return all_input_ids, all_input_mask, all_segment_ids, all_output_mask, all_label

    examples = read_examples(statement_jsonl_path)
    num_workers = 60
    if num_workers <= 1:
        features, concepts_by_sents_list = llm_simple_convert_examples_to_features(examples, list(range(len(examples[0].endings))), max_seq_length, tokenizer, args, debug)
        # features, concepts_by_sents_list = simple_convert_examples_to_features(examples, list(range(len(examples[0].endings))), max_seq_length, tokenizer, debug)
    else:
        from copy import deepcopy
        from multiprocessing import Pool
        c_size = len(examples)//num_workers + 1
        examples_list = [examples[i*c_size: (i+1)*c_size] for i in range(num_workers)]
        label_list_list = [list(range(len(examples[0].endings))) for i in range(num_workers)]
        max_seq_length_list = [max_seq_length for i in range(num_workers)]
        tokenizer_list = [deepcopy(tokenizer) for i in range(num_workers)]
        args_list = [args for i in range(num_workers)]
        with Pool(num_workers) as p:
            ress = list(p.starmap(llm_simple_convert_examples_to_features, zip(examples_list, label_list_list, max_seq_length_list, tokenizer_list, args_list)))
            # ress = list(p.starmap(simple_convert_examples_to_features, zip(examples_list, label_list_list, max_seq_length_list, tokenizer_list)))
        features, concepts_by_sents_list = [], []
        for res in ress:
            features += res[0]
            concepts_by_sents_list += res[1]

    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
