Mistral-CTranslate2-Adaptive-MT.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "view-in-github"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/ymoslem/Adaptive-MT-LLM-Fine-tuning/blob/main/Mistral-CTranslate2-Adaptive-MT.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9l_I-oO92muC"
      },
      "source": [
        "# Translation with Mistral 7B (baseline and fine-tuned models)\n",
        "\n",
        "This notebook is part of the repository [Adaptive-MT-LLM-Fine-tuning](https://github.com/ymoslem/Adaptive-MT-LLM-Fine-tuning)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hnh5nkfxQsyb"
      },
      "source": [
        "# Loading the data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 36,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        },
        "id": "7jIwEIkCQe1b",
        "outputId": "e7887c35-f8f5-4588-9f3e-e48c0226a140"
      },
      "outputs": [
        {
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            },
            "text/plain": [
              "'/content/drive/MyDrive/data/spanish'"
            ]
          },
          "execution_count": 36,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "import os\n",
        "\n",
        "data_path = \"/content/drive/MyDrive/data/\"\n",
        "directory = os.path.join(data_path, \"spanish\")\n",
        "\n",
        "os.chdir(directory)\n",
        "os.getcwd()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 37,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4EznpQIHQhGo",
        "outputId": "49197d44-2441-4600-eac7-7f3dea57a249"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Período de validez después de abierto el envase: 10 horas.\n",
            "Shelf life after first opening the container: 10 hours.\n"
          ]
        }
      ],
      "source": [
        "# Load test datasets\n",
        "\n",
        "source_test_file = \"all-filtered.es.real.test\"\n",
        "target_test_file = \"all-filtered.en.real.test\"\n",
        "\n",
        "with open(source_test_file, encoding=\"utf-8\") as source, open(target_test_file, encoding=\"utf-8\") as target:\n",
        "  source_sentences = [sent.strip() for sent in source.readlines()]\n",
        "  target_sentences = [sent.strip() for sent in target.readlines()]\n",
        "\n",
        "print(source_sentences[0])\n",
        "print(target_sentences[0])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Dzsxxv1tQj3b",
        "outputId": "390b4ba4-0f56-4afb-e5f9-77653d100dac"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Período de validez después de abierto el envase: 4 semanas\n",
            "Período de validez después de abierto el envase: 10 horas.\n",
            "Shelf life after opening the immediate packaging: 4 weeks.\n"
          ]
        }
      ],
      "source": [
        "# Load the fuzzy matches from the Context Dataset\n",
        "\n",
        "online_test_file = \"all-filtered.esen.ms-multi-12.online.test\"\n",
        "\n",
        "with open(online_test_file, encoding=\"utf-8\") as online:\n",
        "  lines = [line.strip().split(\" ||| \") for line in online.readlines()]\n",
        "  scores = [float(line[0].strip()) for line in lines]\n",
        "  fuzzy_source_sentences = [line[1].strip() for line in lines]\n",
        "  online_source_sentences = [line[2].strip() for line in lines]\n",
        "  fuzzy_target_prefixes = [line[3].strip() for line in lines]\n",
        "\n",
        "print(fuzzy_source_sentences[0])\n",
        "print(online_source_sentences[0])\n",
        "print(fuzzy_target_prefixes[0])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DYb9dxsAQtpt"
      },
      "source": [
        "# Create the prompts"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "lhtPPU5zDFp0"
      },
      "outputs": [],
      "source": [
        "# Function to create zero-shot and one-shot prompts\n",
        "\n",
        "def create_prompt(source_lang,\n",
        "                  target_lang,\n",
        "                  fuzzy_sources,\n",
        "                  fuzzy_targets,\n",
        "                  new_sources,\n",
        "                  one_shot=True\n",
        "                  ):\n",
        "\n",
        "  prompts = []\n",
        "\n",
        "  if one_shot:\n",
        "    for fuzzy_src, fuzzy_tgt, new_src in zip(fuzzy_sources, fuzzy_targets, new_sources):\n",
        "      fuzzy_src = source_lang + \": \" + fuzzy_src\n",
        "      fuzzy_tgt = target_lang + \": \" + fuzzy_tgt\n",
        "      new_src = source_lang + \": \" + new_src\n",
        "      segment = fuzzy_src + \"\\n\" + fuzzy_tgt + \"\\n\" + new_src + \"\\n\" + target_lang + \":\"\n",
        "      prompts.append(segment)\n",
        "  else:\n",
        "    for new_src in new_sources:\n",
        "      new_src = source_lang + \": \" + new_src\n",
        "      segment = new_src + \"\\n\" + target_lang + \":\"\n",
        "      prompts.append(segment)\n",
        "\n",
        "  return prompts"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "id": "l_XXEc-cQvno"
      },
      "outputs": [],
      "source": [
        "source_lang = \"Spanish\"\n",
        "target_lang = \"English\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "FVTTpMt8QxKX",
        "outputId": "dc418987-e916-484b-f17e-203db0f3b2f1"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "10000\n",
            "10000\n"
          ]
        }
      ],
      "source": [
        "# Create prompts\n",
        "\n",
        "prompts_zero_shot = create_prompt(source_lang,\n",
        "                                  target_lang,\n",
        "                                  fuzzy_source_sentences,\n",
        "                                  fuzzy_target_prefixes,\n",
        "                                  online_source_sentences,\n",
        "                                  one_shot=False\n",
        "                                  )\n",
        "\n",
        "prompts_one_shot = create_prompt(source_lang,\n",
        "                                  target_lang,\n",
        "                                  fuzzy_source_sentences,\n",
        "                                  fuzzy_target_prefixes,\n",
        "                                  online_source_sentences,\n",
        "                                  one_shot=True\n",
        "                                  )\n",
        "\n",
        "print(len(prompts_zero_shot))\n",
        "print(len(prompts_one_shot))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HAlrExMiEDhk",
        "outputId": "b38b90a1-81b5-4530-f49b-4ea7dbf0e6b1"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Spanish: Período de validez después de abierto el envase: 10 horas.\n",
            "English: \n",
            "\n",
            "Spanish: Período de validez después de abierto el envase: 4 semanas\n",
            "English: Shelf life after opening the immediate packaging: 4 weeks.\n",
            "Spanish: Período de validez después de abierto el envase: 10 horas.\n",
            "English:\n"
          ]
        }
      ],
      "source": [
        "print(prompts_zero_shot[0], \"\\n\")\n",
        "print(prompts_one_shot[0])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dgWEt8FxIdbG"
      },
      "source": [
        "# Loading the model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Zxr4HGJT4LX0",
        "outputId": "808238c3-a19b-4d4d-8266-f17c966842cc"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m36.8/36.8 MB\u001b[0m \u001b[31m44.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h"
          ]
        }
      ],
      "source": [
        "!pip3 install CTranslate2 transformers -q"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZcHklgkgUCXu"
      },
      "outputs": [],
      "source": [
        "# Google Colab switched to CUDA 12 while CTranslate2 still uses CUDA 11\n",
        "# RuntimeError: Library libcublas.so.11 is not found or cannot be loaded\n",
        "# If you received this error during translation, try to install libcublas11\n",
        "\n",
        "# !apt install libcublas11"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "C-1sI-17ADLJ"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "from google.colab import userdata\n",
        "\n",
        "shared_drive = userdata.get(\"shared_drive\")\n",
        "\n",
        "directory = os.path.join(shared_drive, \"models\")\n",
        "\n",
        "os.chdir(directory)\n",
        "os.getcwd()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "dZ28kfD5BT-F"
      },
      "outputs": [],
      "source": [
        "# To convert Mistral baseline (before fine-tuning) to the CTranslate2 format, if you did not already\n",
        "# !ct2-transformers-converter --model mistralai/Mistral-7B-v0.1 --quantization int8 --output_dir ct2-mistral-7B-v0.1\n",
        "\n",
        "# To convert Mistral after FINE-TUNING to the CTranslate2 format, check the steps here:\n",
        "# https://github.com/ymoslem/Adaptive-MT-LLM-Fine-tuning/blob/main/Convert-Mistral-Finetuned-CTranslate2.ipynb"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 39,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7xbLTSSDJiiK",
        "outputId": "5bb6956f-f2cf-4ad9-928b-23b400c3d61d"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "config.json  model.bin\tvocabulary.json\n",
            "config.json  model.bin\tvocabulary.json\n"
          ]
        }
      ],
      "source": [
        "!ls $directory\"/ct2-mistral-finetuned-v1-25Nov\"\n",
        "!ls $directory\"/ct2-mistral-finetuned-v2-26Nov\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ln6SAXa77g7z"
      },
      "outputs": [],
      "source": [
        "# Load the model\n",
        "\n",
        "import ctranslate2\n",
        "import transformers\n",
        "import os\n",
        "\n",
        "# Mistral - Baseline model\n",
        "# model_name = \"ct2-mistral-7B-v0.1\"\n",
        "# tokenizer_name = \"mistralai/Mistral-7B-v0.1\"\n",
        "\n",
        "# Mistral - FINE-TUNED model\n",
        "# model_name = \"ct2-mistral-finetuned-v1-25Nov\"\n",
        "# tokenizer_name = \"mistralai/Mistral-7B-v0.1\"\n",
        "\n",
        "# model_name = \"ct2-mistral-finetuned-v2-26Nov\"\n",
        "# tokenizer_name = \"mistralai/Mistral-7B-v0.1\"\n",
        "\n",
        "model = os.path.join(directory, model_name)\n",
        "\n",
        "generator = ctranslate2.Generator(model, device=\"cuda\")\n",
        "tokenizer = transformers.AutoTokenizer.from_pretrained(tokenizer_name)\n",
        "\n",
        "print(\"Model:\", model_name)\n",
        "print(\"Tokenizer:\", tokenizer_name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eMWOckSmR7pw"
      },
      "source": [
        "# Translation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 41,
      "metadata": {
        "id": "PUyR0h7vOtSy"
      },
      "outputs": [],
      "source": [
        "# Add stopping criteria to avoid over-generation\n",
        "# References:\n",
        "# https://github.com/OpenNMT/CTranslate2/issues/1309\n",
        "# https://github.com/OpenNMT/CTranslate2/issues/1322\n",
        "# https://stackoverflow.com/questions/69403613/how-to-early-stop-autoregressive-model-with-a-list-of-stop-words\n",
        "\n",
        "stopping_criteria = tokenizer.convert_ids_to_tokens(tokenizer.encode(\".\\n\"))\n",
        "# Probably also re-add the default end of sentence token, but maybe it is not nescessary\n",
        "\n",
        "# For Llama-3, try the following:\n",
        "stopping_criteria = tokenizer.convert_ids_to_tokens(tokenizer.encode(\"<|end_of_text|>.\\n\"))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 42,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fZsLGH3l5pxY",
        "outputId": "944fbfa5-71fd-4ae2-c661-611a6d765ebe"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "max_length=40\n",
            "output_length=11\n",
            "\n",
            "Translation:\n",
            "Shelf life after opening: 10 hours.\n"
          ]
        }
      ],
      "source": [
        "# Test: Tokenize and generate (single prompt)\n",
        "\n",
        "n = 0\n",
        "prompt = prompts_zero_shot[n]\n",
        "\n",
        "max_length = len(prompt.split(\"\\n\")[-2].split(\" \")[1:]) * 4\n",
        "\n",
        "\n",
        "tokens = tokenizer.convert_ids_to_tokens(tokenizer.encode(prompt))\n",
        "\n",
        "results = generator.generate_batch([tokens],\n",
        "                                   sampling_topk=1,  # 1 for greedy search\n",
        "                                   max_length=max_length,\n",
        "                                   include_prompt_in_result=False,\n",
        "                                   end_token=stopping_criteria,\n",
        "                                   min_length=1,\n",
        "                                   batch_type=\"tokens\",\n",
        "                                   max_batch_size=8096,\n",
        "                                   )\n",
        "output_ids = results[0].sequences_ids[0]\n",
        "output = tokenizer.decode(output_ids).strip()\n",
        "\n",
        "output_length = len(output_ids)\n",
        "print(f\"{max_length=}\")\n",
        "print(f\"{output_length=}\")\n",
        "\n",
        "print(f\"\\nTranslation:\\n{output}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 43,
      "metadata": {
        "id": "WT7kGz5UZgpe"
      },
      "outputs": [],
      "source": [
        "# Batch translation fuction\n",
        "\n",
        "def translate_batch(prompts,\n",
        "                    tokenizer,\n",
        "                    generator,\n",
        "                    max_length,\n",
        "                    end_token,\n",
        "                    topk=1,\n",
        "                    ):\n",
        "  # Tokenize the prompts\n",
        "  tokenized_inputs = tokenizer(prompts)\n",
        "\n",
        "  # Extract the token IDs for the batch\n",
        "  input_ids_batch = tokenized_inputs['input_ids']\n",
        "\n",
        "  # Convert the batch of token IDs to tokens\n",
        "  tokens_batch = [tokenizer.convert_ids_to_tokens(ids) for ids in input_ids_batch]\n",
        "\n",
        "  # Generate outputs in a batch\n",
        "  results = generator.generate_batch(tokens_batch,\n",
        "                                     sampling_topk=1,  # 1 for greed search\n",
        "                                     max_length=max_length,\n",
        "                                     min_length=1,\n",
        "                                     include_prompt_in_result=False,\n",
        "                                     end_token=end_token,\n",
        "                                     batch_type=\"tokens\",\n",
        "                                     max_batch_size=8096, # 32384 # try smaller numbers if you run out of memory\n",
        "                                     )\n",
        "\n",
        "  # Decode the outputs\n",
        "  sequences_ids = [result.sequences_ids[0] for result in results]\n",
        "  translations = tokenizer.batch_decode(sequences_ids,\n",
        "                                        skip_special_tokens=True,\n",
        "                                        )\n",
        "\n",
        "  return translations"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 25,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "whT9F8gZUAvn",
        "outputId": "e32824e6-76cc-452f-e908-4ccd1af3f755"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "[['▁Shel',\n",
              "  'f',\n",
              "  '▁life',\n",
              "  '▁after',\n",
              "  '▁opening',\n",
              "  ':',\n",
              "  '▁',\n",
              "  '1',\n",
              "  '0',\n",
              "  '▁hours',\n",
              "  '.']]"
            ]
          },
          "execution_count": 25,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "results[0].sequences"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 49,
      "metadata": {
        "id": "ScBp6rpCOZE3"
      },
      "outputs": [],
      "source": [
        "# Parameters\n",
        "length_multiplier = 4\n",
        "topk = 1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 44,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "WKh9naDxMHwl",
        "outputId": "45f8038c-f284-48d6-c6d8-b774e1c9d8e9"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Spanish: Período de validez después de abierto el envase: 4 semanas\n",
            "English: Shelf life after opening the immediate packaging: 4 weeks.\n",
            "Spanish: Período de validez después de abierto el envase: 10 horas.\n",
            "English:\n"
          ]
        }
      ],
      "source": [
        "# @title ✳️ Set prompts (zero-shot, one-shot)\n",
        "\n",
        "# prompts = prompts_zero_shot\n",
        "prompts = prompts_one_shot\n",
        "\n",
        "print(prompts[0])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 50,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "PXy_nPCPT_So",
        "outputId": "2e738a31-25f0-476a-ef7b-6cf6d3efcb55"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Max length: 280\n"
          ]
        }
      ],
      "source": [
        "length = [len(prompt.split(\"\\n\")[-2].split(\" \")[1:]) for prompt in prompts]\n",
        "max_len = max(length) * length_multiplier\n",
        "print(f\"Max length: {max_len}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 51,
      "metadata": {
        "id": "jlmpOSLlSdlg"
      },
      "outputs": [],
      "source": [
        "translations = translate_batch(prompts,\n",
        "                              tokenizer,\n",
        "                              generator,\n",
        "                              max_len,\n",
        "                              stopping_criteria,\n",
        "                              topk\n",
        "                              )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7nOyeq0-WJ7O"
      },
      "outputs": [],
      "source": [
        "print(len(translations))\n",
        "\n",
        "print(*translations[:10], sep=\"\\n\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XZS2np8Iw9Lo"
      },
      "source": [
        "# Save the translations"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 53,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        },
        "id": "e0PryqEOw-mV",
        "outputId": "1d4c9d4d-e6d3-4218-add5-9b229f59ba2a"
      },
      "outputs": [
        {
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            },
            "text/plain": [
              "'/content/drive/MyDrive/data/spanish'"
            ]
          },
          "execution_count": 53,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "import os\n",
        "\n",
        "data_path = \"/content/drive/MyDrive/data/\"\n",
        "directory = os.path.join(data_path, \"spanish\")\n",
        "\n",
        "os.chdir(directory)\n",
        "os.getcwd()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 54,
      "metadata": {
        "id": "uAcRk5kRWsuW"
      },
      "outputs": [],
      "source": [
        "# Save the translations\n",
        "\n",
        "\n",
        "translations_file_name = \"all-filtered.esen.ms-multi-12.online.test.translated-Mistral-zero-shot-batch8096.en\"\n",
        "translations_file_name = \"all-filtered.esen.ms-multi-12.online.test.translated-Mistral-one-shot-batch8096.en\"\n",
        "\n",
        "# translations_file_name = \"all-filtered.esen.ms-multi-12.online.test.translated-Mistral-finetuned-v1-zero-shot.en\"\n",
        "# translations_file_name = \"all-filtered.esen.ms-multi-12.online.test.translated-Mistral-finetuned-v1-one-shot.en\"\n",
        "\n",
        "# translations_file_name = \"all-filtered.esen.ms-multi-12.online.test.translated-Mistral-finetuned-v2-zero-shot.en\"\n",
        "# translations_file_name = \"all-filtered.esen.ms-multi-12.online.test.translated-Mistral-finetuned-v2-one-shot.en\"\n",
        "\n",
        "with open(translations_file_name, \"w+\") as output:\n",
        "  for translation in translations:\n",
        "    output.write(translation + \"\\n\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kXl3b5QuXVun"
      },
      "outputs": [],
      "source": [
        "!wc -l $translations_file_name\n",
        "!head -n 10 $translations_file_name"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "authorship_tag": "ABX9TyPA71ixWQmr5/5fxlon14rg",
      "gpuType": "A100",
      "include_colab_link": true,
      "machine_shape": "hm",
      "mount_file_id": "1c_UHxw-hkS-x7eJma37Zy1uDgf5XKhoF",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}