-
Notifications
You must be signed in to change notification settings - Fork 3
/
data.py
54 lines (46 loc) · 1.75 KB
/
data.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
import os
import torch
import codecs
class Dictionary(object):
def __init__(self):
self.word2idx = {}
self.idx2word = []
self.word_count = {}
def add_word(self, word):
if word not in self.word_count:
self.word_count[word] = 1
else:
self.word_count[word] += 1
if word not in self.word2idx:
self.idx2word.append(word) #[I, do, not, like, ice, cream]
self.word2idx[word] = len(self.idx2word) - 1 # {I: 0, do: 1, not: 2, like: 3, ice: 4, cream: 5}
return self.word2idx[word]
def __len__(self): # function
return len(self.idx2word)
class Corpus(object):
def __init__(self, path):
self.dictionary = Dictionary()
self.train = self.tokenize(os.path.join(path, 'train.txt'))
self.valid = self.tokenize(os.path.join(path, 'valid.txt'))
self.test = self.tokenize(os.path.join(path, 'test.txt'))
def tokenize(self, path):
"""Tokenizes a text file."""
assert os.path.exists(path)
# Add words to the dictionary
with codecs.open(path,'r',encoding='utf8',errors='ignore') as f:
tokens = 0
for line in f:
words = line.split() + ['<eos>']
tokens += len(words)
for word in words:
self.dictionary.add_word(word)
# Tokenize file content
with codecs.open(path,'r',encoding='utf8',errors='ignore') as f:
ids = torch.LongTensor(tokens)
token = 0
for line in f:
words = line.split() + ['<eos>']
for word in words:
ids[token] = self.dictionary.word2idx[word]
token += 1
return ids