demo/AEVNMT.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "from pathlib import Path\n",
    "\n",
    "from aevnmt.models import AEVNMT\n",
    "from aevnmt.train import create_model\n",
    "from aevnmt.train_utils import load_vocabularies\n",
    "from aevnmt.hparams import Hyperparameters\n",
    "from aevnmt.data import create_batch, SOS_TOKEN, EOS_TOKEN, PAD_TOKEN, batch_to_sentences\n",
    "from aevnmt.components import tile_rnn_hidden, ancestral_sample\n",
    "from aevnmt.data.textprocessing import Pipeline, Detokenizer, Recaser, WordDesegmenter"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Restore the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "AEVNMT(\n",
      "  (encoder): RNNEncoder(\n",
      "    (rnn): LSTM(256, 256, batch_first=True, bidirectional=True)\n",
      "  )\n",
      "  (decoder): BahdanauDecoder(\n",
      "    (rnn): LSTM(768, 256, batch_first=True)\n",
      "    (dropout_layer): Dropout(p=0.5)\n",
      "    (pre_output_layer): Linear(in_features=1024, out_features=256, bias=True)\n",
      "    (attention): BahdanauAttention(\n",
      "      (key_layer): Linear(in_features=512, out_features=256, bias=False)\n",
      "      (query_layer): Linear(in_features=256, out_features=256, bias=False)\n",
      "      (scores_layer): Linear(in_features=256, out_features=1, bias=False)\n",
      "    )\n",
      "  )\n",
      "  (language_model): RNNLM(\n",
      "    (embedder): Embedding(9726, 256, padding_idx=1)\n",
      "    (rnn): LSTM(256, 256, batch_first=True)\n",
      "    (dropout_layer): Dropout(p=0.5)\n",
      "  )\n",
      "  (tgt_embedder): Embedding(9726, 256, padding_idx=1)\n",
      "  (dropout_layer): Dropout(p=0.5)\n",
      "  (encoder_init_layer): Sequential(\n",
      "    (0): Linear(in_features=32, out_features=256, bias=True)\n",
      "    (1): Tanh()\n",
      "  )\n",
      "  (decoder_init_layer): Sequential(\n",
      "    (0): Linear(in_features=32, out_features=256, bias=True)\n",
      "    (1): Tanh()\n",
      "  )\n",
      "  (lm_init_layer): Sequential(\n",
      "    (0): Linear(in_features=32, out_features=256, bias=True)\n",
      "    (1): Tanh()\n",
      "  )\n",
      "  (inf_network): InferenceNetwork(\n",
      "    (src_embedder): Embedding(9726, 256, padding_idx=1)\n",
      "    (encoder): RNNEncoder(\n",
      "      (rnn): LSTM(256, 256, batch_first=True, bidirectional=True)\n",
      "    )\n",
      "    (normal_layer): NormalLayer(\n",
      "      (loc_layer): Sequential(\n",
      "        (0): Linear(in_features=512, out_features=256, bias=True)\n",
      "        (1): ReLU()\n",
      "        (2): Linear(in_features=256, out_features=32, bias=True)\n",
      "      )\n",
      "      (scale_layer): Sequential(\n",
      "        (0): Linear(in_features=512, out_features=256, bias=True)\n",
      "        (1): ReLU()\n",
      "        (2): Linear(in_features=256, out_features=32, bias=True)\n",
      "        (3): Softplus(beta=1, threshold=20)\n",
      "      )\n",
      "    )\n",
      "  )\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "device_name = \"cuda:0\" if torch.cuda.is_available() else \"cpu\"\n",
    "device = torch.device(device_name)\n",
    "\n",
    "# Load the original hyperparameters.\n",
    "model_dir = Path(\"./flickr/models/aevnmt/supervised/\")\n",
    "hparams = Hyperparameters(check_required=False)\n",
    "hparams_file = model_dir / \"hparams\"\n",
    "hparams.update_from_file(hparams_file, override=False)\n",
    "\n",
    "# Load the vocabularies.\n",
    "hparams.vocab_prefix = model_dir / \"vocab\"\n",
    "hparams.share_vocab = False\n",
    "vocab_de, vocab_en = load_vocabularies(hparams)\n",
    "\n",
    "# Restore the model.\n",
    "model, _, _, translate = create_model(hparams, vocab_de, vocab_en)\n",
    "model.load_state_dict(torch.load(model_dir / \"model/bleu/de-en.pt\", map_location=device_name))\n",
    "model = model.to(device)\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sample source sentences from the latent space (greedy decoding)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1: Zwei hunde rennen in einer wiese während new york.\n",
      "2: Ein kleiner junge liegt auf einem holzzaun über dem strand.\n",
      "3: Diese menschen lesen verschiedene pinsame tag etwas in der pfanne.\n",
      "4: Ein obdachloser mann sitzt auf einem hocker vor einer skulptur und liest seinen fahrrad.\n",
      "5: Ein mann mit stock, der den kamm in seinem fahrrad fährt.\n"
     ]
    }
   ],
   "source": [
    "num_samples = 5\n",
    "\n",
    "prior = model.prior()\n",
    "z = prior.sample(sample_shape=[num_samples])\n",
    "\n",
    "# Construct the LM initial hidden state and inputs.\n",
    "hidden_lm = tile_rnn_hidden(model.lm_init_layer(z), model.language_model.rnn)\n",
    "x_init = z.new([vocab_de[SOS_TOKEN] for _ in range(num_samples)]).long()\n",
    "x_embed = model.language_model.embedder(x_init)\n",
    "\n",
    "# Keep track of model samples.\n",
    "x_samples = [x_init.unsqueeze(-1)] # List of [num_samples, 1] integers.\n",
    "\n",
    "# Sample num_samples source sentences conditioned on z.\n",
    "for _ in range(hparams.max_decoding_length):\n",
    "    hidden_lm, logits = model.language_model.step(x_embed, hidden_lm)\n",
    "    next_word_dist = torch.distributions.categorical.Categorical(logits=logits)\n",
    "    x = next_word_dist.sample()\n",
    "    x_embed = model.language_model.embedder(x.squeeze())\n",
    "    x_samples.append(x)\n",
    "\n",
    "# Concatenate the samples and convert to sentences.\n",
    "x_samples = torch.cat(x_samples, dim=-1)\n",
    "x_samples = batch_to_sentences(x_samples, vocab_de)\n",
    "\n",
    "# Construct a post-processing pipeline for German.\n",
    "postprocess = [Detokenizer(\"de\"),\n",
    "               Recaser(\"de\"), \n",
    "               WordDesegmenter(separator=hparams.subword_token)] # Executed in reverse order.\n",
    "pipeline_de = Pipeline(pre=[], post=postprocess)\n",
    "\n",
    "# Print the samples.\n",
    "pp_x_samples = [pipeline_de.post(x) for x in x_samples]\n",
    "for idx, x in enumerate(pp_x_samples, 1): print(f\"{idx}: {x}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sample translations from the approximate posterior (greedy decoding)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1: Young blond children, holding a sandwich.\n",
      "2: In pink and little girls holding a pair of sandwich.\n",
      "3: Shirt to short blond-hair, holding a sandwich\n",
      "4: A little girl with blond-hair is holding a sandwich.\n",
      "5: At young blond girl holding a sandwich.\n"
     ]
    }
   ],
   "source": [
    "num_samples = 5\n",
    "x_samples = [\"in kleines blondes mädchen hält ein sandwich .\"] * num_samples\n",
    "x_in, _, seq_mask_x, seq_len_x = create_batch(x_samples, vocab_de, device)\n",
    "\n",
    "# Infer q(z|x). \n",
    "qz = model.approximate_posterior(x_in, seq_mask_x, seq_len_x)\n",
    "z = qz.sample()\n",
    "\n",
    "# Encode the source sentences\n",
    "encoder_outputs, encoder_final = model.encode(x_in, seq_len_x, z)\n",
    "\n",
    "# Create the initial hidden state of the TM.\n",
    "hidden_tm = model.init_decoder(encoder_outputs, encoder_final, z)\n",
    "\n",
    "# Sample target sentences conditional on the source and z.\n",
    "y_samples = ancestral_sample(model.decoder,\n",
    "                             model.tgt_embed,\n",
    "                             model.generate,\n",
    "                             hidden_tm,\n",
    "                             encoder_outputs, encoder_final,\n",
    "                             seq_mask_x,\n",
    "                             vocab_en[SOS_TOKEN],\n",
    "                             vocab_en[EOS_TOKEN],\n",
    "                             vocab_en[PAD_TOKEN],\n",
    "                             hparams.max_decoding_length,\n",
    "                             greedy=False)[\"sample\"]\n",
    "y_samples = batch_to_sentences(y_samples, vocab_en)\n",
    "\n",
    "# Construct a post-processing pipeline for English.\n",
    "postprocess = [Detokenizer(\"en\"),\n",
    "               Recaser(\"en\"), \n",
    "               WordDesegmenter(separator=hparams.subword_token)] # Executed in reverse order.\n",
    "pipeline_en = Pipeline(pre=[], post=postprocess)\n",
    "\n",
    "# Print the samples.\n",
    "pp_y_samples = [pipeline_en.post(y) for y in y_samples]\n",
    "for idx, y in enumerate(pp_y_samples, 1): print(f\"{idx}: {y}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.7.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}