built simple test for dataset

# WILL BE REPLACED WITH JARED'S

class JaredDataset(object):

    def __init__(self):

        self.doc_idx = []

    def __init__(self, doc_idx, sizes, sentences):

        self.doc_idx = doc_idx

        self.num_docs = len(self.doc_idx) - 1

        self.sizes = []

        self.sentences = []

        self.sizes = sizes

        self.sentences = sentences

    def __getitem__(self, idx):

        return self.sentences[idx]

            # Document sentences are in [sent_index_first, sent_index_last).

            sent_index_first = indexed_dataset.doc_idx[doc_index]

            sent_index_last = indexed_dataset.doc_idx[doc_index+1]

            assert sent_index_last >= sent_index_first:

            assert sent_index_last >= sent_index_first

            # Empty docs.

            if (sent_index_last - sent_index_first) == 0:

            # Loop through sentences.

            sent_index = sent_index_first

            target_seq_length = get_target_seq_length(max_num_tokens,

                                                      short_seq_prob, rng)

                                                      short_seq_prob, np_rng)

            size = 0

            while sent_index < sent_index_last:

                exceeded_target_size = (size >= target_seq_length)

                # If only one sentence is left in the document.

                only_one_sent_left = (sent_index == (sent_index_last - 1))

                # If we have at least two sentneces.

                have_more_than_one_sent = (sent_index - sent_index_first) > 1

                # If we have reached end of the document.

                reached_end_of_doc = (sent_index == sent_index_last)

                if (exceeded_target_size and not only_one_sent_left) or \

                   reached_end_of_doc:

                if (exceeded_target_size and not only_one_sent_left and

                    have_more_than_one_sent) or reached_end_of_doc:

                    assert (sent_index - sent_index_first) > 1

                    assert size > 1

                    # Add the sample.

                    samples.append([sent_index_first, sent_index])

                    samples.append([sent_index_first, sent_index,

                                    target_seq_length])

                    # Reset indices

                    sent_index_first = sent_index

                    target_seq_length = get_target_seq_length(max_num_tokens,

                                                              short_seq_prob,

                                                              rng)

                                                              np_rng)

                    size = 0

                    num_sentences = 0

class AlbertDataSet(Dataset):

    def __init__(self, tokenizer, num_epochs, masked_lm_prob, max_seq_length

                 short_seq_prob, seed):

    def __init__(self, indexed_dataset, tokenizer, num_epochs,

                 masked_lm_prob, max_seq_length, short_seq_prob, seed):

        # Params to store.

        self.seed = seed

        self.masked_lm_prob = masked_lm_prob

        self.max_seq_length = max_seq_length

        # Build the indexed dataset.

        self.indexed_dataset = JaredDataset()

        # Indexed dataset.

        self.indexed_dataset = indexed_dataset

        # Build the samples mapping.

        self.samples_mapping = build_training_samples_mapping(

if __name__ == '__main__':

    print('dataset ...')

    from bert_tokenization import FullTokenizer

    import json

    import nltk

    nltk.download('punkt')

    def document_generator_provider(input_file):

        with open(input_file, 'r') as ifile:

            for document in ifile:

                data = json.loads(document)

                text = data['text']

                sentences = []

                for line in text.split('\n'):

                    if line != '\n':

                        sentences.extend(nltk.tokenize.sent_tokenize(line))

                yield sentences

    input_file = '/raid/mshoeybi/data/albert/sample/samples_11.json'

    vocab_file = '/raid/mshoeybi/data/albert/bert_vocab/vocab.txt'

    tokenizer = FullTokenizer(vocab_file, do_lower_case=True)

    document_generator = document_generator_provider(input_file)

    doc_idx = [0]

    sizes = []

    sentences_list = []

    for sentences in document_generator:

        doc_idx.append(len(sentences))

        for sentence in sentences:

            tokens = tokenizer.tokenize(sentence)

            ids = tokenizer.convert_tokens_to_ids(tokens)

            sizes.append(len(ids))

            sentences_list.append(ids)

    for i in range(1, len(doc_idx)):

        doc_idx[i] += doc_idx[i-1]

    indexed_dataset = JaredDataset(doc_idx, sizes, sentences_list)

    dataset = AlbertDataSet(indexed_dataset=indexed_dataset,

                            tokenizer=tokenizer,

                            num_epochs=3,

                            masked_lm_prob=0.15,

                            max_seq_length=512,

                            short_seq_prob=0.1,

                            seed=1234)