dataset.py

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

import torch
import networkx as nx
import dgl
from typing import (
    Any,
    Callable,
    Dict,
    Generator,
    Iterable,
    Iterator,
    List,
    Optional,
    Tuple,
    TypeVar,
    Union,
    Sequence,
    Set,
)

from IPython import embed
from tqdm import tqdm
from allennlp.common.util import pad_sequence_to_length
from allennlp.nn.util import get_mask_from_sequence_lengths

class Dataset(object):
    def __init__(self, batch_size, dataset):
        super(Dataset, self).__init__()

        self.batch_size = batch_size
        self.y_label = {
            'NULL': 0,
            'null': 0,
            'FALLING_ACTION': 1,
            'PRECONDITION': 1,
            'Coref': 1,
            'Cause-Effect': 1,
            'Cause-Effect1': 1,
            'Cause-Effect2': 1,
            'Cause': 1
        }

        self.construct_index(dataset)

    def construct_index(self, dataset):
        self.dataset = dataset
        self.index_length = len(dataset)
        self.shuffle_list = list(range(0, self.index_length))


    def shuffle(self):
        random.shuffle(self.shuffle_list)

    def get_tqdm(self, device, shuffle=True):
        return tqdm(self.reader(device, shuffle), mininterval=2, total=self.index_length // self.batch_size, leave=False, file=sys.stdout, ncols=80)

    def reader(self, device, shuffle):
        cur_idx = 0
        while cur_idx < self.index_length:
            end_index = min(cur_idx + self.batch_size, self.index_length)
            batch = [self.dataset[self.shuffle_list[index]] for index in range(cur_idx, end_index)]
            cur_idx = end_index
            yield self.batchify(batch, device)
        if shuffle:
            self.shuffle()

    def batchify(self, batch, device):
        examples = list()
        sentence_len_s = [len(tup[1]) for tup in batch]
        sentence_len_t = [len(tup[2]) for tup in batch]

        max_sentence_len_s = max(sentence_len_s)
        max_sentence_len_t = max(sentence_len_t)

        event1_lens = [len(tup[2]) for tup in batch]
        event2_lens = [len(tup[3]) for tup in batch]

        sentences_s, sentences_t, event1, event2, data_y = list(), list(), list(), list(), list()
        for data in batch:

            sentences_s.append(data[1])
            sentences_t.append(data[2])
            event1.append(data[3])
            event2.append(data[4])
            y = self.y_label[data[5]] if data[5] in self.y_label else 0
            data_y.append(y)
            examples.append(data)

        sentences_s = list(map(lambda x: pad_sequence_to_length(x, max_sentence_len_s), sentences_s))
        sentences_t = list(map(lambda x: pad_sequence_to_length(x, max_sentence_len_t), sentences_t))

        event1 = list(map(lambda x: pad_sequence_to_length(x, 5), event1))
        event2 = list(map(lambda x: pad_sequence_to_length(x, 5), event2))

        mask_sentences_s = get_mask_from_sequence_lengths(torch.LongTensor(sentence_len_s), max_sentence_len_s)
        mask_sentences_t = get_mask_from_sequence_lengths(torch.LongTensor(sentence_len_t), max_sentence_len_t)

        mask_even1 = get_mask_from_sequence_lengths(torch.LongTensor(event1_lens), 5)
        mask_even2 = get_mask_from_sequence_lengths(torch.LongTensor(event2_lens), 5)

        return [torch.LongTensor(sentences_s).to(device), mask_sentences_s.to(device),
                torch.LongTensor(sentences_t).to(device), mask_sentences_t.to(device),
                torch.LongTensor(event1).to(device), mask_even1.to(device),
                torch.LongTensor(event2).to(device), mask_even2.to(device),
                torch.LongTensor(data_y).to(device), examples]


def get_graph(graph_ngx_file):
        start_time = timeit.default_timer()
        nxgs = []
        dgs = []
        graph_ngx_file =  "/u/wusifan/Event_causality_identification/causaltb/preprocess/"  + graph_ngx_file #"/u/wusifan/Event_causality_identification/cedar/"  + 
        save_file = graph_ngx_file + ".dgl.pk"

        if os.path.exists(save_file):
            df = open(save_file, "rb")
            dgs = pickle.load(df)
            return dgs

        print("loading paths from %s" % graph_ngx_file)
        with open(graph_ngx_file, 'r') as fr:
            for line in fr.readlines():
                line = line.strip()
                nxgs.append(line)
        print('\t Done! Time: ', "{0:.2f} sec".format(float(timeit.default_timer() - start_time)))

        save_file = graph_ngx_file + ".dgl.pk"

        if os.path.exists(save_file):
            import gc
            print("loading pickle for the dgl", save_file)
            start_time = timeit.default_timer()
            with open(save_file, 'rb') as handle:
                gc.disable()
                dgs = pickle.load(handle)
                gc.enable()
            print("finished loading in %.3f secs" % (float(timeit.default_timer() - start_time)))
        else:
            print("len(nxgs):", len(nxgs))
            for index, nxg_str in tqdm(enumerate(nxgs), total=len(nxgs)):
                
                nxg = nx.node_link_graph(json.loads(nxg_str))
                '''
                #dg = dgl.DGLGraph(multigraph=True)
                # dg.from_networkx(nxg, edge_attrs=["rel"])
                # dg.from_networkx(nxg)
                '''
                #dg = dgl.from_networkx(nxg, edge_attrs=["rel"])
                dg = dgl.from_networkx(nxg)
                cids = [nxg.nodes[n_id]['cid']+1 for n_id in range(len(dg))] # -1 --> 0 and 0 stands for a palceholder concept
                
                # rel_types = [nxg.edges[u, v, r]["rel"] + 1 for u, v, r in nxg.edges]  # 0 is used for

                # print(line)
                # node_types = [mapping_type[nxg.nodes[n_id]['type']] for n_id in range(len(dg))]
                # edge_weights = [nxg.edges[u, v].get("weight", 0.0) for u, v in nxg.edges]  # -1 is used for the unk edges
                # dg.edata.update({'weights': torch.FloatTensor(edge_weights)})

                # dg.edata.update({'rel_types': torch.LongTensor(rel_types)})

                dg.ndata.update({'cncpt_ids': torch.LongTensor(cids)})
                
                dgs.append(dg)
            
            print("saving pickle for the dgl", save_file)
            with open(save_file, 'wb') as handle:
                pickle.dump(dgs[0], handle, protocol=pickle.HIGHEST_PROTOCOL)
        return dgs[0]


class DatasetwithGraph(Dataset):
    def __init__(self, batch_size, dataset):
        super(DatasetwithGraph, self).__init__(batch_size, dataset)


    def batchify(self, batch, device):
        examples = list()
        sentence_len_s = [len(tup[1]) for tup in batch]
        sentence_len_t = [len(tup[2]) for tup in batch]

        max_sentence_len_s = max(sentence_len_s)
        max_sentence_len_t = max(sentence_len_t)

        event1_lens = [len(tup[2]) for tup in batch]
        event2_lens = [len(tup[3]) for tup in batch]

        sentences_s, sentences_t, event1, event2, data_y, graph = list(), list(), list(), list(), list(), list()
        for data in batch:

            sentences_s.append(data[1])
            sentences_t.append(data[2])
            event1.append(data[3])
            event2.append(data[4])
            graph.append(get_graph(data[-1]))
            y = self.y_label[data[5]] if data[5] in self.y_label else 0
            data_y.append(y)
            examples.append(data)
        
        
        batch_graph = dgl.batch(graph)
        sentences_s = list(map(lambda x: pad_sequence_to_length(x, max_sentence_len_s), sentences_s))
        sentences_t = list(map(lambda x: pad_sequence_to_length(x, max_sentence_len_t), sentences_t))

        event1 = list(map(lambda x: pad_sequence_to_length(x, 5), event1))
        event2 = list(map(lambda x: pad_sequence_to_length(x, 5), event2))

        mask_sentences_s = get_mask_from_sequence_lengths(torch.LongTensor(sentence_len_s), max_sentence_len_s)
        mask_sentences_t = get_mask_from_sequence_lengths(torch.LongTensor(sentence_len_t), max_sentence_len_t)

        mask_even1 = get_mask_from_sequence_lengths(torch.LongTensor(event1_lens), 5)
        mask_even2 = get_mask_from_sequence_lengths(torch.LongTensor(event2_lens), 5)

        return [torch.LongTensor(sentences_s).to(device), mask_sentences_s.to(device),
                torch.LongTensor(sentences_t).to(device), mask_sentences_t.to(device),
                torch.LongTensor(event1).to(device), mask_even1.to(device),
                torch.LongTensor(event2).to(device), mask_even2.to(device),
                torch.LongTensor(data_y).to(device), batch_graph, examples]


if __name__ == '__main__':
    with open('data.pickle', 'rb') as f:
        # The protocol version used is detected automatically, so we do not
        # have to specify it.
        data = pickle.load(f)
    dataset = Dataset(10, data[:20])
    for batch in dataset.reader('cpu', True):
        sentences_s, mask_s, sentences_t, mask_t, event1, event1_mask, event2, event2_mask, y = batch
        print(sentences_s[0])
        print(mask_s[0])
        print(event1[0])
        print(event2[0])
        break