-
Notifications
You must be signed in to change notification settings - Fork 0
/
data.py
67 lines (48 loc) · 1.79 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
55
56
57
58
59
60
61
62
63
64
65
66
67
import os
import torch
from io import open
from util import load_data
class Dictionary(object):
# Word level or character level
def __init__(self):
self.element2idx = {}
self.idx2element = []
def add_element(self, word):
if word not in self.element2idx:
self.idx2element.append(word)
self.element2idx[word] = len(self.idx2element) - 1
return self.element2idx[word]
def __len__(self):
return len(self.idx2element)
class Corpus(object):
def __init__(self, args):
self.dictionary = Dictionary()
self.use_char_embeds = args.type == 'char'
self.train = self.tokenize(os.path.join(args.data, 'train.json'), use_char_embeds=self.use_char_embeds)
self.valid = self.tokenize(os.path.join(args.data, 'valid.json'), use_char_embeds=self.use_char_embeds)
def tokenize(self, path, use_char_embeds=False):
assert os.path.exists(path)
data = load_data(path)
data = ' '.join(data['reviewText'])
tokens = 0
if use_char_embeds:
data = list(data)
words = data + ['<eos>']
else :
words = data.split() + ['<eos>']
# Add words to the dictionary
tokens += len(words)
for word in words:
self.dictionary.add_element(word)
# Tokenbize file content
ids = torch.LongTensor(tokens)
token = 0
for word in words:
ids[token] = self.dictionary.element2idx[word]
token += 1
return ids
def describe(self):
embed_type = 'char' if self.use_char_embeds else 'word'
s = f'Using {embed_type} level embeddings\n'
s += f'Size of dictionary {len(self.dictionary)} \n'
return s