notebooks/221-machine-translation/221-machine-translation.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Machine translation demo\n",
    "This demo utilizes Intel's pre-trained model that translates from English to German. More information about the model can be found [here](https://github.com/openvinotoolkit/open_model_zoo/blob/master/models/intel/machine-translation-nar-en-de-0002/README.md).\n",
    "\n",
    "This model encodes sentences using the SentecePieceBPETokenizer from HuggingFace. The tokenizer vocabulary is downloaded automatically with the OMZ tool.\n",
    "\n",
    "**Inputs**\n",
    "The model's input is a sequence of up to 150 tokens with the following structure: `<s>` + _tokenized sentence_ + `<s>` + `<pad>` (`<pad>` tokens pad the remaining blank spaces).\n",
    "\n",
    "**Ouput**\n",
    "After the inference, we have a sequence of up to 200 tokens. The structure is the same as the one for the input."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "from openvino.runtime import Core\n",
    "import numpy as np\n",
    "import itertools\n",
    "from tokenizers import SentencePieceBPETokenizer"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Downloading model\n",
    "The following command will download the model to the current directory. Make sure you have run `pip install openvino-dev` beforehand."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": true,
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "! omz_downloader --name  machine-translation-nar-en-de-0002"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Load and configure the model\n",
    "The model is now available in the `intel/` folder. Below, we load and configure its inputs and outputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "core = Core()\n",
    "model = core.read_model('intel/machine-translation-nar-en-de-0002/FP32/machine-translation-nar-en-de-0002.xml')\n",
    "compiled_model = core.compile_model(model)\n",
    "input_name = \"tokens\"\n",
    "output_name = \"pred\"\n",
    "model.output(output_name)\n",
    "max_tokens = model.input(input_name).shape[1]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Load tokenizers\n",
    "\n",
    "NLP models usually take a list of tokens as standard input. A token is a single word converted to some integer. To provide the proper input, we need the vocabulary for such mapping. We use 'merges.txt' to find out what sequences of letters form a token. 'vocab.json' specifies the mapping between tokens and integers.\n",
    "\n",
    "The input needs to be transformed into a token sequence the model understands, and the output must be transformed into a sentence that is human readable.\n",
    "\n",
    "Initialize the tokenizer for the input `src_tokenizer` and the output `tgt_tokenizer`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "src_tokenizer = SentencePieceBPETokenizer.from_file(\n",
    "    'intel/machine-translation-nar-en-de-0002/tokenizer_src/vocab.json',\n",
    "    'intel/machine-translation-nar-en-de-0002/tokenizer_src/merges.txt'\n",
    ")\n",
    "tgt_tokenizer = SentencePieceBPETokenizer.from_file(\n",
    "    'intel/machine-translation-nar-en-de-0002/tokenizer_tgt/vocab.json',\n",
    "    'intel/machine-translation-nar-en-de-0002/tokenizer_tgt/merges.txt'\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Perform translation\n",
    "The following function translates a sentence in English to German."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def translate(sentence: str) -> str:\n",
    "    \"\"\"\n",
    "    Tokenize the sentence using the downloaded tokenizer and run the model,\n",
    "    whose output is decoded into a human readable string.\n",
    "\n",
    "    :param sentence: a string containing the phrase to be translated\n",
    "    :return: the translated string\n",
    "    \"\"\"\n",
    "    # Remove leading and trailing white spaces\n",
    "    sentence = sentence.strip()\n",
    "    assert len(sentence) > 0\n",
    "    tokens = src_tokenizer.encode(sentence).ids\n",
    "    # Transform the tokenized sentence into the model's input format\n",
    "    tokens = [src_tokenizer.token_to_id('<s>')] + \\\n",
    "        tokens + [src_tokenizer.token_to_id('</s>')]\n",
    "    pad_length = max_tokens - len(tokens)\n",
    "\n",
    "    # If the sentence size is less than the maximum allowed tokens,\n",
    "    # fill the remaining tokens with '<pad>'.\n",
    "    if pad_length > 0:\n",
    "        tokens = tokens + [src_tokenizer.token_to_id('<pad>')] * pad_length\n",
    "    assert len(tokens) == max_tokens, \"input sentence is too long\"\n",
    "    encoded_sentence = np.array(tokens).reshape(1, -1)\n",
    "\n",
    "    # Perform inference\n",
    "    enc_translated = compiled_model({input_name: encoded_sentence})\n",
    "    output_key = compiled_model.output(output_name)\n",
    "    enc_translated = enc_translated[output_key][0]\n",
    "\n",
    "    # Decode the sentence\n",
    "    sentence = tgt_tokenizer.decode(enc_translated)\n",
    "\n",
    "    # Remove <pad> tokens, as well as '<s>' and '</s>' tokens which mark the\n",
    "    # beginning and ending of the sentence.\n",
    "    for s in ['</s>', '<s>', '<pad>']:\n",
    "        sentence = sentence.replace(s, '')\n",
    "\n",
    "    # Transform sentence into lower case and join words by a white space\n",
    "    sentence = sentence.lower().split()\n",
    "    sentence = \" \".join(key for key, _ in itertools.groupby(sentence))\n",
    "    return sentence"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Translate the sentence\n",
    "The following function is a basic loop that translates sentences."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def run_translator():\n",
    "    \"\"\"\n",
    "    Run the translation in real time, reading the input from the user.\n",
    "    This function prints the translated sentence and the time\n",
    "    spent during inference.\n",
    "    :return:\n",
    "    \"\"\"\n",
    "    while True:\n",
    "        input_sentence = input()\n",
    "        if input_sentence == \"\":\n",
    "            break\n",
    "\n",
    "        start_time = time.perf_counter()\n",
    "        translated = translate(input_sentence)\n",
    "        end_time = time.perf_counter()\n",
    "        print(f'Translated: {translated}')\n",
    "        print(f'Time: {end_time - start_time:.2f}s')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# uncomment the following line for a real time translation of your input\n",
    "# run_translator()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test your translation\n",
    "Run the following cell with an English sentence to have it translated to German"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "sentence = \"My name is openvino\"\n",
    "print(f'Translated: {translate(sentence)}')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}