| 56 | return [y for x, y in tuples] |
| 57 | |
| 58 | def get_data(): |
| 59 | word2idx = {} |
| 60 | current_idx = 0 |
| 61 | X = [] |
| 62 | Y = [] |
| 63 | for fn, label in zip(('robert_frost.txt', 'edgar_allan_poe.txt'), (0, 1)): |
| 64 | count = 0 |
| 65 | for line in open(fn): |
| 66 | line = line.rstrip() |
| 67 | if line: |
| 68 | print(line) |
| 69 | # tokens = remove_punctuation(line.lower()).split() |
| 70 | tokens = get_tags(line) |
| 71 | if len(tokens) > 1: |
| 72 | # scan doesn't work nice here, technically could fix... |
| 73 | for token in tokens: |
| 74 | if token not in word2idx: |
| 75 | word2idx[token] = current_idx |
| 76 | current_idx += 1 |
| 77 | sequence = np.array([word2idx[w] for w in tokens]) |
| 78 | X.append(sequence) |
| 79 | Y.append(label) |
| 80 | count += 1 |
| 81 | print(count) |
| 82 | if count >= 50: |
| 83 | break |
| 84 | print("Vocabulary:", word2idx.keys()) |
| 85 | return X, Y, current_idx |
| 86 | |
| 87 | |
| 88 | def main(): |