-
Notifications
You must be signed in to change notification settings - Fork 10
/
data_handler.py
217 lines (199 loc) · 8.95 KB
/
data_handler.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
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
import numpy as np
from collections import defaultdict
from collections import deque
np.random.seed(1)
class DataHanlder:
"""
input date handler
for storing/sampling data, etc.
"""
def __init__(self, log_filename: str,
min_count: int=5,
sub_sampling_t: float=1e-5,
neg_sampling_t: float=0.75,
batch_size: int=64,
neg_sample_count: int=5,
half_window_size: int=2,
read_data_method='memory'):
"""
init function.
:param log_filename: word2vec format -> item1 item2 ... itemk\n, each line a user log
:param min_count: the same min_count as in word2vec
:param sub_sampling_t: sub sampling rate, higher than that would be sub sampled using
the word2vec paper using: p(w_i) = 1 - sqrt(sub_sampling / freq)
the word2vec code using: p(w_i) = 1 - (sqrt(sub_sampling / freq) + sub_sampling / freq)
we use word2vec code subsampling method here.
:param neg_sampling_t: the negative sampling t as in the word2vec. seems not shown in the paper, but implemented in
the code:
p(w_i) = f(w_i) ** neg_sampling / sum(f(w_i) ** neg_sampling for w_i in vocab)
:param read_data_method: method to read data:
'memory': load all the sentence to memory, fast but cost memory.
'file': load data from file, slower but save memory
"""
assert read_data_method in ('memory', 'file')
self.log_filename = log_filename
self.min_count = min_count
self.sub_sampling_t = sub_sampling_t
self.neg_sampling_t = neg_sampling_t
self.sentences = deque()
self.read_data_method = read_data_method
print('read dataset...')
self.vocab, self.word2id, self.id2word, self.total_word_count, self.sentence_len = self.gen_vocab()
print(f'got vocab {len(self.vocab)}, total_word_count {self.total_word_count}')
print('gen negative sample table...')
self.neg_sample_table = self.gen_negative_sample_table()
print('done.\ngen sub sampling table...')
self.sub_sampling_table = self.gen_subsample_table()
print('done.')
self.batch_size = batch_size
self.sentence_cursor = 0 # sentence cursor for generating batch
self.neg_sample_count = neg_sample_count
self.half_window_size = half_window_size
def gen_vocab(self):
"""
from log file generate vocabulary
:return: {item_id: freq}
"""
assert self.log_filename != ''
vocab_freq_dict = defaultdict(int)
total_word_count = 0
total_sent_count = 0
with open(self.log_filename, encoding='utf-8') as f:
for line in f:
total_sent_count += 1
item_ids = line.strip().split()
if self.read_data_method == 'memory':
self.sentences.append(item_ids)
for item_id in item_ids:
vocab_freq_dict[item_id] += 1
total_word_count += 1
vocab, word2id, id2word = {}, {}, {}
index = 0
for item_id, freq in vocab_freq_dict.items():
if freq < self.min_count:
continue
vocab[item_id] = freq
word2id[item_id] = index
id2word[index] = item_id
index += 1
return vocab, word2id, id2word, total_word_count, total_sent_count
def gen_subsample_table(self):
"""
sub sampling rate, higher than that would be sub sampled using
the word2vec paper using: p(w_i) = 1 - sqrt(sub_sampling / freq)
the word2vec code using: p(w_i) = 1 - (sqrt(sub_sampling / freq) + sub_sampling / freq)
we use word2vec code sub sampling method here.
:return: {word_id: sample_score}
"""
def sub_sampling(_freq):
return (self.sub_sampling_t / 1.0 / _freq) ** 0.5 + self.sub_sampling_t / 1.0 / _freq
# word freq count to word freq ratio
sub_sample_tbl = {item: freq / 1.0 / self.total_word_count
for item, freq in self.vocab.items()
if freq / 1.0 / self.total_word_count > self.sub_sampling_t}
# freq to score
sub_sample_tbl = {item: sub_sampling(_freq) for item, _freq in sub_sample_tbl.items()}
# word to id
sub_sample_tbl = {self.word2id[i]: j for i, j in sub_sample_tbl.items() if j < 1}
return sub_sample_tbl
def gen_negative_sample_table(self):
"""
implemented same as word2vec c code.
The way this selection is implemented in the C code is interesting. They have a large array with 100M elements
(which they refer to as the unigram table). They fill this table with the index of each word in the vocabulary
multiple times, and the number of times a word’s index appears in the table is given by P(wi) * table_size.
p(w_i) = f(w_i) ** neg_sampling / sum(f(w_i) ** neg_sampling for w_i in vocab)
:return:
"""
sample_tbl_size = 1e8
sample_tbl = []
pow_freq = np.array(list(self.vocab.values())) ** self.neg_sampling_t
pow_total_freq = sum(pow_freq)
r = pow_freq / pow_total_freq
count = np.round(r * sample_tbl_size)
for item_id, _count in enumerate(count):
sample_tbl += [item_id] * int(_count)
sample_tbl = np.array(sample_tbl)
return sample_tbl
def gen_batch(self):
"""
yield batch
:return: pos_pairs -> [(w1, w2) * batch_size ], neg samples -> [self.neg_sample_count * batch_size]
"""
f = open(self.log_filename)
pos_pairs = []
pos_1, pos_2, neg_samples = [], [], []
while True:
while len(pos_pairs) < self.batch_size:
# pos
sentence = []
if self.read_data_method == 'memory':
sentence = self.sentences.popleft()
elif self.read_data_method == 'file':
sentence = f.readline()
if not sentence:
f = open(self.log_filename)
sentence = f.readline()
sentence = sentence.strip().split()
self.sentence_cursor += 1
# to word id
word_ids = [self.word2id[item_id] for item_id in sentence if item_id in self.word2id]
for i, word_id in enumerate(word_ids):
pos_pairs += [(word_id, word_ids[j])
for j in range(max(0, i - self.half_window_size),
min(i + self.half_window_size + 1, len(word_ids) - 1))
if j != i]
if self.read_data_method == 'memory':
self.sentences.append(sentence)
# neg
for pos1, pos2 in pos_pairs[len(neg_samples):]:
pos_1.append(pos1)
pos_2.append(pos2)
neg_samples.append(self.negative_sampling(pos1, pos2))
yield (pos_1[:self.batch_size], pos_2[:self.batch_size], neg_samples[:self.batch_size])
pos_pairs = pos_pairs[self.batch_size:]
pos_1 = pos_1[self.batch_size:]
pos_2 = pos_2[self.batch_size:]
neg_samples = neg_samples[self.batch_size:]
def negative_sampling(self, pos1, pos2):
"""
negative sample, shall not equal to pos1 and pos2
:param pos1:
:param pos2:
:return:
"""
negs = []
while len(negs) < self.neg_sample_count:
_negs = np.random.choice(self.neg_sample_table, size=self.neg_sample_count - len(negs))
negs += [i for i in _negs if i != pos1 and i != pos2]
return negs[:self.neg_sample_count]
def thread_read_data_(self):
"""
useless since GIL exists.
though it is a pure IO thread, it still does not accelerate.
:return:
"""
f = open(self.log_filename)
while True:
if len(self.sentences) < 1000:
sentence = f.readline()
if not sentence:
f = open(self.log_filename)
sentence = f.readline()
item_ids = sentence.strip().split(' ')
self.sentences.append(item_ids)
def test():
handler = DataHanlder('../data/trainset.txt', read_data_method='memory')
i = 0
import time
start = time.time()
for pos_1, pos_2, neg_samples in handler.gen_batch():
if handler.sentence_cursor >= 2 * handler.total_word_count:
break
i += 1
if i > 1000:
break
end = time.time()
print('1000 iter using', end - start)
if __name__ == '__main__':
test()