Word2vec with custom sampling

This is a Python implementation of the word2vec algorithm.

The major part of this code was written by Radim Rehurek and was extracted from the gensim topic modelling library. I added code that abstracts the sampling step of the algorithm, making it more convenient to try out new sampling strategies.

When initializing the Word2Vec class, you now need to specify a sampling function. Given a sentence, it should return a list of pairs of vocabulary items (Vocab). For instance:

model = Word2Vec(Text8Corpus(infile), size=200, min_count=5, workers=1, sampler=sampling_fn)

Although version is not quite as fast as the original (it gives ~50% throughput) it still beats a pure Python implementation by a factor of at least 50.

Demo

run python word2vec.py data/text8-small data/questions-words.txt

Example: The original sampling strategy implemented as a function

defsample_word2vec(sentence, window=5):
pairs= []
foriinrange(len(sentence)):
word=sentence[i]
ifwordisNone:
continuereduced_window=random.randint(window)
j=i-window+reduced_windowifj<0:
j=0k=i+window+1-reduced_windowifk>len(sentence):
k=len(sentence)
forjinrange(j, k):
ifj==iorsentence[j] isNone:
continuepairs.append((sentence[i], sentence[j]))
returnpairs