Processed form of the Penn Treebank dataset.
| 214 | |
| 215 | |
| 216 | class Datasets(object): |
| 217 | """Processed form of the Penn Treebank dataset.""" |
| 218 | |
| 219 | def __init__(self, path): |
| 220 | """Load the Penn Treebank dataset. |
| 221 | |
| 222 | Args: |
| 223 | path: Path to the data/ directory of the dataset from Tomas Mikolov's |
| 224 | webpage - http://www.fit.vutbr.cz/~imikolov/rnnlm/simple-examples.tgz |
| 225 | """ |
| 226 | |
| 227 | self.word2idx = {} # string -> integer id |
| 228 | self.idx2word = [] # integer id -> word string |
| 229 | # Files represented as a list of integer ids (as opposed to list of string |
| 230 | # words). |
| 231 | self.train = self.tokenize(os.path.join(path, "ptb.train.txt")) |
| 232 | self.valid = self.tokenize(os.path.join(path, "ptb.valid.txt")) |
| 233 | |
| 234 | def vocab_size(self): |
| 235 | return len(self.idx2word) |
| 236 | |
| 237 | def add(self, word): |
| 238 | if word not in self.word2idx: |
| 239 | self.idx2word.append(word) |
| 240 | self.word2idx[word] = len(self.idx2word) - 1 |
| 241 | |
| 242 | def tokenize(self, path): |
| 243 | """Read text file in path and return a list of integer token ids.""" |
| 244 | tokens = 0 |
| 245 | with tf.gfile.Open(path, "r") as f: |
| 246 | for line in f: |
| 247 | words = line.split() + ["<eos>"] |
| 248 | tokens += len(words) |
| 249 | for word in words: |
| 250 | self.add(word) |
| 251 | |
| 252 | # Tokenize file content |
| 253 | with tf.gfile.Open(path, "r") as f: |
| 254 | ids = np.zeros(tokens).astype(np.int64) |
| 255 | token = 0 |
| 256 | for line in f: |
| 257 | words = line.split() + ["<eos>"] |
| 258 | for word in words: |
| 259 | ids[token] = self.word2idx[word] |
| 260 | token += 1 |
| 261 | |
| 262 | return ids |
| 263 | |
| 264 | |
| 265 | def small_model(use_cudnn_rnn): |