prediction.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "yrN5xKT31vRy",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!pip install transformers"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JMKaT-L5Sayt",
        "colab_type": "text"
      },
      "source": [
        "# Next Word Prediction Demo\n",
        "\n",
        "At first, all the pre-trained models I could find all use one hot encoding and an LSTM to predict the next characters in the time series where transfer learning could not be achieved since the USE population is very different than the one that trained those models. Next, I thought about finding the most adjacent vector after appending every word in the English vocabulary after the user prompt, but it is quite impractical. \n",
        "\n",
        "Finally, I chose to use Huggingface's pre-trained OpenAI GPT2 because it uses the same attention mechanism in USE as they are both transformers so the original goal of predicting the next word with greater-than-word context is accomplished. The tradeoff here would be the large model size, compared to the DAN version of USE, but given how short the user prompt is expected to be, even if the model evaluation is done on the server-end, latency of 100ms (rough CPU evaluation time shown bellow) + 200ms (estimated additional networking latency) = 300ms shouldn't be a big issue."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EtQYiOdC2Pqx",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import torch\n",
        "from transformers import GPT2Tokenizer, GPT2LMHeadModel\n",
        "import time\n",
        "# Optional logging\n",
        "# import logging\n",
        "# logging.basicConfig(level=logging.INFO)\n",
        "\n",
        "# Tokenizer and model initialization and weight loading\n",
        "tokenizer = GPT2Tokenizer.from_pretrained('gpt2')\n",
        "model = GPT2LMHeadModel.from_pretrained('gpt2')\n",
        "model.eval()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XsF43yld71Is",
        "colab_type": "code",
        "outputId": "be0023bf-7c0a-4cfb-d1c2-094911d0241b",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 121
        }
      },
      "source": [
        "# User prompt\n",
        "prompt = ['This demo can be',\n",
        "          'How to approximate the area of the']\n",
        "\n",
        "for i in prompt:\n",
        "  start = time.time()\n",
        "  # GPT2 encoder works well with GPT2 LM\n",
        "  encoded_prompt = tokenizer.encode(i)\n",
        "\n",
        "  # Disabling auto gradient because it's not necessary for evaluation\n",
        "  with torch.no_grad():\n",
        "    outputs = model(torch.tensor([encoded_prompt]))\n",
        "    predictions = outputs[0]\n",
        "  \n",
        "  # Extract the highest scoring succeeding index\n",
        "  predicted_index = torch.argmax(predictions[0,-1,:]).item()\n",
        "\n",
        "  # The list is just for the printing format\n",
        "  print(i, '+', [tokenizer.decode(predicted_index)])\n",
        "  print('^ that took', time.time()-start, 'secs\\n')"
      ],
      "execution_count": 74,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "This demo can be + [' downloaded']\n",
            "^ that took 0.1089472770690918 secs\n",
            "\n",
            "How to approximate the area of the + [' earth']\n",
            "^ that took 0.12757611274719238 secs\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "Untitled0.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}