skip_gram/Skip-Gram-English-Corpus.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Word2Vec之Skip-Gram模型\n",
    "\n",
    "下面代码将用TensorFlow实现Word2Vec中的Skip-Gram模型。\n",
    "\n",
    "关于Skip-Gram模型请参考上一篇[知乎专栏文章](https://zhuanlan.zhihu.com/p/27234078)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1 导入包"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import time\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import random\n",
    "from collections import Counter"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2 加载数据\n",
    "\n",
    "数据集使用的是来自Matt Mahoney的维基百科文章，数据集已经被清洗过，去除了特殊符号等，并不是全量数据，只是部分数据，所以实际上最后训练出的结果很一般（语料不够）。\n",
    "\n",
    "如果想获取更全的语料数据，可以访问以下网站，这是gensim中Word2Vec提供的语料：\n",
    "\n",
    "- 来自Matt Mahoney预处理后的[文本子集](http://mattmahoney.net/dc/enwik9.zip)，里面包含了10亿个字符。\n",
    "- 与第一条一样的经过预处理的[文本数据](http://dumps.wikimedia.org/enwiki/latest/enwiki-latest-pages-articles.xml.bz2)，但是包含了30个亿的字符。\n",
    "- 多种语言的[训练文本](http://www.statmt.org/wmt11/translation-task.html#download)。\n",
    "- [UMBC webbase corpus](http://ebiquity.umbc.edu/redirect/to/resource/id/351/UMBC-webbase-corpus)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with open('data/text8') as f:\n",
    "    text = f.read()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3 数据预处理\n",
    "\n",
    "数据预处理过程主要包括：\n",
    "\n",
    "- 替换文本中特殊符号并去除低频词\n",
    "- 对文本分词\n",
    "- 构建语料\n",
    "- 单词映射表"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# 定义函数来完成数据的预处理\n",
    "def preprocess(text, freq=5):\n",
    "    '''\n",
    "    对文本进行预处理\n",
    "    \n",
    "    参数\n",
    "    ---\n",
    "    text: 文本数据\n",
    "    freq: 词频阈值\n",
    "    '''\n",
    "    # 对文本中的符号进行替换\n",
    "    text = text.lower()\n",
    "    text = text.replace('.', ' <PERIOD> ')\n",
    "    text = text.replace(',', ' <COMMA> ')\n",
    "    text = text.replace('\"', ' <QUOTATION_MARK> ')\n",
    "    text = text.replace(';', ' <SEMICOLON> ')\n",
    "    text = text.replace('!', ' <EXCLAMATION_MARK> ')\n",
    "    text = text.replace('?', ' <QUESTION_MARK> ')\n",
    "    text = text.replace('(', ' <LEFT_PAREN> ')\n",
    "    text = text.replace(')', ' <RIGHT_PAREN> ')\n",
    "    text = text.replace('--', ' <HYPHENS> ')\n",
    "    text = text.replace('?', ' <QUESTION_MARK> ')\n",
    "    # text = text.replace('\\n', ' <NEW_LINE> ')\n",
    "    text = text.replace(':', ' <COLON> ')\n",
    "    words = text.split()\n",
    "    \n",
    "    # 删除低频词，减少噪音影响\n",
    "    word_counts = Counter(words)\n",
    "    trimmed_words = [word for word in words if word_counts[word] > freq]\n",
    "\n",
    "    return trimmed_words"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['anarchism', 'originated', 'as', 'a', 'term', 'of', 'abuse', 'first', 'used', 'against', 'early', 'working', 'class', 'radicals', 'including', 'the', 'diggers', 'of', 'the', 'english']\n"
     ]
    }
   ],
   "source": [
    "# 清洗文本并分词\n",
    "words = preprocess(text)\n",
    "print(words[:20])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# 构建映射表\n",
    "vocab = set(words)\n",
    "vocab_to_int = {w: c for c, w in enumerate(vocab)}\n",
    "int_to_vocab = {c: w for c, w in enumerate(vocab)}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total words: 16680599\n",
      "unique words: 63641\n"
     ]
    }
   ],
   "source": [
    "print(\"total words: {}\".format(len(words)))\n",
    "print(\"unique words: {}\".format(len(set(words))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# 对原文本进行vocab到int的转换\n",
    "int_words = [vocab_to_int[w] for w in words]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4 采样\n",
    "\n",
    "对停用词进行采样，例如“the”， “of”以及“for”这类单词进行剔除。剔除这些单词以后能够加快我们的训练过程，同时减少训练过程中的噪音。\n",
    "\n",
    "我们采用以下公式:\n",
    "$$ P(w_i) = 1 - \\sqrt{\\frac{t}{f(w_i)}} $$\n",
    "\n",
    "其中$ t $是一个阈值参数，一般为1e-3至1e-5。  \n",
    "$f(w_i)$ 是单词 $w_i$ 在整个数据集中的出现频次。  \n",
    "$P(w_i)$ 是单词被删除的概率。\n",
    "\n",
    ">这个公式和论文中描述的那个公式有一些不同"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "t = 1e-5 # t值\n",
    "threshold = 0.8 # 剔除概率阈值\n",
    "\n",
    "# 统计单词出现频次\n",
    "int_word_counts = Counter(int_words)\n",
    "total_count = len(int_words)\n",
    "# 计算单词频率\n",
    "word_freqs = {w: c/total_count for w, c in int_word_counts.items()}\n",
    "# 计算被删除的概率\n",
    "prob_drop = {w: 1 - np.sqrt(t / word_freqs[w]) for w in int_word_counts}\n",
    "# 对单词进行采样\n",
    "train_words = [w for w in int_words if prob_drop[w] < threshold]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "6925252"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(train_words)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# 从上面数据可以看到，我们本身有1670万的文本，经过采样后剩下693万。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 5 构造batch\n",
    "\n",
    "Skip-Gram模型是通过输入词来预测上下文。因此我们要构造我们的训练样本，具体思想请参考知乎专栏，这里不再重复。\n",
    "\n",
    "对于一个给定词，离它越近的词可能与它越相关，离它越远的词越不相关，这里我们设置窗口大小为5，对于每个训练单词，我们还会在[1:5]之间随机生成一个整数R，用R作为我们最终选择output word的窗口大小。这里之所以多加了一步随机数的窗口重新选择步骤，是为了能够让模型更聚焦于当前input word的邻近词。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def get_targets(words, idx, window_size=5):\n",
    "    '''\n",
    "    获得input word的上下文单词列表\n",
    "    \n",
    "    参数\n",
    "    ---\n",
    "    words: 单词列表\n",
    "    idx: input word的索引号\n",
    "    window_size: 窗口大小\n",
    "    '''\n",
    "    target_window = np.random.randint(1, window_size+1)\n",
    "    # 这里要考虑input word前面单词不够的情况\n",
    "    start_point = idx - target_window if (idx - target_window) > 0 else 0\n",
    "    end_point = idx + target_window\n",
    "    # output words(即窗口中的上下文单词)\n",
    "    targets = set(words[start_point: idx] + words[idx+1: end_point+1])\n",
    "    return list(targets)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def get_batches(words, batch_size, window_size=5):\n",
    "    '''\n",
    "    构造一个获取batch的生成器\n",
    "    '''\n",
    "    n_batches = len(words) // batch_size\n",
    "    \n",
    "    # 仅取full batches\n",
    "    words = words[:n_batches*batch_size]\n",
    "    \n",
    "    for idx in range(0, len(words), batch_size):\n",
    "        x, y = [], []\n",
    "        batch = words[idx: idx+batch_size]\n",
    "        for i in range(len(batch)):\n",
    "            batch_x = batch[i]\n",
    "            batch_y = get_targets(batch, i, window_size)\n",
    "            # 由于一个input word会对应多个output word，因此需要长度统一\n",
    "            x.extend([batch_x]*len(batch_y))\n",
    "            y.extend(batch_y)\n",
    "        yield x, y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 6 构建网络\n",
    "\n",
    "该部分主要包括：\n",
    "\n",
    "- 输入层\n",
    "- Embedding\n",
    "- Negative Sampling"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 输入"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "train_graph = tf.Graph()\n",
    "with train_graph.as_default():\n",
    "    inputs = tf.placeholder(tf.int32, shape=[None], name='inputs')\n",
    "    labels = tf.placeholder(tf.int32, shape=[None, None], name='labels')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Embedding\n",
    "\n",
    "嵌入矩阵的矩阵形状为 $ vocab\\_size\\times hidden\\_units\\_size$ \n",
    "\n",
    "TensorFlow中的tf.nn.embedding_lookup函数可以实现lookup的计算方式"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "vocab_size = len(int_to_vocab)\n",
    "embedding_size = 200 # 嵌入维度"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with train_graph.as_default():\n",
    "    # 嵌入层权重矩阵\n",
    "    embedding = tf.Variable(tf.random_uniform([vocab_size, embedding_size], -1, 1))\n",
    "    # 实现lookup\n",
    "    embed = tf.nn.embedding_lookup(embedding, inputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Negative Sampling\n",
    "\n",
    "负采样主要是为了解决梯度下降计算速度慢的问题，详情同样参考我的上一篇知乎专栏文章。\n",
    "\n",
    "TensorFlow中的tf.nn.sampled_softmax_loss会在softmax层上进行采样计算损失，计算出的loss要比full softmax loss低。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "n_sampled = 100\n",
    "\n",
    "with train_graph.as_default():\n",
    "    softmax_w = tf.Variable(tf.truncated_normal([vocab_size, embedding_size], stddev=0.1))\n",
    "    softmax_b = tf.Variable(tf.zeros(vocab_size))\n",
    "    \n",
    "    # 计算negative sampling下的损失\n",
    "    loss = tf.nn.sampled_softmax_loss(softmax_w, softmax_b, labels, embed, n_sampled, vocab_size)\n",
    "    \n",
    "    cost = tf.reduce_mean(loss)\n",
    "    optimizer = tf.train.AdamOptimizer().minimize(cost)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 验证\n",
    "\n",
    "为了更加直观的看到我们训练的结果，我们将查看训练出的相近语义的词。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with train_graph.as_default():\n",
    "    # 随机挑选一些单词\n",
    "    valid_size = 16 \n",
    "    valid_window = 100\n",
    "    # 从不同位置各选8个单词\n",
    "    valid_examples = np.array(random.sample(range(valid_window), valid_size//2))\n",
    "    valid_examples = np.append(valid_examples, \n",
    "                               random.sample(range(1000,1000+valid_window), valid_size//2))\n",
    "    \n",
    "    valid_size = len(valid_examples)\n",
    "    # 验证单词集\n",
    "    valid_dataset = tf.constant(valid_examples, dtype=tf.int32)\n",
    "    \n",
    "    # 计算每个词向量的模并进行单位化\n",
    "    norm = tf.sqrt(tf.reduce_sum(tf.square(embedding), 1, keep_dims=True))\n",
    "    normalized_embedding = embedding / norm\n",
    "    # 查找验证单词的词向量\n",
    "    valid_embedding = tf.nn.embedding_lookup(normalized_embedding, valid_dataset)\n",
    "    # 计算余弦相似度\n",
    "    similarity = tf.matmul(valid_embedding, tf.transpose(normalized_embedding))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "epochs = 10 # 迭代轮数\n",
    "batch_size = 1000 # batch大小\n",
    "window_size = 10 # 窗口大小\n",
    "\n",
    "with train_graph.as_default():\n",
    "    saver = tf.train.Saver() # 文件存储\n",
    "\n",
    "with tf.Session(graph=train_graph) as sess:\n",
    "    iteration = 1\n",
    "    loss = 0\n",
    "    sess.run(tf.global_variables_initializer())\n",
    "\n",
    "    for e in range(1, epochs+1):\n",
    "        batches = get_batches(train_words, batch_size, window_size)\n",
    "        start = time.time()\n",
    "        # \n",
    "        for x, y in batches:\n",
    "            \n",
    "            feed = {inputs: x,\n",
    "                    labels: np.array(y)[:, None]}\n",
    "            train_loss, _ = sess.run([cost, optimizer], feed_dict=feed)\n",
    "            \n",
    "            loss += train_loss\n",
    "            \n",
    "            if iteration % 100 == 0: \n",
    "                end = time.time()\n",
    "                print(\"Epoch {}/{}\".format(e, epochs),\n",
    "                      \"Iteration: {}\".format(iteration),\n",
    "                      \"Avg. Training loss: {:.4f}\".format(loss/100),\n",
    "                      \"{:.4f} sec/batch\".format((end-start)/100))\n",
    "                loss = 0\n",
    "                start = time.time()\n",
    "            \n",
    "            # 计算相似的词\n",
    "            if iteration % 1000 == 0:\n",
    "                # 计算similarity\n",
    "                sim = similarity.eval()\n",
    "                for i in range(valid_size):\n",
    "                    valid_word = int_to_vocab[valid_examples[i]]\n",
    "                    top_k = 8 # 取最相似单词的前8个\n",
    "                    nearest = (-sim[i, :]).argsort()[1:top_k+1]\n",
    "                    log = 'Nearest to [%s]:' % valid_word\n",
    "                    for k in range(top_k):\n",
    "                        close_word = int_to_vocab[nearest[k]]\n",
    "                        log = '%s %s,' % (log, close_word)\n",
    "                    print(log)\n",
    "            \n",
    "            iteration += 1\n",
    "            \n",
    "    save_path = saver.save(sess, \"checkpoints/text8.ckpt\")\n",
    "    embed_mat = sess.run(normalized_embedding)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "%config InlineBackend.figure_format = 'retina'\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.manifold import TSNE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "viz_words = 500\n",
    "tsne = TSNE()\n",
    "embed_tsne = tsne.fit_transform(embed_mat[:viz_words, :])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(14, 14))\n",
    "for idx in range(viz_words):\n",
    "    plt.scatter(*embed_tsne[idx, :], color='steelblue')\n",
    "    plt.annotate(int_to_vocab[idx], (embed_tsne[idx, 0], embed_tsne[idx, 1]), alpha=0.7)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}