Transformer.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from model import Transformer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import io\n",
    "import spacy\n",
    "import numpy as np\n",
    "\n",
    "from torchtext import data\n",
    "from torchtext import datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "spacy_de = spacy.load('de')\n",
    "spacy_en = spacy.load('en')\n",
    "\n",
    "\n",
    "def tokenize_en(text):\n",
    "    return [tok.text for tok in spacy_en.tokenizer(text)]\n",
    "def tokenize_de(text):\n",
    "    return [tok.text for tok in spacy_de.tokenizer(text)]\n",
    "\n",
    "DE = data.Field(tokenize=tokenize_de,\n",
    "                init_token='<SOS>',\n",
    "                eos_token='<EOS>',\n",
    "                fix_length=20,\n",
    "                lower=True,\n",
    "                batch_first=True)\n",
    "EN = data.Field(tokenize=tokenize_en,\n",
    "                init_token='<SOS>',\n",
    "                eos_token='<EOS>',\n",
    "                lower=True,\n",
    "                fix_length=20,\n",
    "                batch_first=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "train = datasets.TranslationDataset(path='./data/train', \n",
    "                                    exts=('.de.txt', '.en.txt'),\n",
    "                                    fields=(DE, EN))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "test = datasets.TranslationDataset(path='./data/test', \n",
    "                                   exts=('.de.txt', '.en.txt'),\n",
    "                                   fields=(DE, EN))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "DE.build_vocab(train.src, min_freq=3)\n",
    "EN.build_vocab(train, max_size=50000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "train_iter = data.BucketIterator(dataset=train, \n",
    "                                 batch_size=32,\n",
    "                                 shuffle=True,\n",
    "                                 sort_key=lambda x: data.interleave_keys(len(x.src), len(x.trg)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "test_iter = data.BucketIterator(dataset=test, \n",
    "                                batch_size=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def make_batch(batch):\n",
    "    src = batch.src\n",
    "    src_mask = 1-(src==1)\n",
    "    src_mask.unsqueeze_(2)\n",
    "    src_position = torch.LongTensor([range(src.size(1))]*src.size(0)).to(src.device)\n",
    "    \n",
    "    trg = batch.trg\n",
    "    trg_mask = 1-(trg==1)\n",
    "    trg_mask.unsqueeze_(2)\n",
    "    trg_position = torch.LongTensor([range(trg.size(1))]*trg.size(0)).to(trg.device)\n",
    "    return src,src_mask,src_position,trg,trg_mask,trg_position"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "model = Transformer(len(DE.vocab),len(EN.vocab))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "cuda = torch.device('cuda')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "model = model.to(cuda)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "train_step = optim.Adam(model.parameters(),lr=1e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "loss_ = []\n",
    "\n",
    "for i,batch in enumerate(train_iter):\n",
    "    src,src_mask,src_position,trg,trg_mask,trg_position = make_batch(batch)\n",
    "    logit = model(src,src_mask,src_position,trg,trg_mask,trg_position)\n",
    "    loss = F.cross_entropy(logit[:,:-1,:].contiguous().view(-1,50004),\n",
    "                           trg[:,1:].contiguous().view(-1),ignore_index=1)\n",
    "    train_step.zero_grad()\n",
    "    loss.backward()\n",
    "    train_step.step()\n",
    "    loss_.append(loss.item())\n",
    "    \n",
    "    if i % 1000 == 0:\n",
    "        print \"epoch:%s , loss:%s\" % (i,np.mean(loss_)) \n",
    "        loss_ = []\n",
    "    \n",
    "    if i == 30000:\n",
    "        torch.save(model.state_dict(), \"model-30000.pth\")\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "_,p = logit.max(-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "' '.join([EN.vocab.itos[i] for i in trg[3]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "' '.join([EN.vocab.itos[i] for i in p[3]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "test_batch = next(iter(test_iter))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "src,src_mask,src_position,trg,trg_mask,trg_position = make_batch(test_batch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "trg_ = torch.zeros_like(src)\n",
    "trg_mask_ = torch.zeros_like(src_mask)\n",
    "trg_mask_[:,0,:] = 2\n",
    "trg_position_ = src_position"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "o = model.inference(src,src_mask,src_position.cuda(),trg_,trg_mask_,trg_position_.cuda())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "' '.join([EN.vocab.itos[i] for i in trg[1]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "' '.join([EN.vocab.itos[i] for i in o[1]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}