-
Notifications
You must be signed in to change notification settings - Fork 1
/
word2vec_cbow_hs.py
707 lines (571 loc) · 31.3 KB
/
word2vec_cbow_hs.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
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Copyright (C) 2013 Radim Rehurek <me@radimrehurek.com>
# Licensed under the GNU LGPL v2.1 - http://www.gnu.org/licenses/lgpl.html
#
# Modified by Sébastien Jean
"""
Module for deep learning via *hierarchical softmax CBOW* from [1]_.
The training algorithm was originally ported from the C package https://code.google.com/p/word2vec/
and extended with additional functionality.
**Install Cython with `pip install cython` to use optimized word2vec training** (70x speedup [2]_).
Initialize a model with e.g.::
>>> model = Word2Vec(sentences, size=100, window=5, min_count=5, workers=4)
Persist a model to disk with::
>>> model.save(fname)
>>> model = Word2Vec.load(fname) # you can continue training with the loaded model!
The model can also be instantiated from an existing file on disk in the word2vec C format::
>>> model = Word2Vec.load_word2vec_format('/tmp/vectors.txt', binary=False) # C text format
>>> model = Word2Vec.load_word2vec_format('/tmp/vectors.bin', binary=True) # C binary format
You can perform various syntactic/semantic NLP word tasks with the model. Some of them
are already built-in::
>>> model.most_similar(positive=['woman', 'king'], negative=['man'])
[('queen', 0.50882536), ...]
>>> model.doesnt_match("breakfast cereal dinner lunch".split())
'cereal'
>>> model.similarity('woman', 'man')
0.73723527
>>> model['computer'] # raw numpy vector of a word
array([-0.00449447, -0.00310097, 0.02421786, ...], dtype=float32)
and so on.
.. [1] Tomas Mikolov, Kai Chen, Greg Corrado, and Jeffrey Dean. Efficient Estimation of Word Representations in Vector Space. In Proceedings of Workshop at ICLR, 2013.
.. [2] Optimizing word2vec in gensim, http://radimrehurek.com/2013/09/word2vec-in-python-part-two-optimizing/
"""
import logging
import sys
import os
import heapq
import time
import itertools
import threading
from multiprocessing.pool import ThreadPool
from Queue import Queue
from numpy import zeros_like, empty, exp, dot, outer, random, dtype, get_include,\
float32 as REAL, uint32, seterr, array, uint8, vstack, argsort, fromstring, inf
from numpy import shape #to delete
logger = logging.getLogger("word2vec_mod")
from gensim import utils, matutils # utility fnc for pickling, common scipy operations etc
try:
# try to compile and use the faster cython version
import pyximport
pyximport.install(setup_args={"include_dirs": get_include()})
from word2vec_inner_cbow_hs import train_sentence, FAST_VERSION
except:
# failed... fall back to plain numpy (20-80x slower training than the above)
FAST_VERSION = -1
def train_sentence(model, sentence, alpha, work=None, neu1=None): #mod
"""
Update CBOW hierarchical softmax model by training on a single sentence.
The sentence is a list of Vocab objects (or None, where the corresponding
word is not in the vocabulary. Called internally from `Word2Vec.train()`.
"""
for pos, word in enumerate(sentence):
if word is None:
continue # OOV word in the input sentence => skip
if model.reduce > 0:
reduced_window = random.randint(model.window) # `b` in the original word2vec code
else:
reduced_window = 0
# Combine all surrounding words into an appropriate input
start = max(0, pos - model.window + reduced_window)
l1 = matutils.zeros_aligned((model.layer1_size), dtype=REAL)#Initialize input
weights = 0.
weights_2 = 0.
for pos2, word2 in enumerate(sentence[start : pos + model.window + 1 - reduced_window], start):
if pos2 == pos or word2 is None:
pass
else:
weights += word2.count_power
weights_2 += word2.count_power_2
l1 = l1 + word2.count_power*model.syn0[word2.index]
if weights > 0.0000000000000001:
regularization = weights/weights_2
l1 = l1 / weights
l2a = model.syn1[word.point] # 2d matrix, codelen x layer1_size
fa = 1.0 / (1.0 + exp(-dot(l1, l2a.T))) # propagate hidden -> output
ga = (1 - word.code - fa) * alpha # vector of error gradients multiplied by the learning rate
model.syn1[word.point] += outer(ga, l1) # learn hidden -> output
for pos2, word2 in enumerate(sentence[start : pos + model.window + 1 - reduced_window], start):
if pos2 == pos or word2 is None:
pass
else:
model.syn0[word2.index] += regularization*word2.count_power*dot(ga, l2a) #MUST BE MODIFIED
return len([word for word in sentence if word is not None])
class Vocab(object):
"""A single vocabulary item, used internally for constructing binary trees (incl. both word leaves and inner nodes)."""
def __init__(self, **kwargs):
self.count = 0
self.count_power = 0.
self.count_power_2 = 0.
self.__dict__.update(kwargs)
def __lt__(self, other): # used for sorting in a priority queue
return self.count < other.count
def __str__(self):
vals = ['%s:%r' % (key, self.__dict__[key]) for key in sorted(self.__dict__) if not key.startswith('_')]
return "<" + ', '.join(vals) + ">"
class Word2Vec(utils.SaveLoad):
"""
Class for training, using and evaluating neural networks described in https://code.google.com/p/word2vec/
The model can be stored/loaded via its `save()` and `load()` methods, or stored/loaded in a format
compatible with the original word2vec implementation via `save_word2vec_format()` and `load_word2vec_format()`.
"""
def __init__(self, sentences=None, size=100, alpha=0.025, window=5, min_count=5, max_count=inf, seed=1, workers=1, min_alpha=0.0001, reduce=1, alpha_decay=1.0, c_power=0., direction=0, auto_train=1):
"""
Initialize the model from an iterable of `sentences`. Each sentence is a
list of words (utf8 strings) that will be used for training.
See :class:`BrownCorpus` in this module for an example.
If you don't supply `sentences`, the model is left uninitialized -- use if
you plan to initialize it in some other way.
`size` is the dimensionality of the feature vectors.
`window` is the maximum distance between the current and predicted word within a sentence.
`alpha` is the initial learning rate (will linearly drop to zero as training progresses).
`seed` = for the random number generator.
`min_count` = ignore all words with total frequency lower than this.
`workers` = use this many worker threads to train the model (=faster training with multicore machines)
"""
self.vocab = {} # mapping from a word (string) to a Vocab object
self.index2word = [] # map from a word's matrix index (int) to word (string)
self.layer1_size = int(size)
self.alpha = float(alpha)
self.window = int(window)
self.seed = seed
self.min_count = min_count
self.max_count = max_count
self.workers = workers
self.min_alpha = min_alpha
self.reduce = int(reduce)
self.alpha_decay = float(alpha_decay)
self.c_power = float(c_power)
self.direction = int(direction) # direction is only implemented in Cython, not in the pure python/numpy version
self.auto_train = int(auto_train)
if sentences is not None:
self.build_vocab(sentences)
if self.auto_train > 0:
self.train(sentences)
def create_binary_tree(self):
"""
Create a binary Huffman tree using stored vocabulary word counts. Frequent words
will have shorter binary codes. Called internally from `build_vocab()`.
"""
logger.info("constructing a huffman tree from %i words" % len(self.vocab))
# build the huffman tree
heap = self.vocab.values()
heapq.heapify(heap)
for i in xrange(len(self.vocab) - 1):
min1, min2 = heapq.heappop(heap), heapq.heappop(heap)
heapq.heappush(heap, Vocab(count=min1.count + min2.count, index=i + len(self.vocab), left=min1, right=min2))
# recurse over the tree, assigning a binary code to each vocabulary word
if heap:
max_depth, stack = 0, [(heap[0], [], [])]
while stack:
node, codes, points = stack.pop()
if node.index < len(self.vocab):
# leaf node => store its path from the root
node.code, node.point = codes, points
max_depth = max(len(codes), max_depth)
else:
# inner node => continue recursion
points = array(list(points) + [node.index - len(self.vocab)], dtype=uint32)
stack.append((node.left, array(list(codes) + [0], dtype=uint8), points))
stack.append((node.right, array(list(codes) + [1], dtype=uint8), points))
logger.info("built huffman tree with maximum node depth %i" % max_depth)
def build_vocab(self, sentences):
"""
Build vocabulary from a sequence of sentences (can be a once-only generator stream).
Each sentence must be a list of utf8 strings.
"""
logger.info("collecting all words and their counts")
sentence_no, vocab = -1, {}
total_words = lambda: sum(v.count for v in vocab.itervalues())
for sentence_no, sentence in enumerate(sentences):
if sentence_no % 10000 == 0:
logger.info("PROGRESS: at sentence #%i, processed %i words and %i word types" %
(sentence_no, total_words(), len(vocab)))
for word in sentence:
if word in vocab:
vocab[word].count += 1
else:
vocab[word] = Vocab(count=1)
logger.info("collected %i word types from a corpus of %i words and %i sentences" %
(len(vocab), total_words(), sentence_no + 1))
# assign a unique index to each word
self.vocab, self.index2word = {}, []
for word, v in vocab.iteritems():
if v.count >= self.min_count and v.count <= self.max_count:
try:
utils.to_utf8(word)
v.index = len(self.vocab)
self.index2word.append(word)
self.vocab[word] = v
self.vocab[word].count_power = pow(v.count, self.c_power)
self.vocab[word].count_power_2 = pow(v.count, 2*self.c_power)
except:
pass
logger.info("total %i word types after removing those with count<%s" % (len(self.vocab), self.min_count))
# add info about each word's Huffman encoding
self.create_binary_tree()
self.reset_weights()
def train(self, sentences, total_words=None, word_count=0, chunksize=100):
"""
Update the model's neural weights from a sequence of sentences (can be a once-only generator stream).
Each sentence must be a list of utf8 strings.
"""
if FAST_VERSION < 0:
import warnings
warnings.warn("Cython compilation failed, training will be slow. Do you have Cython installed? `pip install cython`")
logger.info("training model with %i workers on %i vocabulary and %i features" % (self.workers, len(self.vocab), self.layer1_size))
if not self.vocab:
raise RuntimeError("you must first build vocabulary before training the model")
start, next_report = time.time(), [1.0]
word_count, total_words = [word_count], total_words or sum(v.count for v in self.vocab.itervalues())
jobs = Queue(maxsize=2 * self.workers) # buffer ahead only a limited number of jobs.. this is the reason we can't simply use ThreadPool :(
lock = threading.Lock() # for shared state (=number of words trained so far, log reports...)
def worker_train():
"""Train the model, lifting lists of sentences from the jobs queue."""
work = matutils.zeros_aligned(self.layer1_size, dtype=REAL) # each thread must have its own work memory
neu1 = matutils.zeros_aligned(self.layer1_size, dtype=REAL) # each thread must have its own work memory
while True:
job = jobs.get()
if job is None: # data finished, exit
break
# update the learning rate before every job
alpha = max(self.min_alpha, self.alpha * (1 - 1.0 * self.alpha_decay * word_count[0] / total_words))
# how many words did we train on? out-of-vocabulary (unknown) words do not count
job_words = sum(train_sentence(self, sentence, alpha, work, neu1) for sentence in job)
with lock:
word_count[0] += job_words
elapsed = time.time() - start
if elapsed >= next_report[0]:
logger.info("PROGRESS: at %.2f%% words, alpha %.05f, %.0f words/s" %
(100.0 * word_count[0] / total_words, alpha, word_count[0] / elapsed if elapsed else 0.0))
next_report[0] = elapsed + 1.0 # don't flood the log, wait at least a second between progress reports
workers = [threading.Thread(target=worker_train) for _ in xrange(self.workers)]
for thread in workers:
thread.daemon = True # make interrupting the process with ctrl+c easier
thread.start()
# convert input strings to Vocab objects (or None for OOV words), and start filling the jobs queue
no_oov = ([self.vocab.get(word, None) for word in sentence] for sentence in sentences)
for job_no, job in enumerate(utils.grouper(no_oov, chunksize)):
logger.debug("putting job #%i in the queue, qsize=%i" % (job_no, jobs.qsize()))
jobs.put(job)
logger.info("reached the end of input; waiting to finish %i outstanding jobs" % jobs.qsize())
for _ in xrange(self.workers):
jobs.put(None) # give the workers heads up that they can finish -- no more work!
for thread in workers:
thread.join()
elapsed = time.time() - start
logger.info("training on %i words took %.1fs, %.0f words/s" %
(word_count[0], elapsed, word_count[0] / elapsed if elapsed else 0.0))
return word_count[0]
def reset_weights(self):
"""Reset all projection weights to an initial (untrained) state, but keep the existing vocabulary."""
random.seed(self.seed)
self.syn0 = matutils.zeros_aligned((len(self.vocab), self.layer1_size), dtype=REAL)
self.syn1 = matutils.zeros_aligned((len(self.vocab), self.layer1_size), dtype=REAL)
self.syn0 += (random.rand(len(self.vocab), self.layer1_size) - 0.5) / self.layer1_size
self.syn0norm = None
def save_word2vec_format(self, fname, binary=False):
"""
Store the input-hidden weight matrix in the same format used by the original
C word2vec-tool, for compatibility.
"""
logger.info("storing %sx%s projection weights into %s" % (len(self.vocab), self.layer1_size, fname))
assert (len(self.vocab), self.layer1_size) == self.syn0.shape
with open(fname, 'wb') as fout:
fout.write("%s %s\n" % self.syn0.shape)
# store in sorted order: most frequent words at the top
for word, vocab in sorted(self.vocab.iteritems(), key=lambda item: -item[1].count):
word = utils.to_utf8(word) # always store in utf8
row = self.syn0[vocab.index]
if binary:
fout.write("%s %s\n" % (word, row.tostring()))
else:
fout.write("%s %s\n" % (word, ' '.join("%f" % val for val in row)))
@classmethod
def load_word2vec_format(cls, fname, binary=False):
"""
Load the input-hidden weight matrix from the original C word2vec-tool format.
Note that the information loaded is incomplete (the binary tree is missing),
so while you can query for word similarity etc., you cannot continue training
with a model loaded this way.
"""
logger.info("loading projection weights from %s" % (fname))
with open(fname) as fin:
header = fin.readline()
vocab_size, layer1_size = map(int, header.split()) # throws for invalid file format
result = Word2Vec(size=layer1_size)
result.syn0 = empty((vocab_size, layer1_size), dtype=REAL)
if binary:
binary_len = dtype(REAL).itemsize * layer1_size
for line_no in xrange(vocab_size):
# mixed text and binary: read text first, then binary
word = []
while True:
ch = fin.read(1)
if ch == ' ':
word = ''.join(word)
break
word.append(ch)
result.vocab[word] = Vocab(index=line_no, count=vocab_size - line_no)
result.index2word.append(word)
result.syn0[line_no] = fromstring(fin.read(binary_len), dtype=REAL)
fin.read(1) # newline
else:
for line_no, line in enumerate(fin):
parts = line.split()
assert len(parts) == layer1_size + 1
word, weights = parts[0], map(REAL, parts[1:])
result.vocab[word] = Vocab(index=line_no, count=vocab_size - line_no)
result.index2word.append(word)
result.syn0[line_no] = weights
logger.info("loaded %s matrix from %s" % (result.syn0.shape, fname))
result.init_sims()
return result
def most_similar(self, positive=[], negative=[], topn=10):
"""
Find the top-N most similar words. Positive words contribute positively towards the
similarity, negative words negatively.
This method computes cosine similarity between a simple mean of the projection
weight vectors of the given words, and corresponds to the `word-analogy` and
`distance` scripts in the original word2vec implementation.
Example::
>>> trained_model.most_similar(positive=['woman', 'king'], negative=['man'])
[('queen', 0.50882536), ...]
"""
self.init_sims()
if isinstance(positive, basestring) and not negative:
# allow calls like most_similar('dog'), as a shorthand for most_similar(['dog'])
positive = [positive]
# add weights for each word, if not already present; default to 1.0 for positive and -1.0 for negative words
positive = [(word, 1.0) if isinstance(word, basestring) else word for word in positive]
negative = [(word, -1.0) if isinstance(word, basestring) else word for word in negative]
# compute the weighted average of all words
all_words, mean = set(), []
for word, weight in positive + negative:
if word in self.vocab:
mean.append(weight * matutils.unitvec(self.syn0[self.vocab[word].index]))
all_words.add(self.vocab[word].index)
else:
raise KeyError("word '%s' not in vocabulary" % word)
if not mean:
raise ValueError("cannot compute similarity with no input")
mean = matutils.unitvec(array(mean).mean(axis=0)).astype(REAL)
dists = dot(self.syn0norm, mean)
if not topn:
return dists
best = argsort(dists)[::-1][:topn + len(all_words)]
# ignore (don't return) words from the input
result = [(self.index2word[sim], dists[sim]) for sim in best if sim not in all_words]
return result[:topn]
def doesnt_match(self, words):
"""
Which word from the given list doesn't go with the others?
Example::
>>> trained_model.doesnt_match("breakfast cereal dinner lunch".split())
'cereal'
"""
words = [word for word in words if word in self.vocab] # filter out OOV words
logger.debug("using words %s" % words)
if not words:
raise ValueError("cannot select a word from an empty list")
vectors = vstack(matutils.unitvec(self.syn0[self.vocab[word].index]) for word in words).astype(REAL)
mean = matutils.unitvec(vectors.mean(axis=0)).astype(REAL)
dists = dot(vectors, mean)
return sorted(zip(dists, words))[0][1]
def __getitem__(self, word):
"""
Return a word's representations in vector space, as a 1D numpy array.
Example::
>>> trained_model['woman']
array([ -1.40128313e-02, ...]
"""
return self.syn0[self.vocab[word].index]
def __contains__(self, word):
return word in self.vocab
def similarity(self, w1, w2):
"""
Compute cosine similarity between two words.
Example::
>>> trained_model.similarity('woman', 'man')
0.73723527
>>> trained_model.similarity('woman', 'woman')
1.0
"""
return dot(matutils.unitvec(self[w1]), matutils.unitvec(self[w2]))
def init_sims(self):
if getattr(self, 'syn0norm', None) is None:
logger.info("precomputing L2-norms of word weight vectors")
self.syn0norm = vstack(matutils.unitvec(vec) for vec in self.syn0).astype(REAL)
def accuracy(self, questions, restrict_vocab=30000):
"""
Compute accuracy of the model. `questions` is a filename where lines are
4-tuples of words, split into sections by ": SECTION NAME" lines.
See https://code.google.com/p/word2vec/source/browse/trunk/questions-words.txt for an example.
The accuracy is reported (=printed to log and returned as a list) for each
section separately, plus there's one aggregate summary at the end.
Use `restrict_vocab` to ignore all questions containing a word whose frequency
is not in the top-N most frequent words (default top 30,000).
This method corresponds to the `compute-accuracy` script of the original C word2vec.
"""
ok_vocab = dict(sorted(self.vocab.iteritems(), key=lambda item: -item[1].count)[:restrict_vocab])
ok_index = set(v.index for v in ok_vocab.itervalues())
def log_accuracy(section):
correct, incorrect = section['correct'], section['incorrect']
if correct + incorrect > 0:
logger.info("%s: %.1f%% (%i/%i)" %
(section['section'], 100.0 * correct / (correct + incorrect),
correct, correct + incorrect))
sections, section = [], None
for line_no, line in enumerate(open(questions)):
# TODO: use level3 BLAS (=evaluate multiple questions at once), for speed
if line.startswith(': '):
# a new section starts => store the old section
if section:
sections.append(section)
log_accuracy(section)
section = {'section': line.lstrip(': ').strip(), 'correct': 0, 'incorrect': 0}
else:
if not section:
raise ValueError("missing section header before line #%i in %s" % (line_no, questions))
try:
a, b, c, expected = [word.lower() for word in line.split()] # TODO assumes vocabulary preprocessing uses lowercase, too...
except:
logger.info("skipping invalid line #%i in %s" % (line_no, questions))
if a not in ok_vocab or b not in ok_vocab or c not in ok_vocab or expected not in ok_vocab:
logger.debug("skipping line #%i with OOV words: %s" % (line_no, line))
continue
ignore = set(self.vocab[v].index for v in [a, b, c]) # indexes of words to ignore
predicted = None
# find the most likely prediction, ignoring OOV words and input words
for index in argsort(self.most_similar(positive=[b, c], negative=[a], topn=False))[::-1]:
if index in ok_index and index not in ignore:
predicted = self.index2word[index]
if predicted != expected:
logger.debug("%s: expected %s, predicted %s" % (line.strip(), expected, predicted))
break
section['correct' if predicted == expected else 'incorrect'] += 1
if section:
# store the last section, too
sections.append(section)
log_accuracy(section)
total = {'section': 'total', 'correct': sum(s['correct'] for s in sections), 'incorrect': sum(s['incorrect'] for s in sections)}
log_accuracy(total)
sections.append(total)
return sections
def __str__(self):
return "Word2Vec(vocab=%s, size=%s, alpha=%s)" % (len(self.index2word), self.layer1_size, self.alpha)
class BrownCorpus(object):
"""Iterate over sentences from the Brown corpus (part of NLTK data)."""
def __init__(self, dirname):
self.dirname = dirname
def __iter__(self):
for fname in os.listdir(self.dirname):
fname = os.path.join(self.dirname, fname)
if not os.path.isfile(fname):
continue
for line in open(fname):
# each file line is a single sentence in the Brown corpus
# each token is WORD/POS_TAG
token_tags = [t.split('/') for t in line.split() if len(t.split('/')) == 2]
# ignore words with non-alphabetic tags like ",", "!" etc (punctuation, weird stuff)
words = ["%s/%s" % (token.lower(), tag[:2]) for token, tag in token_tags if tag[:2].isalpha()]
if not words: # don't bother sending out empty sentences
continue
yield words
class Text8Corpus(object):
"""Iterate over sentences from the "text8" corpus, unzipped from http://mattmahoney.net/dc/text8.zip ."""
def __init__(self, fname):
self.fname = fname
def __iter__(self):
# the entire corpus is one gigantic line -- there are no sentence marks at all
# so just split the sequence of tokens arbitrarily: 1 sentence = 1000 tokens
sentence, rest, max_sentence_length = [], '', 1000
with open(self.fname) as fin:
while True:
text = rest + fin.read(8192) # avoid loading the entire file (=1 line) into RAM
if text == rest: # EOF
sentence.extend(rest.split()) # return the last chunk of words, too (may be shorter/longer)
if sentence:
yield sentence
break
last_token = text.rfind(' ') # the last token may have been split in two... keep it for the next iteration
words, rest = (text[:last_token].split(), text[last_token:].strip()) if last_token >= 0 else ([], text)
sentence.extend(words)
while len(sentence) >= max_sentence_length:
yield sentence[:max_sentence_length]
sentence = sentence[max_sentence_length:]
class Text8Corpus_small(object):
"""Iterate over sentences from the "text8" corpus, unzipped from http://mattmahoney.net/dc/text8.zip ."""
def __init__(self, fname):
self.fname = fname
def __iter__(self):
# the entire corpus is one gigantic line -- there are no sentence marks at all
# so just split the sequence of tokens arbitrarily: 1 sentence = 1000 tokens
sentence, rest, max_sentence_length = [], '', 20
with open(self.fname) as fin:
while True:
text = rest + fin.read(8192) # avoid loading the entire file (=1 line) into RAM
if text == rest: # EOF
sentence.extend(rest.split()) # return the last chunk of words, too (may be shorter/longer)
if sentence:
yield sentence
break
last_token = text.rfind(' ') # the last token may have been split in two... keep it for the next iteration
words, rest = (text[:last_token].split(), text[last_token:].strip()) if last_token >= 0 else ([], text)
sentence.extend(words)
while len(sentence) >= max_sentence_length:
yield sentence[:max_sentence_length]
sentence = sentence[max_sentence_length:]
class Holmes(object):
"""Iterate over sentences from the Holmes training dataset ."""
def __init__(self, dir):
self.dir = dir
def __iter__(self):
sentence, rest, max_sentence_length = [], '', 1000
for i in sorted(os.listdir(self.dir)):
with open(os.path.join(self.dir,i)) as fin:
while True:
text = rest + fin.read(8192) # avoid loading the entire file (=1 line) into RAM
if text == rest: # EOF
sentence.extend(rest.split()) # return the last chunk of words, too (may be shorter/longer)
if sentence:
yield sentence
break
last_token = text.rfind(' ') # the last token may have been split in two... keep it for the next iteration
words, rest = (text[:last_token].split(), text[last_token:].strip()) if last_token >= 0 else ([], text)
sentence.extend(words)
while len(sentence) >= max_sentence_length:
yield sentence[:max_sentence_length]
sentence = sentence[max_sentence_length:]
class LineSentence(object):
def __init__(self, fname):
"""Simple format: one sentence = one line; words already preprocessed and separated by whitespace."""
self.fname = fname
def __iter__(self):
for line in open(self.fname):
yield line.split()
# Example: ./word2vec.py ~/workspace/word2vec/text8 ~/workspace/word2vec/questions-words.txt ./text8
if __name__ == "__main__":
logging.basicConfig(format='%(asctime)s : %(threadName)s : %(levelname)s : %(message)s', level=logging.INFO)
logging.info("running %s" % " ".join(sys.argv))
logging.info("using optimization %s" % FAST_VERSION)
# check and process cmdline input
program = os.path.basename(sys.argv[0])
if len(sys.argv) < 2:
print globals()['__doc__'] % locals()
sys.exit(1)
infile = sys.argv[1]
from gensim.models.word2vec_mod import Word2Vec # avoid referencing __main__ in pickle
seterr(all='raise') # don't ignore numpy errors
# model = Word2Vec(LineSentence(infile), size=200, min_count=5, workers=4)
model = Word2Vec(Text8Corpus(infile), size=200, min_count=5, workers=1)
if len(sys.argv) > 3:
outfile = sys.argv[3]
model.save(outfile + '.model')
model.save_word2vec_format(outfile + '.model.bin', binary=True)
model.save_word2vec_format(outfile + '.model.txt', binary=False)
if len(sys.argv) > 2:
questions_file = sys.argv[2]
model.accuracy(sys.argv[2])
logging.info("finished running %s" % program)