meena_chatbot_inference.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "meena_chatbot_inference.ipynb",
      "provenance": [],
      "toc_visible": true,
      "authorship_tag": "ABX9TyOxh1UURaXFl9to/fWC5kjE",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/frankplus/meena-chatbot/blob/main/meena_chatbot_inference.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ROp4-qZF7hWB"
      },
      "source": [
        "## Initialize"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "46wppQZ4LHGn"
      },
      "source": [
        "!pip install -q -U tensorflow-gpu==1.15.2\n",
        "!pip install -q -U tensorflow-datasets==3.2.1\n",
        "!pip install -q -U tensor2tensor\n",
        " \n",
        "import tensorflow as tf\n",
        "import os\n",
        "\n",
        "tf.get_logger().propagate = False"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gnq9RQpY7kUJ"
      },
      "source": [
        "## Download pretrained model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dZR6cKWaQoIs"
      },
      "source": [
        "model_name = \"Italian_108M\"\n",
        "!gdown  https://drive.google.com/uc?id=1y0abt3nOKPo5DBfKx3b7A7pjm3GH3wi1\n",
        "!unzip {model_name}.zip\n",
        "MODEL_DIR = model_name + '/'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "us239gXxK5P0"
      },
      "source": [
        "from tensor2tensor import models\n",
        "from tensor2tensor import problems\n",
        "from tensor2tensor.utils import hparams_lib\n",
        "from tensor2tensor.utils import registry\n",
        "from tensor2tensor.data_generators import text_problems\n",
        "import numpy as np\n",
        "import re\n",
        "\n",
        "MODEL = \"evolved_transformer\"\n",
        "VOCAB_SIZE = 2**13\n",
        "\n",
        "# sampling parameters\n",
        "CONVERSATION_TURNS = 3\n",
        "SAMPLING_TEMPERATURE = 0.88\n",
        "NUM_SAMPLES = 5\n",
        "MAX_LCS_RATIO = 0.9\n",
        "\n",
        "tfe = tf.contrib.eager\n",
        "tfe.enable_eager_execution()\n",
        "Modes = tf.estimator.ModeKeys\n",
        "\n",
        "@registry.register_problem\n",
        "class ChatBot(text_problems.Text2TextProblem):\n",
        "    @property\n",
        "    def approx_vocab_size(self):\n",
        "        return VOCAB_SIZE\n",
        "\n",
        "chat_bot_problem = problems.problem(\"chat_bot\")\n",
        "ckpt_path = MODEL_DIR\n",
        "encoders = chat_bot_problem.feature_encoders(MODEL_DIR)\n",
        "hparams = hparams_lib.create_hparams_from_json(MODEL_DIR + 'hparams.json')\n",
        "hparams.data_dir = MODEL_DIR\n",
        "hparams_lib.add_problem_hparams(hparams, \"chat_bot\")\n",
        "hparams.sampling_method = \"random\"\n",
        "hparams.sampling_temp = SAMPLING_TEMPERATURE\n",
        "\n",
        "chatbot_model = registry.model(MODEL)(hparams, Modes.PREDICT)\n",
        "\n",
        "def preprocess_sentence(sentence):\n",
        "    sentence = sentence.lower().strip()\n",
        "    # creating a space between a word and the punctuation following it\n",
        "    # eg: \"he is a boy.\" => \"he is a boy .\"\n",
        "    sentence = re.sub(r\"([?.!,])\", r\" \\1 \", sentence)\n",
        "    sentence = sentence.replace(\"'\", \"' \")\n",
        "    sentence = re.sub(r'[\" \"]+', \" \", sentence)\n",
        "    sentence = re.sub(r\"[^a-zA-Z0-9?.!,àèìòùáéíóú']+\", \" \", sentence)\n",
        "    sentence = sentence.strip()\n",
        "    return sentence\n",
        "\n",
        "def postprocess_sentence(sentence):\n",
        "    # remove space before punctuation\n",
        "    sentence = sentence.rstrip(\" .\")\n",
        "    return re.sub(r\"\\s+(\\W)\", r\"\\1\", sentence)\n",
        "\n",
        "def encode(conversation, output_str=None):\n",
        "    \"\"\"Input str to features dict, ready for inference\"\"\"\n",
        "    encoded_inputs = []\n",
        "    for conversation_turn in conversation:\n",
        "        encoded_inputs += encoders[\"inputs\"].encode(conversation_turn) + [2]\n",
        "    encoded_inputs.pop()\n",
        "    encoded_inputs += [1]\n",
        "    if len(encoded_inputs) > hparams.max_length:\n",
        "        encoded_inputs = encoded_inputs[-hparams.max_length:]\n",
        "    batch_inputs = tf.reshape(encoded_inputs, [1, -1, 1])  # Make it 3D.\n",
        "    return {\"inputs\": batch_inputs}\n",
        "\n",
        "def decode(integers):\n",
        "    \"\"\"List of ints to str\"\"\"\n",
        "    integers = list(np.squeeze(integers))\n",
        "    if 1 in integers:\n",
        "        integers = integers[:integers.index(1)]\n",
        "    decoded = encoders[\"inputs\"].decode(integers)\n",
        "    return postprocess_sentence(decoded)\n",
        "\n",
        "def lcs_ratio(context, predicted): \n",
        "    m = len(context) \n",
        "    n = len(predicted) \n",
        "    L = [[None]*(n + 1) for i in range(m + 1)] \n",
        "    for i in range(m + 1): \n",
        "        for j in range(n + 1): \n",
        "            if i == 0 or j == 0 : \n",
        "                L[i][j] = 0\n",
        "            elif context[i-1] == predicted[j-1]: \n",
        "                L[i][j] = L[i-1][j-1]+1\n",
        "            else: \n",
        "                L[i][j] = max(L[i-1][j], L[i][j-1]) \n",
        "    return L[m][n] / n\n",
        "\n",
        "def predict(conversation):\n",
        "    preprocessed = [preprocess_sentence(x) for x in conversation]\n",
        "    encoded_inputs = encode(preprocessed)\n",
        "    print(\"decoded input: \" + decode(encoded_inputs[\"inputs\"]))\n",
        "    with tfe.restore_variables_on_create(ckpt_path):\n",
        "        while True:\n",
        "            output_candidates = [chatbot_model.infer(encoded_inputs, decode_length=1) for _ in range(NUM_SAMPLES)]\n",
        "            output_candidates.sort(key = lambda x: -float(x[\"scores\"]))\n",
        "\n",
        "            for x in output_candidates:\n",
        "                print(str(float(x[\"scores\"])) + \"\\t\" + decode(x[\"outputs\"]))\n",
        "\n",
        "            for candidate in output_candidates:\n",
        "                decoded = decode(candidate[\"outputs\"])\n",
        "                if lcs_ratio(\" \".join(preprocessed), decoded) < MAX_LCS_RATIO:\n",
        "                    return decoded\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FjIhTHgD99-5"
      },
      "source": [
        "## Predict"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xRJ781Bs5hhz"
      },
      "source": [
        "conversation = []\n",
        "while True:\n",
        "    sentence = input(\"Input: \")\n",
        "    conversation.append(sentence)\n",
        "    while len(conversation) > CONVERSATION_TURNS: \n",
        "        conversation.pop(0)\n",
        "    response = predict(conversation)\n",
        "    conversation.append(response)\n",
        "    print(response)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}