Castor/mp_cnn/dataset.py

from collections import defaultdict
from enum import Enum
import math
import os

import nltk
import numpy as np
import torch
from torch.autograd import Variable
import torch.utils.data as data

import preprocessing

nltk.download('stopwords', quiet=True)
from nltk.corpus import stopwords

# logging setup
import logging
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

ch = logging.StreamHandler()
ch.setLevel(logging.DEBUG)
formatter = logging.Formatter('%(levelname)s - %(message)s')
ch.setFormatter(formatter)
logger.addHandler(ch)


class DatasetType(Enum):
    TRAIN = 1
    TEST = 2
    DEV = 3


class MPCNNDatasetFactory(object):
    """
    Get the corresponding Dataset class for a particular dataset.
    """
    @staticmethod
    def get_dataset(dataset_name, word_vectors_file, batch_size, cuda, sample):
        extra_args = {'shuffle': True}
        dev_loader = None
        if sample:
            sample_indices = list(range(sample))
            subset_random_sampler = data.sampler.SubsetRandomSampler(sample_indices)
            extra_args['sampler'] = subset_random_sampler
            extra_args['shuffle'] = False
        if dataset_name == 'sick':
            train_loader = torch.utils.data.DataLoader(SICKDataset(DatasetType.TRAIN, cuda), batch_size=batch_size, **extra_args)
            test_loader = torch.utils.data.DataLoader(SICKDataset(DatasetType.TEST, cuda), batch_size=batch_size, **extra_args)
            dev_loader = torch.utils.data.DataLoader(SICKDataset(DatasetType.DEV, cuda), batch_size=batch_size, **extra_args)
        elif dataset_name == 'msrvid':
            train_loader = torch.utils.data.DataLoader(MSRVIDDataset(DatasetType.TRAIN, cuda), batch_size=batch_size, **extra_args)
            test_loader = torch.utils.data.DataLoader(MSRVIDDataset(DatasetType.TEST, cuda), batch_size=batch_size, **extra_args)
        else:
            raise ValueError('{} is not a valid dataset.'.format(dataset_name))

        word_index, embedding = preprocessing.get_glove_embedding(word_vectors_file, train_loader.dataset.dataset_root)
        logger.info('Finished loading GloVe embedding for vocab in data...')

        train_loader.dataset.initialize(word_index, embedding)
        test_loader.dataset.initialize(word_index, embedding)
        if dev_loader is not None:
            dev_loader.dataset.initialize(word_index, embedding)
        return train_loader, test_loader, dev_loader


class MPCNNDataset(data.Dataset):
    train_folder = 'train'
    test_folder = 'test'
    dev_folder = 'dev'
    # subclass will override fields below
    dataset_root = None
    num_classes = None

    def __init__(self, dataset_type, cuda):
        if not isinstance(dataset_type, DatasetType):
            raise ValueError('dataset_type ({}) must be of type DatasetType enum'.format(dataset_type))

        if dataset_type == DatasetType.TRAIN:
            subfolder = MPCNNDataset.train_folder
        elif dataset_type == DatasetType.TEST:
            subfolder = MPCNNDataset.test_folder
        else:
            subfolder = MPCNNDataset.dev_folder

        self.dataset_dir = os.path.join(self.dataset_root, subfolder)
        if not os.path.exists(self.dataset_dir):
            raise RuntimeError('{} does not exist'.format(self.dataset_dir))

        self.cuda = cuda
        self.max_length = -10000
        self.unk = torch.Tensor(300)
        self.unk.normal_(0, 0.01)

    def initialize(self, word_index, embedding):
        """
        Convert sentences into sentence embeddings.
        """
        sent_a = self._load(self.dataset_dir, 'a.txt')
        sent_b = self._load(self.dataset_dir, 'b.txt')
        word_to_doc_cnt = defaultdict(int)

        # obtain max sentence length to use as dimension for padding to support batching
        sent_a_tokens, sent_b_tokens = [], []
        for i in range(len(sent_a)):
            sa_tokens = sent_a[i].split(' ')
            sb_tokens = sent_b[i].split(' ')
            self.max_length = max(self.max_length, len(sa_tokens), len(sb_tokens))
            sent_a_tokens.append(sa_tokens)
            sent_b_tokens.append(sb_tokens)

            unique_tokens = set(sa_tokens) | set(sb_tokens)
            for t in unique_tokens:
                word_to_doc_cnt[t] += 1

        self.sentences = []
        stoplist = set(stopwords.words('english'))
        num_docs = len(word_to_doc_cnt)
        for i in range(len(sent_a)):
            sent_pair = {}
            sent_pair['a'] = self._get_sentence_embeddings(sent_a_tokens[i], word_index, embedding)
            sent_pair['b'] = self._get_sentence_embeddings(sent_b_tokens[i], word_index, embedding)

            tokens_a_set, tokens_b_set = set(sent_a_tokens[i]), set(sent_b_tokens[i])
            intersect = tokens_a_set & tokens_b_set
            overlap = len(intersect) / (len(tokens_a_set) + len(tokens_b_set))
            idf_intersect = sum(np.math.log(num_docs / word_to_doc_cnt[w]) for w in intersect)
            idf_weighted_overlap = idf_intersect / (len(tokens_a_set) + len(tokens_b_set))

            tokens_a_set_no_stop = set(w for w in sent_a_tokens[i] if w not in stoplist)
            tokens_b_set_no_stop = set(w for w in sent_b_tokens[i] if w not in stoplist)
            intersect_no_stop = tokens_a_set_no_stop & tokens_b_set_no_stop
            overlap_no_stop = len(intersect_no_stop) / (len(tokens_a_set_no_stop) + len(tokens_b_set_no_stop))
            idf_intersect_no_stop = sum(np.math.log(num_docs / word_to_doc_cnt[w]) for w in intersect_no_stop)
            idf_weighted_overlap_no_stop = idf_intersect_no_stop / (len(tokens_a_set_no_stop) + len(tokens_b_set_no_stop))
            ext_feats = torch.Tensor([overlap, idf_weighted_overlap, overlap_no_stop, idf_weighted_overlap_no_stop])
            ext_feats = ext_feats.cuda() if self.cuda else ext_feats
            sent_pair['ext_feats'] = ext_feats

            self.sentences.append(sent_pair)

        self.labels = self._load(self.dataset_dir, 'sim.txt', float)

    def _load(self, dataset_dir, fname, type_converter=str):
        data = []
        with open(os.path.join(dataset_dir, fname), 'r') as f:
            for line in f:
                stripped_line = line.rstrip('.\n')
                item = type_converter(stripped_line)
                data.append(item)
        return data

    def _get_sentence_embeddings(self, tokens, word_index, embedding):
        sentence_embedding = torch.zeros(300, self.max_length)
        sentence_embedding[:, :len(tokens)].normal_(0, 1)
        found_pos, found_emb_idx = [], []
        for i, token in enumerate(tokens):
            if token in word_index:
                found_pos.append(i)
                found_emb_idx.append(word_index[token])
            else:
                sentence_embedding[:, i] = self.unk

        found_word_vecs = embedding(Variable(torch.LongTensor(found_emb_idx)))
        for i, v in enumerate(found_pos):
            sentence_embedding[:, v] = found_word_vecs[i].data
        return sentence_embedding.cuda() if self.cuda else sentence_embedding

    def __getitem__(self, idx):
        return self.sentences[idx], self.labels[idx]

    def __len__(self):
        return len(self.labels)


class SICKDataset(MPCNNDataset):

    dataset_root = os.path.join(os.pardir, os.pardir, 'data', 'sick')
    num_classes = 5

    def __init__(self, dataset_type, cuda):
        super(SICKDataset, self).__init__(dataset_type, cuda)

    def initialize(self, word_index, embedding):
        super(SICKDataset, self).initialize(word_index, embedding)
        new_labels = torch.zeros(self.__len__(), self.num_classes)
        for i, sim in enumerate(self.labels):
            ceil, floor = math.ceil(sim), math.floor(sim)
            if ceil == floor:
                new_labels[i][floor - 1] = 1
            else:
                new_labels[i][floor - 1] = ceil - sim
                new_labels[i][ceil - 1] = sim - floor

        self.labels = new_labels.cuda() if self.cuda else new_labels


class MSRVIDDataset(MPCNNDataset):

    dataset_root = os.path.join(os.pardir, os.pardir, 'data', 'msrvid')
    num_classes = 6

    def __init__(self, dataset_type, cuda):
        super(MSRVIDDataset, self).__init__(dataset_type, cuda)

    def initialize(self, word_index, embedding):
        super(MSRVIDDataset, self).initialize(word_index, embedding)
        new_labels = torch.zeros(self.__len__(), self.num_classes)
        for i, sim in enumerate(self.labels):
            ceil, floor = math.ceil(sim), math.floor(sim)
            if ceil == floor:
                new_labels[i][floor] = 1
            else:
                new_labels[i][floor] = ceil - sim
                new_labels[i][ceil] = sim - floor

        self.labels = new_labels.cuda() if self.cuda else new_labels