|
import codecs |
|
import os |
|
import collections |
|
from six.moves import cPickle |
|
import numpy as np |
|
|
|
class TextLoader(): |
|
def __init__(self, data_dir, batch_size, seq_length, encoding='utf-8'): |
|
self.data_dir = data_dir |
|
self.batch_size = batch_size |
|
self.seq_length = seq_length |
|
self.encoding = encoding |
|
|
|
input_file = os.path.join(data_dir, "input.txt") |
|
vocab_file = os.path.join(data_dir, "vocab.pkl") |
|
tensor_file = os.path.join(data_dir, "data.npy") |
|
|
|
if not (os.path.exists(vocab_file) and os.path.exists(tensor_file)): |
|
print("reading text file") |
|
self.preprocess(input_file, vocab_file, tensor_file) |
|
else: |
|
print("loading preprocessed files") |
|
self.load_preprocessed(vocab_file, tensor_file) |
|
self.create_batches() |
|
self.reset_batch_pointer() |
|
|
|
def preprocess(self, input_file, vocab_file, tensor_file): |
|
with codecs.open(input_file, "r", encoding=self.encoding) as f: |
|
data = f.read() |
|
counter = collections.Counter(data) |
|
count_pairs = sorted(counter.items(), key=lambda x: -x[1]) |
|
self.chars, _ = zip(*count_pairs) |
|
self.vocab_size = len(self.chars) |
|
self.vocab = dict(zip(self.chars, range(len(self.chars)))) |
|
with open(vocab_file, 'wb') as f: |
|
cPickle.dump(self.chars, f) |
|
self.tensor = np.array(list(map(self.vocab.get, data))) |
|
np.save(tensor_file, self.tensor) |
|
|
|
def load_preprocessed(self, vocab_file, tensor_file): |
|
with open(vocab_file, 'rb') as f: |
|
self.chars = cPickle.load(f) |
|
self.vocab_size = len(self.chars) |
|
self.vocab = dict(zip(self.chars, range(len(self.chars)))) |
|
self.tensor = np.load(tensor_file) |
|
self.num_batches = int(self.tensor.size / (self.batch_size * |
|
self.seq_length)) |
|
|
|
def create_batches(self): |
|
self.num_batches = int(self.tensor.size / (self.batch_size * |
|
self.seq_length)) |
|
|
|
|
|
if self.num_batches==0: |
|
assert False, "Not enough data. Make seq_length and batch_size small." |
|
|
|
self.tensor = self.tensor[:self.num_batches * self.batch_size * self.seq_length] |
|
xdata = self.tensor |
|
ydata = np.copy(self.tensor) |
|
ydata[:-1] = xdata[1:] |
|
ydata[-1] = xdata[0] |
|
self.x_batches = np.split(xdata.reshape(self.batch_size, -1), self.num_batches, 1) |
|
self.y_batches = np.split(ydata.reshape(self.batch_size, -1), self.num_batches, 1) |
|
|
|
|
|
def next_batch(self): |
|
x, y = self.x_batches[self.pointer], self.y_batches[self.pointer] |
|
self.pointer += 1 |
|
return x, y |
|
|
|
def reset_batch_pointer(self): |
|
self.pointer = 0 |
|
|
