wordModel.py

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import math
import os
import sys
import argparse
import random
import zipfile
import pickle

import numpy as np
from six.moves import urllib, xrange
import tensorflow as tf

from tensorflow.contrib.tensorboard.plugins import projector

def maybe_download(url, filename, expected_bytes):
	"""Download a file if not present, and make sure it's the right size."""
	local_filename = os.path.join(path, filename)
	if not os.path.exists(local_filename):
		local_filename, _ = urllib.request.urlretrieve(url + filename, local_filename)
	statinfo = os.stat(local_filename)
	if statinfo.st_size == expected_bytes:
		print('Found and verified', filename)
	else:
		print(statinfo.st_size)
		raise Exception('Failed to verify ' + local_filename +
							'. Can you get to it with a browser?')
	return local_filename

# Read the data into a list of strings.
def read_data(filename):
	"""Extract the first file enclosed in a zip file as a list of words."""
	with zipfile.ZipFile(filename) as f:
		data = tf.compat.as_str(f.read(f.namelist()[0])).split()
	return data


def build_dataset(words, n_words):
	"""Process raw inputs into a dataset."""
	count = [['UNK', -1]]
	count.extend(collections.Counter(words).most_common(n_words - 1))
	dictionary = dict()
	for word, _ in count:
		dictionary[word] = len(dictionary)
	data = list()
	unk_count = 0
	for word in words:
		index = dictionary.get(word, 0)
		if index == 0:	# dictionary['UNK']
			unk_count += 1
		data.append(index)
	count[0][1] = unk_count
	reversed_dictionary = dict(zip(dictionary.values(), dictionary.keys()))
	return data, count, dictionary, reversed_dictionary

def generate_batch(batch_size, num_skips, skip_window):
	global data_index
	assert batch_size % num_skips == 0
	assert num_skips <= 2 * skip_window
	batch = np.ndarray(shape=(batch_size), dtype=np.int32)
	labels = np.ndarray(shape=(batch_size, 1), dtype=np.int32)
	span = 2 * skip_window + 1	# [ skip_window target skip_window ]
	buffer = collections.deque(maxlen=span)	# pylint: disable=redefined-builtin
	if data_index + span > len(data):
		data_index = 0
	buffer.extend(data[data_index:data_index + span])
	data_index += span
	for i in range(batch_size // num_skips):
		context_words = [w for w in range(span) if w != skip_window]
		words_to_use = random.sample(context_words, num_skips)
		for j, context_word in enumerate(words_to_use):
			batch[i * num_skips + j] = buffer[skip_window]
			labels[i * num_skips + j, 0] = buffer[context_word]
		if data_index == len(data):
			buffer.extend(data[0:span])
			data_index = span
		else:
			buffer.append(data[data_index])
			data_index += 1
	# Backtrack a little bit to avoid skipping words in the end of a batch
	data_index = (data_index + len(data) - span) % len(data)
	return batch, labels

def makeWordModel():
	# Dowload files
	url = 'http://mattmahoney.net/dc/'
	filename = maybe_download(url,'text8.zip', 31344016)

	vocabulary = read_data(filename)
	print('Data size', len(vocabulary))

	# Build the dictionary and replace rare words with UNK token.
	vocabulary_size = 50000

	# Filling 4 global variables:
	# data - list of codes (integers from 0 to vocabulary_size-1). 
	#	This is the original text but words are replaced by their codes
	# count - map of words(strings) to count of occurrences
	# dictionary - map of words(strings) to their codes(integers)
	# reverse_dictionary - maps codes(integers) to words(strings)
	global data
	global dictionary
	global reverse_dictionary
	data, count, dictionary, reverse_dictionary = build_dataset(
		vocabulary, vocabulary_size)
	del vocabulary	# Hint to reduce memory.
	print('Most common words (+UNK)', count[:5])
	print('Sample data', data[:10], [reverse_dictionary[i] for i in data[:10]])

	global data_index
	data_index=0
	batch, labels = generate_batch(batch_size=8, num_skips=2, skip_window=1)
	for i in range(8):
		print(batch[i], reverse_dictionary[batch[i]], '->', labels[i, 0],
			reverse_dictionary[labels[i, 0]])

	# Build and train a skip-gram model.

	batch_size = 128
	embedding_size = 128	# Dimension of the embedding vector.
	skip_window = 1	# How many words to consider left and right.
	num_skips = 2	# How many times to reuse an input to generate a label.
	num_sampled = 64	# Number of negative examples to sample.

	# We pick a random validation set to sample nearest neighbors. Here we limit the
	# validation samples to the words that have a low numeric ID, which by
	# construction are also the most frequent. These 3 variables are used only for
	# displaying model accuracy, they don't affect calculation.
	valid_size = 16	# Random set of words to evaluate similarity on.
	valid_window = 100	# Only pick dev samples in the head of the distribution.
	valid_examples = np.random.choice(valid_window, valid_size, replace=False)

	graph = tf.Graph()

	with graph.as_default():
		# Input data.
		with tf.name_scope('inputs'):
			train_inputs = tf.placeholder(tf.int32, shape=[batch_size])
			train_labels = tf.placeholder(tf.int32, shape=[batch_size, 1])
			valid_dataset = tf.constant(valid_examples, dtype=tf.int32)

		# Ops and variables pinned to the CPU because of missing GPU implementation
		with tf.device('/cpu:0'):
			# Look up embeddings for inputs.
			with tf.name_scope('embeddings'):
				embeddings = tf.Variable(
					tf.random_uniform([vocabulary_size, embedding_size], -1.0, 1.0))
				embed = tf.nn.embedding_lookup(embeddings, train_inputs)

			# Construct the variables for the NCE loss
			with tf.name_scope('weights'):
				nce_weights = tf.Variable(
					tf.truncated_normal(
						[vocabulary_size, embedding_size],
						stddev=1.0 / math.sqrt(embedding_size)))
			with tf.name_scope('biases'):
				nce_biases = tf.Variable(tf.zeros([vocabulary_size]))

		# Compute the average NCE loss for the batch.
		# tf.nce_loss automatically draws a new sample of the negative labels each
		# time we evaluate the loss.
		# Explanation of the meaning of NCE loss:
		#	http://mccormickml.com/2016/04/19/word2vec-tutorial-the-skip-gram-models/nearest/
		with tf.name_scope('loss'):
			loss = tf.reduce_mean(
				tf.nn.nce_loss(
					weights=nce_weights,
					biases=nce_biases,
					labels=train_labels,
					inputs=embed,
					num_sampled=num_sampled,
					num_classes=vocabulary_size))

		# Add the loss value as a scalar to summary.
		tf.summary.scalar('loss', loss)

		# Construct the SGD optimizer using a learning rate of 1.0.
		with tf.name_scope('optimizer'):
			optimizer = tf.train.GradientDescentOptimizer(1.0).minimize(loss)

		# Compute the cosine similarity between minibatch examples and all embeddings.
		norm = tf.sqrt(tf.reduce_sum(tf.square(embeddings), 1, keepdims=True))
		normalized_embeddings = embeddings / norm
		valid_embeddings = tf.nn.embedding_lookup(normalized_embeddings,valid_dataset)
		similarity = tf.matmul(valid_embeddings, normalized_embeddings, transpose_b=True)

		# Merge all summaries.
		merged = tf.summary.merge_all()

		# Add variable initializer.
		init = tf.global_variables_initializer()

		# Create a saver.
		saver = tf.train.Saver()
	# Begin training.
	num_steps = 100001

	with tf.Session(graph=graph) as session:
		# We must initialize all variables before we use them.
		init.run()
		print('Initialized')

		average_loss = 0
		for step in xrange(num_steps):
			batch_inputs, batch_labels = generate_batch(batch_size, num_skips, skip_window)
			feed_dict = {train_inputs: batch_inputs, train_labels: batch_labels}

			# Define metadata variable.
			run_metadata = tf.RunMetadata()

			# We perform one update step by evaluating the optimizer op (including it
			# in the list of returned values for session.run()
			# Also, evaluate the merged op to get all summaries from the returned "summary" variable.
			# Feed metadata variable to session for visualizing the graph in TensorBoard.
			_, summary, loss_val = session.run(
					[optimizer, merged, loss],
					feed_dict=feed_dict,
					run_metadata=run_metadata)
			average_loss += loss_val

			if step % 2000 == 0:
				if step > 0:
					average_loss /= 2000
				# The average loss is an estimate of the loss over the last 2000 batches.
				print('Average loss at step ', step, ': ', average_loss)
				average_loss = 0
			# Note that this is expensive (~20% slowdown if computed every 500 steps)
			if step % 10000 == 0:
				sim = similarity.eval()
				for i in xrange(valid_size):
					valid_word = reverse_dictionary[valid_examples[i]]
					top_k = 8	# number of nearest neighbors
					nearest = (-sim[i,:]).argsort()[1:top_k + 1]
					log_str = 'Nearest to %s:' % valid_word
					for k in xrange(top_k):
						close_word = reverse_dictionary[nearest[k]]
						log_str = '%s %s,' % (log_str, close_word)
					print(log_str)
		save_obj(dictionary,'metadata')
		save_obj(normalized_embeddings.eval(),'theGrail')
		saver.save(session, os.path.join(path,'wordModel'))

def save_obj(obj, name):
    with open('models/nearest/'+ name + '.pkl', 'wb') as f:
        pickle.dump(obj, f, pickle.HIGHEST_PROTOCOL)

def load_obj(name):
    with open('models/nearest/' + name + '.pkl', 'rb') as f:
        return pickle.load(f)

def loadNearestModel():
	global path
	path = os.path.join(os.path.dirname(os.path.realpath(sys.argv[0])),'models/nearest')

	if not os.path.exists(os.path.join(path,'wordModel.index')):
		makeWordModel()
	dictionary = load_obj('metadata')
	reverse_dictionary = dict(zip(dictionary.values(), dictionary.keys()))

	vecs = load_obj('theGrail')
	return vecs, dictionary, reverse_dictionary

def searchSimilar(model, word, kn):
	#model 0 -> vecs, 1 -> dictionary, 2 -> reverseDictionary
	if word not in model[1]:
		return None
	j=model[1][word]
	a = np.array(np.matmul(model[0][j,:],np.transpose(model[0])))
	nearest = (-a).argsort()[1:kn + 1]
	close_word=list()
	for k in nearest:
		if a[k]<0.615:
			break
		close_word.append(model[2][k])
	if close_word==[]: return None
	return close_word

if __name__ == '__main__':
	model = loadNearestModel()
	print(searchSimilar(model,'is',5))