[Incredible Word2Vec] 6. Keras Version of Word2Vec

#! -*- coding:utf-8 -*-
# Keras version of Word2Vec, Author: Su Jianlin, http://kexue.fm
# Tested on Keras 2.0.6 + Tensorflow

import numpy as np
from keras.layers import Input, Embedding, Lambda
from keras.models import Model
import keras.backend as K

word_size = 128 # Word vector dimension
window = 5 # Window size
nb_negative = 16 # Number of samples for random negative sampling
min_count = 10 # Words with frequency less than min_count will be discarded
nb_worker = 4 # Concurrency for data reading
nb_epoch = 2 # Number of iterations; with Adam, 1-2 epochs yield good results
subsample_t = 1e-5 # Words with frequency > subsample_t will be downsampled
nb_sentence_per_batch = 20
# Currently using sentences as the batch unit. 
# Note: sample count is proportional to the number of characters.

import pymongo
class Sentences: # Corpus generator, must be written this way to be reusable
    def __init__(self):
        self.db = pymongo.MongoClient().weixin.text_articles
    def __iter__(self):
        for t in self.db.find(no_cursor_timeout=True).limit(100000):
            yield t['words'] # Returns tokenized results

sentences = Sentences()
words = {} # Word frequency table
nb_sentence = 0 # Total sentences
total = 0. # Total word frequency

for d in sentences:
    nb_sentence += 1
    for w in d:
        if w not in words:
            words[w] = 0
        words[w] += 1
        total += 1
    if nb_sentence % 10000 == 0:
        print u'Found %s articles'%nb_sentence

words = {i:j for i,j in words.items() if j >= min_count} # Truncate frequency
id2word = {i+1:j for i,j in enumerate(words)} # ID to word mapping, 0 is UNK
word2id = {j:i for i,j in id2word.items()} # Word to ID mapping
nb_word = len(words)+1 # Total words (including padding symbol 0)

subsamples = {i:j/total for i,j in words.items() if j/total > subsample_t}
# Downsampling formula following the original word2vec source code
subsamples = {i:subsample_t/j+(subsample_t/j)**0.5 for i,j in subsamples.items()}
subsamples = {word2id[i]:j for i,j in subsamples.items() if j < 1.} # Downsampling table

def data_generator(): # Training data generator
    while True:
        x,y = [],[]
        _ = 0
        for d in sentences:
            d = [0]*window + [word2id[w] for w in d if w in word2id] + [0]*window
            r = np.random.random(len(d))
            for i in range(window, len(d)-window):
                if d[i] in subsamples and r[i] > subsamples[d[i]]: 
                    continue
                x.append(d[i-window:i]+d[i+1:i+1+window])
                y.append([d[i]])
            _ += 1
            if _ == nb_sentence_per_batch:
                x,y = np.array(x),np.array(y)
                z = np.zeros((len(x), 1))
                yield [x,y],z
                x,y = [],[]
                _ = 0

# CBOW Input
input_words = Input(shape=(window*2,), dtype='int32')
input_vecs = Embedding(nb_word, word_size, name='word2vec')(input_words)
input_vecs_sum = Lambda(lambda x: K.sum(x, axis=1))(input_vecs) # Sum context vectors

# Construct random negative samples and combine with target for sampling
target_word = Input(shape=(1,), dtype='int32')
negatives = Lambda(lambda x: K.random_uniform((K.shape(x)[0], nb_negative), 0, nb_word, 'int32'))(target_word)
samples = Lambda(lambda x: K.concatenate(x))([target_word, negatives])

# Perform Dense and Softmax only within the samples
softmax_weights = Embedding(nb_word, word_size, name='W')(samples)
softmax_biases = Embedding(nb_word, 1, name='b')(samples)
softmax = Lambda(lambda x: 
                    K.softmax((K.batch_dot(x[0], K.expand_dims(x[1],2))+x[2])[:,:,0])
                )([softmax_weights, input_vecs_sum, softmax_biases]) 
# Parameters stored in Embedding layers; matrix multiplication via K backend

# Note: target is at index 0 in samples, so softmax target ID is always 0 (see z in generator)

model = Model(inputs=[input_words, target_word], outputs=softmax)
model.compile(loss='sparse_categorical_crossentropy', optimizer='adam', metrics=['accuracy'])

model.fit_generator(data_generator(), 
                    steps_per_epoch=nb_sentence/nb_sentence_per_batch, 
                    epochs=nb_epoch,
                    workers=nb_worker,
                    use_multiprocessing=True
                   )

model.save_weights('word2vec.model')

# Check word vector quality via similarity
embeddings = model.get_weights()[0]
normalized_embeddings = embeddings / (embeddings**2).sum(axis=1).reshape((-1,1))**0.5

def most_similar(w):
    v = normalized_embeddings[word2id[w]]
    sims = np.dot(normalized_embeddings, v)
    sort = sims.argsort()[::-1]
    sort = sort[sort > 0]
    return [(id2word[i], sims[i]) for i in sort[:10]]

import pandas as pd
pd.Series(most_similar(u'science'))