Castor/mp_cnn/dataset.py

from enum import Enum
import math
import os

import torch
from torch.autograd import Variable
import torch.nn.functional as F
import torch.utils.data as data

import preprocessing

# 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

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

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

        self.sentences = []
        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)
            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])

        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