docs/notebook_tutorial/wrap_MT_inference_pipeline.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Machine Translation Inference Pipeline\n",
    "## Packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import shutil\n",
    "from typing import Dict\n",
    "from transformers import T5Tokenizer, T5ForConditionalGeneration\n",
    "from forte import Pipeline\n",
    "from forte.data import DataPack\n",
    "from forte.common import Resources, Config\n",
    "from forte.processors.base import PackProcessor\n",
    "from forte.data.readers import PlainTextReader"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Background\n",
    "\n",
    "After a Data Scientist is satisfied with the results of a training model, they will had their notebook over to an MLE who has to convert their model into an inference model. \n",
    "\n",
    "## Inference Workflow\n",
    "\n",
    "### Pipeline\n",
    "To simplify the example, we consider `t5-small` as a trained MT model. We should always consider pipeline first when it comes to an inference workflow. As the [glossary](https://asyml-forte.readthedocs.io/en/latest/index_appendices.html#glossary) suggests, it's an inference system that contains a set of processing components. \n",
    "\n",
    "Therefore, we initialize a `pipeline` below.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pipeline: Pipeline = Pipeline[DataPack]()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Reader\n",
    "After observing the dataset, it's a plain `txt` file. Therefore, we can use `PlainTextReader` directly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pipeline.set_reader(PlainTextReader())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "However, it's still beneficial to take a deeper look at how this class is designed so that users can customize a reader when needed.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "### Processor\n",
    "We already have inference model which is the `t5-small`, and we need a component to make inference. Therefore, besides model itself, there are several behaviors needed.\n",
    "1. tokenization that transforms input text to sequences of tokens.\n",
    "2. since T5 has a better performance given a task prompt, we also want to include the prompt in our data.\n",
    "\n",
    "In forte, we have a generic class `PackProcessor` that wraps model and inference related components and behaviors to process `DataPack`, and we need to create a class inherit generic method and customize the behaviors.\n",
    "\n",
    "The generic method to process `DataPack` is `_process(self, input_pack: DataPack)`. It should tokenize the input text, use class model to make an inference, decode the output token ids and finally writes the output to a target file.\n",
    "\n",
    "Given what we discussed, we have a processor class below, and we need to add it to pipeline after define it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class MachineTranslationProcessor(PackProcessor):\n",
    "    \"\"\"\n",
    "    Translate the input text and output to a file.\n",
    "    \"\"\"\n",
    "    def initialize(self, resources: Resources, configs: Config):\n",
    "        super().initialize(resources, configs)\n",
    "\n",
    "        # Initialize the tokenizer and model\n",
    "        model_name: str = self.configs.pretrained_model\n",
    "        self.tokenizer = T5Tokenizer.from_pretrained(model_name)\n",
    "        self.model = T5ForConditionalGeneration.from_pretrained(model_name)\n",
    "        self.task_prefix = \"translate English to German: \"\n",
    "        self.tokenizer.padding_side = \"left\"\n",
    "        self.tokenizer.pad_token = self.tokenizer.eos_token\n",
    "\n",
    "        if not os.path.isdir(self.configs.output_folder):\n",
    "            os.mkdir(self.configs.output_folder)\n",
    "\n",
    "    def _process(self, input_pack: DataPack):\n",
    "        file_name: str = os.path.join(\n",
    "            self.configs.output_folder, os.path.basename(input_pack.pack_name)\n",
    "        )\n",
    "\n",
    "        # en2de machine translation \n",
    "        inputs = self.tokenizer([\n",
    "            self.task_prefix + sentence\n",
    "            for sentence in input_pack.text.split('\\n')\n",
    "        ], return_tensors=\"pt\", padding=True)\n",
    "\n",
    "        output_sequences = self.model.generate(\n",
    "            input_ids=inputs[\"input_ids\"],\n",
    "            attention_mask=inputs[\"attention_mask\"],\n",
    "            do_sample=False,\n",
    "        )\n",
    "\n",
    "        outputs = self.tokenizer.batch_decode(\n",
    "            output_sequences, skip_special_tokens=True\n",
    "        )\n",
    "\n",
    "        # Write output to the specified file\n",
    "        with open(file=file_name, mode='w') as f:\n",
    "            f.write('\\n'.join(outputs))\n",
    "\n",
    "    @classmethod\n",
    "    def default_configs(cls) -> Dict:\n",
    "        return {\n",
    "            \"pretrained_model\": \"t5-small\",\n",
    "            \"output_folder\": \"mt_test_output\"\n",
    "        }\n",
    "\n",
    "pipeline.add(MachineTranslationProcessor(), config={\n",
    "    \"pretrained_model\": \"t5-small\"\n",
    "})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "Examples\n",
    "---------\n",
    "\n",
    "We have an working [MT translation pipeline example](https://github.com/asyml/forte/blob/master/docs/notebook_tutorial/wrap_MT_inference_pipeline.ipynb).\n",
    "\n",
    "There are several basic functions of processor and internal functions are defined in this example.\n",
    "\n",
    "* ``initialize()``: Pipeline will call it at the start of processing. The processor will be initialized with\n",
    "        ``configs``, and register global resources into ``resource``. The\n",
    "        implementation should set up the states of the component.\n",
    "    - initialize a pre-trained model\n",
    "    - initialize tokenizer\n",
    "    - initialize model-specific attributes such as task prefix\n",
    "* ``process()``: using the loaded model to make predictions and write the prediction results out.\n",
    "    - we first tokenize the input text\n",
    "    - then use model to generate output sequence ids\n",
    "    - then we decode output sequence ids into tokens and write the output into a file"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After setting up pipeline's components, we can run the pipeline on the input directory as below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dir_path = os.path.abspath(\n",
    "            os.path.join(\"data_samples\", \"machine_translation\")\n",
    "        ) # notebook should be running from project root folder\n",
    "\n",
    "pipeline.run(dir_path)\n",
    "print(\"Done successfully\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One can investigate the machine translation output in folder `mt_test_output` located at the script location.\n",
    "Then we remove the output folder below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "shutil.rmtree(MachineTranslationProcessor.default_configs()[\"output_folder\"])"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "ef24553d1198fce4a0be9f455df40aff4e6272106653c30479d64479a9d4460b"
  },
  "kernelspec": {
   "display_name": "Python 3.7.11 ('forte_env')",
   "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.11"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}