# pylint: disable=C0111,too-many-arguments,too-many-instance-attributes,too-many-locals,redefined-outer-name,fixme # pylint: disable=superfluous-parens, no-member, invalid-name from __future__ import print_function import sys sys.path.insert(0, "../../python") import numpy as np import mxnet as mx # The interface of a data iter that works for bucketing # # DataIter # - default_bucket_key: the bucket key for the default symbol. # # DataBatch # - provide_data: same as DataIter, but specific to this batch # - provide_label: same as DataIter, but specific to this batch # - bucket_key: the key for the bucket that should be used for this batch def default_read_content(path): with open(path) as ins: content = ins.read() content = content.replace('\n', ' ').replace('. ', ' ') return content def default_build_vocab(path): content = default_read_content(path) content = content.split(' ') words = set([x for x in content if len(x) > 0]) words = [x for x in words] words = sorted(words) the_vocab = {} idx = 1 # 0 is left for zero-padding the_vocab[' '] = 0 # put a dummy element here so that len(vocab) is correct for word in words: if len(word) == 0: continue if not word in the_vocab: the_vocab[word] = idx idx += 1 return the_vocab def default_text2id(sentence, the_vocab): words = sentence.split(' ') words = [the_vocab[w] for w in words if len(w) > 0] return words def default_gen_buckets(sentences, batch_size, the_vocab): len_dict = {} max_len = -1 for sentence in sentences: words = default_text2id(sentence, the_vocab) lw = len(words) if lw == 0: continue if lw > max_len: max_len = lw if lw in len_dict: len_dict[lw] += 1 else: len_dict[lw] = 1 print(len_dict) tl = 0 buckets = [] for l, n in len_dict.items(): # TODO: There are better heuristic ways to do this if n + tl >= batch_size: buckets.append(l) tl = 0 else: tl += n if tl > 0: buckets.append(max_len) return buckets class SimpleBatch(object): def __init__(self, data_names, data, label_names, label, bucket_key): self.data = data self.label = label self.data_names = data_names self.label_names = label_names self.bucket_key = bucket_key self.pad = 0 self.index = None # TODO: what is index? @property def provide_data(self): return [(n, x.shape) for n, x in zip(self.data_names, self.data)] @property def provide_label(self): return [(n, x.shape) for n, x in zip(self.label_names, self.label)] class DummyIter(mx.io.DataIter): "A dummy iterator that always return the same batch, used for speed testing" def __init__(self, real_iter): super(DummyIter, self).__init__() self.real_iter = real_iter self.provide_data = real_iter.provide_data self.provide_label = real_iter.provide_label self.batch_size = real_iter.batch_size for batch in real_iter: self.the_batch = batch break def __iter__(self): return self def next(self): return self.the_batch class BucketSentenceIter(mx.io.DataIter): def __init__(self, path, vocab, buckets, batch_size, init_states, data_name='data', label_name='label', seperate_char=' ', text2id=None, read_content=None): super(BucketSentenceIter, self).__init__() if text2id is None: self.text2id = default_text2id else: self.text2id = text2id if read_content is None: self.read_content = default_read_content else: self.read_content = read_content content = self.read_content(path) sentences = content.split(seperate_char) if len(buckets) == 0: buckets = default_gen_buckets(sentences, batch_size, vocab) print(buckets) self.vocab_size = len(vocab) self.data_name = data_name self.label_name = label_name buckets.sort() self.buckets = buckets self.data = [[] for _ in buckets] # pre-allocate with the largest bucket for better memory sharing self.default_bucket_key = max(buckets) for sentence in sentences: sentence = self.text2id(sentence, vocab) if len(sentence) == 0: continue for i, bkt in enumerate(buckets): if bkt >= len(sentence): self.data[i].append(sentence) break # we just ignore the sentence it is longer than the maximum # bucket size here # convert data into ndarrays for better speed during training data = [np.zeros((len(x), buckets[i])) for i, x in enumerate(self.data)] for i_bucket in range(len(self.buckets)): for j in range(len(self.data[i_bucket])): sentence = self.data[i_bucket][j] data[i_bucket][j, :len(sentence)] = sentence self.data = data # Get the size of each bucket, so that we could sample # uniformly from the bucket bucket_sizes = [len(x) for x in self.data] print("Summary of dataset ==================") for bkt, size in zip(buckets, bucket_sizes): print("bucket of len %3d : %d samples" % (bkt, size)) self.batch_size = batch_size self.make_data_iter_plan() self.init_states = init_states self.init_state_arrays = [mx.nd.zeros(x[1]) for x in init_states] self.provide_data = [('data', (batch_size, self.default_bucket_key))] + init_states self.provide_label = [('softmax_label', (self.batch_size, self.default_bucket_key))] def make_data_iter_plan(self): "make a random data iteration plan" # truncate each bucket into multiple of batch-size bucket_n_batches = [] for i in range(len(self.data)): bucket_n_batches.append(len(self.data[i]) / self.batch_size) self.data[i] = self.data[i][:int(bucket_n_batches[i]*self.batch_size)] bucket_plan = np.hstack([np.zeros(n, int)+i for i, n in enumerate(bucket_n_batches)]) np.random.shuffle(bucket_plan) bucket_idx_all = [np.random.permutation(len(x)) for x in self.data] self.bucket_plan = bucket_plan self.bucket_idx_all = bucket_idx_all self.bucket_curr_idx = [0 for x in self.data] self.data_buffer = [] self.label_buffer = [] for i_bucket in range(len(self.data)): data = np.zeros((self.batch_size, self.buckets[i_bucket])) label = np.zeros((self.batch_size, self.buckets[i_bucket])) self.data_buffer.append(data) self.label_buffer.append(label) def __iter__(self): init_state_names = [x[0] for x in self.init_states] for i_bucket in self.bucket_plan: data = self.data_buffer[i_bucket] label = self.label_buffer[i_bucket] i_idx = self.bucket_curr_idx[i_bucket] idx = self.bucket_idx_all[i_bucket][i_idx:i_idx+self.batch_size] self.bucket_curr_idx[i_bucket] += self.batch_size data[:] = self.data[i_bucket][idx] for k in range(len(data)): label[k] = sorted(data[k]) #count = len(data[k]) / 2 #for j in range(count): # data[j+count] = data[j] #label[:, :-1] = data[:, 1:] #label[:, -1] = 0 data_all = [mx.nd.array(data)] + self.init_state_arrays label_all = [mx.nd.array(label)] data_names = ['data'] + init_state_names label_names = ['softmax_label'] data_batch = SimpleBatch(data_names, data_all, label_names, label_all, self.buckets[i_bucket]) yield data_batch def reset(self): self.bucket_curr_idx = [0 for x in self.data]