word2vec.py

# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Basic word2vec example."""

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

import collections
import math
import os
import random
from tempfile import gettempdir
import json

import numpy as np
from six.moves import xrange  # pylint: disable=redefined-builtin
import tensorflow as tf

import matplotlib as mpl

# mpl.rcParams['font.sans-serif'] = ['SimHei']
# mpl.rcParams['font.serif'] = ['SimHei']
# import seaborn as sns
# sns.set_style("darkgrid", {"font.sans-serif": ['simhei', 'Arial']})

# Read the data into a list of strings.
with open('QuanSongci.txt', 'r', encoding='utf-8') as f:
    vocabulary = f.read()
print('Data size', len(vocabulary))

# Step 2: Build the dictionary and replace rare words with UNK token.
vocabulary_size = 5000


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


# 生成四个全局变量:
# data - 用字符所对应的序号来代替原有的文本.
# count - 字符出现次数的一个列表映射
# dictionary - 字符与其对应的索引号的映射
# reverse_dictionary - 字符与其对应的索引号的逆映射
data, count, dictionary, reverse_dictionary = build_dataset(vocabulary,
                                                            vocabulary_size)
print('--------save dictionary-----------------')
with open('E:\\ailearning\\word_writer\\dictionary_new.json', 'w', encoding='utf-8') as f:
    json.dump(dictionary, f)
with open('E:\\ailearning\\word_writer\\reversed_dictionary.json', 'w', encoding='utf-8') as f:
    json.dump(reverse_dictionary, f)
with open('E:\\ailearning\\word_writer\\data.json', 'w', encoding='utf-8') as f:
    json.dump(data, f)


del vocabulary  # Hint to reduce memory.
print('Top 20 frequency words (+UNK)', count[:20])
print('Sample data', data[:10], [reverse_dictionary[i] for i in data[:10]])

data_index = 0


# Step 3: Function to generate a training batch for the skip-gram model.
# 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)
#     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):
#         #   print('pause!')
#         if data_index == len(data):
#             buffer.clear()
#             buffer.extend(data[: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
#
#
# batch, labels = generate_batch(batch_size=16, num_skips=2, skip_window=1)
# for i in range(16):
#     print(batch[i], reverse_dictionary[batch[i]],
#           '->', labels[i, 0], reverse_dictionary[labels[i, 0]])
#
# # Step 4: 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.
#     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.
#         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
#         nce_weights = tf.Variable(
#             tf.truncated_normal([vocabulary_size, embedding_size],
#                                 stddev=1.0 / math.sqrt(embedding_size)))
#         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-model/
#     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))
#
#     # Construct the SGD optimizer using a learning rate of 1.0.
#     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, keep_dims=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)
#
#     # Add variable initializer.
#     init = tf.global_variables_initializer()
#
# # Step 5: Begin training.
# num_steps = 800001
#
# 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 range(num_steps):
#         batch_inputs, batch_labels = generate_batch(
#             batch_size, num_skips, skip_window)
#         feed_dict = {train_inputs: batch_inputs, train_labels: batch_labels}
#
#         # We perform one update step by evaluating the optimizer op (including it
#         # in the list of returned values for session.run()
#         _, loss_val = session.run([optimizer, loss], feed_dict=feed_dict)
#         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 range(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 range(top_k):
#                     close_word = reverse_dictionary[nearest[k]]
#                     log_str = '%s %s,' % (log_str, close_word)
#                 print(log_str)
#     final_embeddings = normalized_embeddings.eval()
#     np.save('embedding.npy', final_embeddings)
#
#
# # Step 6: Visualize the embeddings.
#
#
# # pylint: disable=missing-docstring
# # Function to draw visualization of distance between embeddings.
# def plot_with_labels(low_dim_embs, labels, filename):
#     assert low_dim_embs.shape[0] >= len(labels), 'More labels than embeddings'
#     plt.figure(figsize=(18, 18))  # in inches
#     for i, label in enumerate(labels):
#         x, y = low_dim_embs[i, :]
#         plt.scatter(x, y)
#         plt.annotate(label,
#                      xy=(x, y),
#                      xytext=(5, 2),
#                      textcoords='offset points',
#                      ha='right',
#                      va='bottom')
#
#     plt.savefig(filename)
#
#
# try:
#     # pylint: disable=g-import-not-at-top
#     from sklearn.manifold import TSNE
#     import matplotlib.pyplot as plt
#
#     print('begin to plot picture~')
#     tsne = TSNE(perplexity=30, n_components=2, init='pca', n_iter=5000, method='exact')
#     plot_only = 500
#     low_dim_embs = tsne.fit_transform(final_embeddings[:plot_only, :])
#     labels = [reverse_dictionary[i] for i in range(plot_only)]
#     plot_with_labels(low_dim_embs, labels, os.path.join(gettempdir(), 'E:\\ailearning\\word_writer\\tsne.png'))
#
# except ImportError as ex:
#     print('Please install sklearn, matplotlib, and scipy to show embeddings.')
#     print(ex)