115 lines
4.5 KiB
Plaintext
115 lines
4.5 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "code",
|
|
"metadata": {},
|
|
"source": [
|
|
"# code by Tae Hwan Jung @graykode\n",
|
|
"import numpy as np\n",
|
|
"import torch\n",
|
|
"import torch.nn as nn\n",
|
|
"import torch.optim as optim\n",
|
|
"import matplotlib.pyplot as plt\n",
|
|
"\n",
|
|
"def random_batch():\n",
|
|
" random_inputs = []\n",
|
|
" random_labels = []\n",
|
|
" random_index = np.random.choice(range(len(skip_grams)), batch_size, replace=False)\n",
|
|
"\n",
|
|
" for i in random_index:\n",
|
|
" random_inputs.append(np.eye(voc_size)[skip_grams[i][0]]) # target\n",
|
|
" random_labels.append(skip_grams[i][1]) # context word\n",
|
|
"\n",
|
|
" return random_inputs, random_labels\n",
|
|
"\n",
|
|
"# Model\n",
|
|
"class Word2Vec(nn.Module):\n",
|
|
" def __init__(self):\n",
|
|
" super(Word2Vec, self).__init__()\n",
|
|
" # W and WT is not Traspose relationship\n",
|
|
" self.W = nn.Linear(voc_size, embedding_size, bias=False) # voc_size > embedding_size Weight\n",
|
|
" self.WT = nn.Linear(embedding_size, voc_size, bias=False) # embedding_size > voc_size Weight\n",
|
|
"\n",
|
|
" def forward(self, X):\n",
|
|
" # X : [batch_size, voc_size]\n",
|
|
" hidden_layer = self.W(X) # hidden_layer : [batch_size, embedding_size]\n",
|
|
" output_layer = self.WT(hidden_layer) # output_layer : [batch_size, voc_size]\n",
|
|
" return output_layer\n",
|
|
"\n",
|
|
"if __name__ == '__main__':\n",
|
|
" batch_size = 2 # mini-batch size\n",
|
|
" embedding_size = 2 # embedding size\n",
|
|
"\n",
|
|
" sentences = [\"apple banana fruit\", \"banana orange fruit\", \"orange banana fruit\",\n",
|
|
" \"dog cat animal\", \"cat monkey animal\", \"monkey dog animal\"]\n",
|
|
"\n",
|
|
" word_sequence = \" \".join(sentences).split()\n",
|
|
" word_list = \" \".join(sentences).split()\n",
|
|
" word_list = list(set(word_list))\n",
|
|
" word_dict = {w: i for i, w in enumerate(word_list)}\n",
|
|
" voc_size = len(word_list)\n",
|
|
"\n",
|
|
" # Make skip gram of one size window\n",
|
|
" skip_grams = []\n",
|
|
" for i in range(1, len(word_sequence) - 1):\n",
|
|
" target = word_dict[word_sequence[i]]\n",
|
|
" context = [word_dict[word_sequence[i - 1]], word_dict[word_sequence[i + 1]]]\n",
|
|
" for w in context:\n",
|
|
" skip_grams.append([target, w])\n",
|
|
"\n",
|
|
" model = Word2Vec()\n",
|
|
"\n",
|
|
" criterion = nn.CrossEntropyLoss()\n",
|
|
" optimizer = optim.Adam(model.parameters(), lr=0.001)\n",
|
|
"\n",
|
|
" # Training\n",
|
|
" for epoch in range(5000):\n",
|
|
" input_batch, target_batch = random_batch()\n",
|
|
" input_batch = torch.Tensor(input_batch)\n",
|
|
" target_batch = torch.LongTensor(target_batch)\n",
|
|
"\n",
|
|
" optimizer.zero_grad()\n",
|
|
" output = model(input_batch)\n",
|
|
"\n",
|
|
" # output : [batch_size, voc_size], target_batch : [batch_size] (LongTensor, not one-hot)\n",
|
|
" loss = criterion(output, target_batch)\n",
|
|
" if (epoch + 1) % 1000 == 0:\n",
|
|
" print('Epoch:', '%04d' % (epoch + 1), 'cost =', '{:.6f}'.format(loss))\n",
|
|
"\n",
|
|
" loss.backward()\n",
|
|
" optimizer.step()\n",
|
|
"\n",
|
|
" for i, label in enumerate(word_list):\n",
|
|
" W, WT = model.parameters()\n",
|
|
" x, y = W[0][i].item(), W[1][i].item()\n",
|
|
" plt.scatter(x, y)\n",
|
|
" plt.annotate(label, xy=(x, y), xytext=(5, 2), textcoords='offset points', ha='right', va='bottom')\n",
|
|
" plt.show()\n"
|
|
],
|
|
"outputs": [],
|
|
"execution_count": null
|
|
}
|
|
],
|
|
"metadata": {
|
|
"anaconda-cloud": {},
|
|
"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.6.1"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 4
|
|
} |