MT/NLLB.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/ymoslem/Adaptive-MT-LLM/blob/main/MT/NLLB.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Translation with NLLB\n",
        "* Repository: https://github.com/facebookresearch/fairseq/tree/nllb\n",
        "* Paper: https://arxiv.org/abs/2207.04672"
      ],
      "metadata": {
        "id": "wuxyY59gSSZn"
      },
      "id": "wuxyY59gSSZn"
    },
    {
      "cell_type": "markdown",
      "id": "de8d5db2",
      "metadata": {
        "id": "de8d5db2"
      },
      "source": [
        "# Method 1: CTranslate2 for Translation and SentencePiece for tokenization (faster)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "341f9ea5",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2022-11-11T00:39:37.320843Z",
          "start_time": "2022-11-11T00:39:37.318871Z"
        },
        "id": "341f9ea5"
      },
      "source": [
        "## Download the NLLB model(s)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "99933dea",
      "metadata": {
        "id": "99933dea"
      },
      "outputs": [],
      "source": [
        "# Models required for using NLLB with CTranslate2\n",
        "\n",
        "# NLLB 600M - CTranslate2 int8\n",
        "# !wget https://pretrained-nmt-models.s3.us-west-2.amazonaws.com/CTranslate2/nllb/nllb-200_600M_int8_ct2.zip\n",
        "\n",
        "# NLLB 3.3B - CTranslate2 int8\n",
        "# !wget https://pretrained-nmt-models.s3.us-west-2.amazonaws.com/CTranslate2/nllb/nllb-200_3.3B_int8_ct2.zip\n",
        "\n",
        "# SentencePiece\n",
        "# !wget https://pretrained-nmt-models.s3.us-west-2.amazonaws.com/CTranslate2/nllb/flores200_sacrebleu_tokenizer_spm.model\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "4ea0aa2a",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2022-11-07T20:13:20.610986Z",
          "start_time": "2022-11-07T20:12:58.241811Z"
        },
        "id": "4ea0aa2a"
      },
      "outputs": [],
      "source": [
        "# Example of converting an NLLB model to CTranslate2 with int8 quantization\n",
        "# Note: You do not need this step, if you downloaed the NLLB model from the previous cell.\n",
        "\n",
        "# !ct2-transformers-converter --model facebook/nllb-200-distilled-600M --quantization int8 --output_dir ct2/nllb-200-distilled-600M-int8\n",
        "\n",
        "# If you are using a GPU for inference, float16 quantization can be faster.\n",
        "\n",
        "# !ct2-transformers-converter --model facebook/nllb-200-distilled-600M --quantization float16 --output_dir ct2/nllb-200-distilled-600M-float16\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "2115051f",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2022-11-11T21:43:51.959892Z",
          "start_time": "2022-11-11T21:43:49.711868Z"
        },
        "id": "2115051f"
      },
      "outputs": [],
      "source": [
        "# Load the model and tokenizer\n",
        "# Important: This should be done only once\n",
        "\n",
        "import ctranslate2\n",
        "import sentencepiece as spm\n",
        "\n",
        "device = \"cuda\"  # or \"cpu\"\n",
        "\n",
        "# [Modify] Set paths to the CTranslate2 and SentencePiece models\n",
        "ct_model_path = \"ct2/nllb-200-3.3B-int8\"\n",
        "sp_model_path = \"flores200_sacrebleu_tokenizer_spm.model\"\n",
        "\n",
        "# Load the source SentecePiece model\n",
        "sp = spm.SentencePieceProcessor()\n",
        "sp.load(sp_model_path)\n",
        "\n",
        "translator = ctranslate2.Translator(ct_model_path, device)\n",
        "# translator = ctranslate2.Translator(ct_model_path, device=\"cuda\", device_index=[0,1])  # multiple GPUs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "3e9e7254",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2022-11-11T21:15:21.932986Z",
          "start_time": "2022-11-11T21:15:21.930338Z"
        },
        "id": "3e9e7254",
        "outputId": "f9e0b933-10a5-46b8-99cc-d66ca948479a"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "3.0.1\n"
          ]
        }
      ],
      "source": [
        "print(ctranslate2.__version__)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "0a3a118f",
      "metadata": {
        "id": "0a3a118f"
      },
      "source": [
        "## Translation of a list of sentences"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "9922c527",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2022-11-11T21:43:56.675775Z",
          "start_time": "2022-11-11T21:43:52.891825Z"
        },
        "id": "9922c527",
        "outputId": "e3921543-5d22-44b1-affb-26470bdc2f18"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Translations:\n",
            "•  Los ensayos clínicos chinos en Wuhan y Shenzhen afirmaron haber demostrado que el favipiravir era \"claramente efectivo\".\n",
            "•  Podemos ver que la prueba de PCR en tiempo real para el ácido nucleico en el tracto respiratorio o muestras de sangre se agregó a la segunda (18 de enero de 2020) y tercera (22 de enero de 2020) ediciones.\n",
            "•  El virus SARS-CoV-2 es la causa de COVID-19 (enfermedad por coronavirus 2019), una enfermedad respiratoria contagiosa que se identificó por primera vez en diciembre de 2019, en Wuhan, Hubei, China.\n"
          ]
        }
      ],
      "source": [
        "# Translate a list of sentences\n",
        "\n",
        "# Source and target langauge codes\n",
        "src_lang = \"eng_Latn\"\n",
        "tgt_lang = \"spa_Latn\"\n",
        "\n",
        "beam_size = 4\n",
        "\n",
        "source_sents = ['Chinese clinical trials in Wuhan and Shenzhen claimed to show that favipiravir was \"clearly effective\".',\n",
        "                \"We can see that the real-time PCR test for nucleic acid in respiratory tract or blood samples was added to the second (18 January 2020) and third (22 January 2020) editions.\",\n",
        "                \"The SARS-CoV-2 virus is the cause of COVID-19 (coronavirus disease 2019), a contagious respiratory disease that was first identified in December 2019, in Wuhan, Hubei, China.\"\n",
        "               ]\n",
        "\n",
        "source_sents = [sent.strip() for sent in source_sents]\n",
        "target_prefix = [[tgt_lang]] * len(source_sents)\n",
        "\n",
        "# Subword the source sentences\n",
        "source_sents_subworded = sp.encode_as_pieces(source_sents)\n",
        "source_sents_subworded = [[src_lang] + sent + [\"</s>\"] for sent in source_sents_subworded]\n",
        "\n",
        "# Translate the source sentences\n",
        "translator = ctranslate2.Translator(ct_model_path, device=device)\n",
        "translations = translator.translate_batch(source_sents_subworded, batch_type=\"tokens\", max_batch_size=2024, beam_size=beam_size, target_prefix=target_prefix)\n",
        "translations = [translation.hypotheses[0] for translation in translations]\n",
        "\n",
        "# Desubword the target sentences\n",
        "translations_desubword = sp.decode(translations)\n",
        "translations_desubword = [sent[len(tgt_lang):].strip() for sent in translations_desubword]\n",
        "\n",
        "\n",
        "print(\"Translations:\", *translations_desubword, sep=\"\\n• \")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "66c37385",
      "metadata": {
        "id": "66c37385"
      },
      "source": [
        "## File Translation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "914e2ec3",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2022-11-11T21:18:29.815557Z",
          "start_time": "2022-11-11T21:15:23.441841Z"
        },
        "id": "914e2ec3"
      },
      "outputs": [],
      "source": [
        "# Translate a file\n",
        "\n",
        "# Source and target langauge codes\n",
        "# arb_Arab deu_Latn ita_Latn rus_Cyrl spa_Latn zho_Hans\n",
        "src_lang = \"eng_Latn\"\n",
        "tgt_lang = \"spa_Latn\"\n",
        "\n",
        "beam_size = 4\n",
        "\n",
        "file_path = \"tico/tico-19.final.en.test\"\n",
        "\n",
        "with open(file_path) as source_file:\n",
        "    source_sents = source_file.readlines()\n",
        "\n",
        "source_sents = [sent.strip() for sent in source_sents]\n",
        "print(src_lang, source_sents[0], sep=\" --> \")\n",
        "target_prefix = [[tgt_lang]] * len(source_sents)\n",
        "\n",
        "# Subword the source sentences\n",
        "source_sents_subworded = sp.encode_as_pieces(source_sents)\n",
        "source_sents_subworded = [[src_lang] + sent + [\"</s>\"] for sent in source_sents_subworded]\n",
        "\n",
        "# Translate the source sentences\n",
        "translator = ctranslate2.Translator(ct_model_path, device=device)\n",
        "translations = translator.translate_batch(source_sents_subworded, batch_type=\"tokens\", max_batch_size=2048, beam_size=beam_size, target_prefix=target_prefix)\n",
        "translations = [translation.hypotheses[0] for translation in translations]\n",
        "\n",
        "# Desubword the target sentences\n",
        "translations_desubword = sp.decode(translations)\n",
        "translations_desubword = [sent[len(tgt_lang):].strip() for sent in translations_desubword]\n",
        "print(tgt_lang, translations_desubword[0], sep=\" --> \")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "386fc3ae",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2022-11-11T21:18:29.823870Z",
          "start_time": "2022-11-11T21:18:29.816880Z"
        },
        "id": "386fc3ae"
      },
      "outputs": [],
      "source": [
        "# Save the translations to the a file\n",
        "\n",
        "target_file_path = file_path + \".\" + tgt_lang\n",
        "\n",
        "with open(target_file_path, \"w+\", encoding=\"utf-8\") as target:\n",
        "    for line in translations_desubword:\n",
        "        target.write(line.strip() + \"\\n\")\n",
        "\n",
        "print(\"Done! Target file saved at:\", target_file_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "35b9fc66",
      "metadata": {
        "id": "35b9fc66"
      },
      "source": [
        "# Method 2: Hugging face for both tokenization and translation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "6b895a9e",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2022-11-11T21:42:20.549792Z",
          "start_time": "2022-11-11T21:41:34.418751Z"
        },
        "id": "6b895a9e",
        "outputId": "f5981f99-c58c-48c4-96a3-48a0186154cc"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'Los ensayos clínicos chinos en Wuhan y Shenzhen afirmaron mostrar que el favipiravir era \"claramente efectivo\".'"
            ]
          },
          "execution_count": 5,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "from transformers import AutoModelForSeq2SeqLM, AutoTokenizer\n",
        "import torch\n",
        "\n",
        "src_lang = \"eng_Latn\"\n",
        "tgt_lang = \"spa_Latn\"\n",
        "\n",
        "tokenizer = AutoTokenizer.from_pretrained(\"facebook/nllb-200-3.3B\", src_lang=src_lang)\n",
        "model = AutoModelForSeq2SeqLM.from_pretrained(\"facebook/nllb-200-3.3B\",\n",
        "                                              torch_dtype=torch.float16,\n",
        "                                              low_cpu_mem_usage=True,)\n",
        "model = model.half()\n",
        "model.to(\"cuda\")\n",
        "\n",
        "source_text = 'Chinese clinical trials in Wuhan and Shenzhen claimed to show that favipiravir was \"clearly effective\".'\n",
        "inputs = tokenizer(source_text, return_tensors=\"pt\").to(\"cuda\")\n",
        "\n",
        "translated_tokens = model.generate(\n",
        "        **inputs, forced_bos_token_id=tokenizer.convert_tokens_to_ids(tgt_lang), num_beams=5, max_length=100\n",
        ")\n",
        "tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "eef9e952",
      "metadata": {
        "id": "eef9e952"
      },
      "source": [
        "# Method 3: CTranslate2 for translation, but Hugging Face for tokenization"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "3bfc96e9",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2022-11-11T21:46:16.630948Z",
          "start_time": "2022-11-11T21:46:12.960744Z"
        },
        "id": "3bfc96e9"
      },
      "outputs": [],
      "source": [
        "# Load the model and tokenizer\n",
        "# Important: This should be done only once\n",
        "\n",
        "import ctranslate2\n",
        "import transformers\n",
        "\n",
        "src_lang = \"eng_Latn\"\n",
        "tgt_lang = \"spa_Latn\"\n",
        "\n",
        "device = \"cuda\"  # or \"cpu\"\n",
        "beam_size = 4\n",
        "\n",
        "translator = ctranslate2.Translator(\"ct2/nllb-200-3.3B-int8\", device)\n",
        "tokenizer = transformers.AutoTokenizer.from_pretrained(\"facebook/nllb-200-3.3B\", src_lang=src_lang)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e00b4d2d",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2022-11-11T21:46:19.346340Z",
          "start_time": "2022-11-11T21:46:16.632248Z"
        },
        "id": "e00b4d2d",
        "outputId": "2b2dfa91-a791-45fa-f701-7c5b7a3b4033"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Los ensayos clínicos chinos en Wuhan y Shenzhen afirmaron mostrar que el favipiravir era \"claramente efectivo\".\n"
          ]
        }
      ],
      "source": [
        "source_text = 'Chinese clinical trials in Wuhan and Shenzhen claimed to show that favipiravir was \"clearly effective\".'\n",
        "\n",
        "source = tokenizer.convert_ids_to_tokens(tokenizer.encode(source_text))\n",
        "target_prefix = [tgt_lang]\n",
        "results = translator.translate_batch([source], target_prefix=[target_prefix], beam_size=beam_size)\n",
        "target = results[0].hypotheses[0][1:]\n",
        "\n",
        "print(tokenizer.decode(tokenizer.convert_tokens_to_ids(target)))"
      ]
    }
  ],
  "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.7.5"
    },
    "colab": {
      "provenance": [],
      "include_colab_link": true
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}