fastNLP/reproduction/CNN-sentence_classification/dataset.py

import codecs
import random
import re

import gensim
import numpy as np
from gensim import corpora
from torch.utils.data import Dataset


def clean_str(string):
    """
    Tokenization/string cleaning for all datasets except for SST.
    Original taken from https://github.com/yoonkim/CNN_sentence/blob/master/process_data.py
    """
    string = re.sub(r"[^A-Za-z0-9(),!?\'\`]", " ", string)
    string = re.sub(r"\'s", " \'s", string)
    string = re.sub(r"\'ve", " \'ve", string)
    string = re.sub(r"n\'t", " n\'t", string)
    string = re.sub(r"\'re", " \'re", string)
    string = re.sub(r"\'d", " \'d", string)
    string = re.sub(r"\'ll", " \'ll", string)
    string = re.sub(r",", " , ", string)
    string = re.sub(r"!", " ! ", string)
    string = re.sub(r"\(", " \( ", string)
    string = re.sub(r"\)", " \) ", string)
    string = re.sub(r"\?", " \? ", string)
    string = re.sub(r"\s{2,}", " ", string)
    return string.strip()


def pad_sentences(sentence, padding_word=" <PAD/>"):
    sequence_length = 64
    sent = sentence.split()
    padded_sentence = sentence + padding_word * (sequence_length - len(sent))
    return padded_sentence


# data loader
class MRDataset(Dataset):
    def __init__(self):

        # load positive and negative sentenses from files
        with codecs.open("./rt-polaritydata/rt-polarity.pos", encoding='ISO-8859-1') as f:
            positive_examples = list(f.readlines())
        with codecs.open("./rt-polaritydata/rt-polarity.neg", encoding='ISO-8859-1') as f:
            negative_examples = list(f.readlines())
        # s.strip: clear "\n"; clear_str; pad
        positive_examples = [pad_sentences(clean_str(s.strip())) for s in positive_examples]
        negative_examples = [pad_sentences(clean_str(s.strip())) for s in negative_examples]
        self.examples = positive_examples + negative_examples
        self.sentences_texts = [sample.split() for sample in self.examples]

        # word dictionary
        dictionary = corpora.Dictionary(self.sentences_texts)
        self.word2id_dict = dictionary.token2id  # transform to dict, like {"human":0, "a":1,...}

        # set lables: postive is 1; negative is 0
        positive_labels = [1 for _ in positive_examples]
        negative_labels = [0 for _ in negative_examples]
        self.lables = positive_labels + negative_labels
        examples_lables = list(zip(self.examples, self.lables))
        random.shuffle(examples_lables)
        self.MRDataset_frame = examples_lables

        # transform word to id
        self.MRDataset_wordid = \
            [(
                np.array([self.word2id_dict[word] for word in sent[0].split()], dtype=np.int64),
                sent[1]
            ) for sent in self.MRDataset_frame]

    def word_embeddings(self, path="./GoogleNews-vectors-negative300.bin/GoogleNews-vectors-negative300.bin"):
        # establish from google
        model = gensim.models.KeyedVectors.load_word2vec_format(path, binary=True)

        print('Please wait ... (it could take a while to load the file : {})'.format(path))
        word_dict = self.word2id_dict
        embedding_weights = np.random.uniform(-0.25, 0.25, (len(self.word2id_dict), 300))

        for word in word_dict:
            word_id = word_dict[word]
            if word in model.wv.vocab:
                embedding_weights[word_id, :] = model[word]
            return embedding_weights

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

    def __getitem__(self, idx):

        sample = self.MRDataset_wordid[idx]
        return sample

    def getsent(self, idx):

        sample = self.MRDataset_wordid[idx][0]
        return sample

    def getlabel(self, idx):

        label = self.MRDataset_wordid[idx][1]
        return label

    def word2id(self):

        return self.word2id_dict

    def id2word(self):

        id2word_dict = dict([val, key] for key, val in self.word2id_dict.items())
        return id2word_dict


class train_set(Dataset):

    def __init__(self, samples):
        self.train_frame = samples

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

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


class test_set(Dataset):

    def __init__(self, samples):
        self.test_frame = samples

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

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