Word2Vec.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import collections\n",
    "import os\n",
    "import random\n",
    "import urllib\n",
    "import zipfile\n",
    "\n",
    "import numpy as np\n",
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "learning_rate = 0.1\n",
    "batch_size = 128\n",
    "num_steps = 3000000\n",
    "display_step = 10000\n",
    "eval_step = 200000\n",
    "\n",
    "# 词向量评估词.\n",
    "eval_words = ['five', 'of', 'going', 'hardware', 'american', 'britain']\n",
    "\n",
    "# Word2Vec Parameters.\n",
    "embedding_size = 200 # 词向量的维度，暂定200维\n",
    "max_vocabulary_size = 50000 # 总词数\n",
    "min_occurrence = 10 # 最小出现次数，若语料库出现少于十次则不要\n",
    "skip_window = 3 # 前后窗口大小，考虑前后的三个单词\n",
    "num_skips = 2 # How many times to reuse an input to generate a label.\n",
    "num_sampled = 64 # Number of negative examples to sample."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 下载数据，已手动下载\n",
    "# Download a small chunk of Wikipedia articles collection.\n",
    "url = 'http://mattmahoney.net/dc/text8.zip'\n",
    "data_path = 'text8.zip'\n",
    "if not os.path.exists(data_path):\n",
    "    print(\"Downloading the dataset... (It may take some time)\")\n",
    "    filename, _ = urllib.urlretrieve(url, data_path)\n",
    "    print(\"Done!\")\n",
    "# Unzip the dataset file. Text has already been processed.\n",
    "with zipfile.ZipFile(data_path) as f:\n",
    "    text_words = f.read(f.namelist()[0]).lower().split()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 转换编码格式\n",
    "test_word = []\n",
    "for word in text_words:\n",
    "    test_word.append(word.decode('utf8'))\n",
    "text_words = test_word"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 开始创造字典，先定义过少词的位置编码\n",
    "\n",
    "count = [('UNK', -1)]\n",
    "# 统计每个单词出现的次数，并且取\n",
    "# collections.Counter统计每个词的熟练，most_common 取前一百\n",
    "count.extend(collections.Counter(text_words).most_common(max_vocabulary_size - 1))\n",
    "for i in range(len(count) - 1, -1, -1):\n",
    "    if count[i][1] < min_occurrence:\n",
    "        count.pop(i)\n",
    "    else:\n",
    "        # The collection is ordered, so stop when 'min_occurrence' is reached.\n",
    "        break\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "count = [('UNK', -1)]\n",
    "# 统计每个单词出现的次数，并且取\n",
    "# collections.Counter统计每个词的熟练，most_common 取前一百\n",
    "count.extend(collections.Counter(text_words).most_common(max_vocabulary_size - 1))\n",
    "\n",
    "for i in range(len(count) - 1, -1, -1):\n",
    "    # 检验50K个单词中是否有次数少于最低次数的单词\n",
    "    if count[i][1] < min_occurrence:\n",
    "        count.pop(i)\n",
    "    else:\n",
    "        # The collection is ordered, so stop when 'min_occurrence' is reached.\n",
    "        break\n",
    "        \n",
    "# Compute the vocabulary size.\n",
    "vocabulary_size = len(count)\n",
    "# Assign an id to each word.\n",
    "word2id = dict()\n",
    "# enumerate 组合函数， 数据格式为[(1,\"a\")]\n",
    "for i, (word, _)in enumerate(count):\n",
    "    word2id[word] = i\n",
    "\n",
    "# 构造索引表\n",
    "data = list()\n",
    "unk_count = 0\n",
    "for word in text_words:\n",
    "    # Retrieve a word id, or assign it index 0 ('UNK') if not in dictionary.\n",
    "    index = word2id.get(word, 0)\n",
    "    if index == 0:\n",
    "        unk_count += 1\n",
    "    data.append(index)\n",
    "count[0] = ('UNK', unk_count)\n",
    "id2word = dict(zip(word2id.values(), word2id.keys()))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Words count: 17005207\n",
      "Unique words: 253854\n",
      "Vocabulary size: 47135\n",
      "Most common words: [('UNK', 444176), ('the', 1061396), ('of', 593677), ('and', 416629), ('one', 411764), ('in', 372201), ('a', 325873), ('to', 316376), ('zero', 264975), ('nine', 250430)]\n"
     ]
    }
   ],
   "source": [
    "print(\"Words count:\", len(text_words))\n",
    "print(\"Unique words:\", len(set(text_words)))\n",
    "print(\"Vocabulary size:\", vocabulary_size)\n",
    "print(\"Most common words:\", count[:10])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 计算逻辑，我没看懂\n",
    "data_index = 0\n",
    "# Generate training batch for the skip-gram model.\n",
    "def next_batch(batch_size, num_skips, skip_window):\n",
    "    global data_index\n",
    "    assert batch_size % num_skips == 0\n",
    "    assert num_skips <= 2 * skip_window\n",
    "    batch = np.ndarray(shape=(batch_size), dtype=np.int32)\n",
    "    labels = np.ndarray(shape=(batch_size, 1), dtype=np.int32)\n",
    "    # get window size (words left and right + current one).\n",
    "    span = 2 * skip_window + 1\n",
    "    buffer = collections.deque(maxlen=span)\n",
    "    if data_index + span > len(data):\n",
    "        data_index = 0\n",
    "    buffer.extend(data[data_index:data_index + span])\n",
    "    data_index += span\n",
    "    for i in range(batch_size // num_skips):\n",
    "        context_words = [w for w in range(span) if w != skip_window]\n",
    "        words_to_use = random.sample(context_words, num_skips)\n",
    "        for j, context_word in enumerate(words_to_use):\n",
    "            batch[i * num_skips + j] = buffer[skip_window]\n",
    "            labels[i * num_skips + j, 0] = buffer[context_word]\n",
    "        if data_index == len(data):\n",
    "            buffer.extend(data[0:span])\n",
    "            data_index = span\n",
    "        else:\n",
    "            buffer.append(data[data_index])\n",
    "            data_index += 1\n",
    "    # Backtrack a little bit to avoid skipping words in the end of a batch.\n",
    "    data_index = (data_index + len(data) - span) % len(data)\n",
    "    return batch, labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Ensure the following ops & var are assigned on CPU\n",
    "# (some ops are not compatible on GPU).\n",
    "with tf.device('/cpu:0'):\n",
    "    # Create the embedding variable (each row represent a word embedding vector).\n",
    "    # 随机初始化 参数为0-1之间\n",
    "    embedding = tf.Variable(tf.random.normal([vocabulary_size, embedding_size]))\n",
    "    # Construct the variables for the NCE loss.\n",
    "    nce_weights = tf.Variable(tf.random.normal([vocabulary_size, embedding_size]))\n",
    "    nce_biases = tf.Variable(tf.zeros([vocabulary_size]))\n",
    "\n",
    "def get_embedding(x):\n",
    "    with tf.device('/cpu:0'):\n",
    "        # Lookup the corresponding embedding vectors for each sample in X.\n",
    "        x_embed = tf.nn.embedding_lookup(embedding, x)\n",
    "        return x_embed\n",
    "\n",
    "def nce_loss(x_embed, y):\n",
    "    with tf.device('/cpu:0'):\n",
    "        # Compute the average NCE loss for the batch.\n",
    "        y = tf.cast(y, tf.int64)\n",
    "        loss = tf.reduce_mean(\n",
    "            tf.nn.nce_loss(weights=nce_weights,\n",
    "                           biases=nce_biases,\n",
    "                           labels=y,\n",
    "                           inputs=x_embed,\n",
    "                           num_sampled=num_sampled,\n",
    "                           num_classes=vocabulary_size))\n",
    "        return loss\n",
    "\n",
    "# Evaluation.\n",
    "def evaluate(x_embed):\n",
    "    with tf.device('/cpu:0'):\n",
    "        # Compute the cosine similarity between input data embedding and every embedding vectors\n",
    "        x_embed = tf.cast(x_embed, tf.float32)\n",
    "        x_embed_norm = x_embed / tf.sqrt(tf.reduce_sum(tf.square(x_embed)))\n",
    "        embedding_norm = embedding / tf.sqrt(tf.reduce_sum(tf.square(embedding), 1, keepdims=True), tf.float32)\n",
    "        cosine_sim_op = tf.matmul(x_embed_norm, embedding_norm, transpose_b=True)\n",
    "        return cosine_sim_op\n",
    "\n",
    "# Define the optimizer.\n",
    "optimizer = tf.optimizers.SGD(learning_rate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Optimization process. \n",
    "def run_optimization(x, y):\n",
    "    with tf.device('/cpu:0'):\n",
    "        # Wrap computation inside a GradientTape for automatic differentiation.\n",
    "        with tf.GradientTape() as g:\n",
    "            emb = get_embedding(x)\n",
    "            loss = nce_loss(emb, y)\n",
    "\n",
    "        # Compute gradients.\n",
    "        gradients = g.gradient(loss, [embedding, nce_weights, nce_biases])\n",
    "\n",
    "        # Update W and b following gradients.\n",
    "        optimizer.apply_gradients(zip(gradients, [embedding, nce_weights, nce_biases]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "step: 1, loss: 106.166237\n",
      "Evaluation...\n",
      "\"five\" nearest neighbors: four, three, six, two, seven, with, at, by,\n",
      "\"of\" nearest neighbors: the, and, to, is, a, in, by, that,\n",
      "\"going\" nearest neighbors: about, most, with, natural, from, time, many, including,\n",
      "\"hardware\" nearest neighbors: volcanic, anglia, pes, laboratory, andy, indicated, succeeded, wenceslaus,\n",
      "\"american\" nearest neighbors: not, which, with, from, have, on, or, nine,\n",
      "\"britain\" nearest neighbors: than, about, against, new, other, more, european, such,\n",
      "step: 10000, loss: 54.022804\n",
      "step: 20000, loss: 58.068230\n",
      "step: 30000, loss: 60.818146\n",
      "step: 40000, loss: 37.207851\n",
      "step: 50000, loss: 25.253174\n",
      "step: 60000, loss: 30.059118\n",
      "step: 70000, loss: 15.227887\n",
      "step: 80000, loss: 24.829811\n",
      "step: 90000, loss: 11.404577\n",
      "step: 100000, loss: 19.067738\n",
      "step: 110000, loss: 19.203407\n",
      "step: 120000, loss: 19.266972\n",
      "step: 130000, loss: 22.038090\n",
      "step: 140000, loss: 15.218485\n",
      "step: 150000, loss: 19.419878\n",
      "step: 160000, loss: 11.832603\n",
      "step: 170000, loss: 9.694306\n",
      "step: 180000, loss: 13.145454\n",
      "step: 190000, loss: 10.242126\n",
      "step: 200000, loss: 17.161320\n",
      "Evaluation...\n",
      "\"five\" nearest neighbors: four, three, six, seven, eight, nine, one, two,\n",
      "\"of\" nearest neighbors: and, the, UNK, by, in, to, also, on,\n",
      "\"going\" nearest neighbors: also, them, such, government, been, there, against, which,\n",
      "\"hardware\" nearest neighbors: may, between, out, see, who, people, over, from,\n",
      "\"american\" nearest neighbors: in, s, d, after, first, on, by, over,\n",
      "\"britain\" nearest neighbors: called, on, these, by, which, an, see, as,\n",
      "step: 210000, loss: 9.658041\n",
      "step: 220000, loss: 11.079557\n",
      "step: 230000, loss: 16.474039\n",
      "step: 240000, loss: 13.149895\n",
      "step: 250000, loss: 9.017632\n",
      "step: 260000, loss: 11.744461\n",
      "step: 270000, loss: 13.206314\n",
      "step: 280000, loss: 9.824510\n",
      "step: 290000, loss: 12.028326\n",
      "step: 300000, loss: 13.147823\n",
      "step: 310000, loss: 10.856438\n",
      "step: 320000, loss: 9.291737\n",
      "step: 330000, loss: 5.578553\n",
      "step: 340000, loss: 10.291031\n",
      "step: 350000, loss: 8.780314\n",
      "step: 360000, loss: 8.299170\n",
      "step: 370000, loss: 5.889017\n",
      "step: 380000, loss: 8.957944\n",
      "step: 390000, loss: 7.310172\n",
      "step: 400000, loss: 8.507822\n",
      "Evaluation...\n",
      "\"five\" nearest neighbors: six, three, seven, four, eight, two, one, nine,\n",
      "\"of\" nearest neighbors: and, the, in, while, with, on, were, was,\n",
      "\"going\" nearest neighbors: political, them, would, they, more, very, those, also,\n",
      "\"hardware\" nearest neighbors: may, see, who, e, out, de, that, through,\n",
      "\"american\" nearest neighbors: d, b, john, in, british, s, english, since,\n",
      "\"britain\" nearest neighbors: history, see, by, UNK, during, under, was, including,\n"
     ]
    }
   ],
   "source": [
    "# Words for testing.\n",
    "x_test = np.array([word2id[w] for w in eval_words])\n",
    "\n",
    "# Run training for the given number of steps.\n",
    "for step in range(1, num_steps + 1):\n",
    "    batch_x, batch_y = next_batch(batch_size, num_skips, skip_window)\n",
    "    run_optimization(batch_x, batch_y)\n",
    "    \n",
    "    if step % display_step == 0 or step == 1:\n",
    "        loss = nce_loss(get_embedding(batch_x), batch_y)\n",
    "        print(\"step: %i, loss: %f\" % (step, loss))\n",
    "        \n",
    "    # Evaluation.\n",
    "    if step % eval_step == 0 or step == 1:\n",
    "        print(\"Evaluation...\")\n",
    "        sim = evaluate(get_embedding(x_test)).numpy()\n",
    "        for i in range(len(eval_words)):\n",
    "            top_k = 8  # number of nearest neighbors.\n",
    "            nearest = (-sim[i, :]).argsort()[1:top_k + 1]\n",
    "            log_str = '\"%s\" nearest neighbors:' % eval_words[i]\n",
    "            for k in range(top_k):\n",
    "                log_str = '%s %s,' % (log_str, id2word[nearest[k]])\n",
    "            print(log_str)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_optimization(x, y):\n",
    "\n",
    "    with tf.GradientTape() as g:\n",
    "        emb = get_embedding(x)\n",
    "        loss = nce_loss(emb, y)\n",
    "\n",
    "    # Compute gradients.\n",
    "    gradients = g.gradient(loss, [embedding, nce_weights, nce_biases])\n",
    "\n",
    "    # Update W and b following gradients.\n",
    "    optimizer.apply_gradients(zip(gradients, [embedding, nce_weights, nce_biases]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "wjx_env",
   "language": "python",
   "name": "wjx_env"
  },
  "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}