private/NlpInPracticeCourse/NamedEntityRecognition: Train_your_own_BERT_NER.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "#Training NER model from BERT + WikiAnn\n",
        "\n",
        "In this colab, we will use the WikiAnn corpora. WikiAnn is annotated from Wikipedia pages and their categories. Check the WikiAnn paper at https://aclanthology.org/P17-1178.pdf\n",
        "\n",
        "WikiAnn were used for tranfer learning of NER from well-resourced languages into under-resources languages. Check the paper at https://github.com/afshinrahimi/mmner\n",
        "\n",
        "The WikiAnn corpora are described at https://huggingface.co/datasets/wikiann\n",
        "\n",
        "For training, we will use the BERT model for downstream tasks.\n",
        "\n",
        "Both the model and the dataset are stored in huggingface, so we will use huggingface modules `datasets`, `tokenizers` for training, `sequeval` for evaluation, and `transformers` for prediction.\n"
      ],
      "metadata": {
        "id": "pmBLtRRTWEWm"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "llooKX2-WCRm"
      },
      "outputs": [],
      "source": [
        "!pip install datasets\n",
        "!pip install tokenizers\n",
        "!pip install transformers\n",
        "!pip install seqeval"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Make sure, we are using the GPU. If GPU is not set up, go to `Runtime`/`Change runtime type` and select `GPU`."
      ],
      "metadata": {
        "id": "_53cGbGEE-TA"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import tensorflow as tf\n",
        "tf.test.gpu_device_name()"
      ],
      "metadata": {
        "id": "7_EUzq-boboJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Here, we load the WikiAnn corpus. We can use the huggingface `dataset` module. Check https://huggingface.co/datasets/wikiann for available langauges and data sizes."
      ],
      "metadata": {
        "id": "kXd9oSPEFNVa"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from datasets import load_dataset\n",
        "\n",
        "dataset = load_dataset(\"wikiann\", \"sk\")"
      ],
      "metadata": {
        "id": "Vw-CwOigYSPE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "By loading the WikiAnn datasets, we obtain the `DatasetDict`. The data itself is under `DatasetDict.data`, however, we will work with the dictionary."
      ],
      "metadata": {
        "id": "xN2cYTFwFmtA"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "type(dataset)"
      ],
      "metadata": {
        "id": "9974i-vSY6YT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "dataset"
      ],
      "metadata": {
        "id": "MVnPJl2TYa7i"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "type(dataset['train'])"
      ],
      "metadata": {
        "id": "6877QCS5g9Ps"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "dataset[\"train\"].features"
      ],
      "metadata": {
        "id": "6jBiyLRHYeax"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "label_names = dataset[\"train\"].features[\"ner_tags\"].feature.names\n",
        "label_names"
      ],
      "metadata": {
        "id": "OwrH1k1OYaD1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "**TASK 1**: Display some examples in your language to get familiar with the WikiAnn data. Write down some examples and your observations."
      ],
      "metadata": {
        "id": "52f-Gm5bF41i"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "example_no = 405\n",
        "dataset.data['train']['tokens'][example_no]"
      ],
      "metadata": {
        "id": "DdBWVGwGZLXi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "dataset.data['train']['ner_tags'][example_no]"
      ],
      "metadata": {
        "id": "30ycCdcwfyDz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "dataset['train'][example_no]"
      ],
      "metadata": {
        "id": "LDVEk2Kwkhjm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Next, we have to use *the same tokenizer* as for the pretrained model. Different tokenizers can split sentences in different ways but we need the data to be split exactly the same way it is in the pretrained model."
      ],
      "metadata": {
        "id": "_5Vz3KeVGlJt"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from transformers import AutoTokenizer\n",
        "tokenizer = AutoTokenizer.from_pretrained(\"bert-base-multilingual-cased\")\n"
      ],
      "metadata": {
        "id": "FD_U6dMtZh5e"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "**TASK 2**: Check tokenization on few sentences in your language. Write down your observations."
      ],
      "metadata": {
        "id": "xyHF_msoG82n"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "text = \"JA som tu! Bývám v Liptovskom Mikuláši.\"\n",
        "tokenized = tokenizer(text)\n",
        "tokenized"
      ],
      "metadata": {
        "id": "V2EEYO6SZyH_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "tokenizer.tokenize(text)"
      ],
      "metadata": {
        "id": "TA_hOGIAyYkY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The tokens are converted to token IDs, and these are converted to tensors.\n",
        "\n",
        "We can see the tokens are often smaller units than words. However, we have NER tags for words. The next function spreads the token class (the NER tag) for all subwords of a token.\n",
        "\n",
        "The code is copied from https://www.freecodecamp.org/news/getting-started-with-ner-models-using-huggingface/"
      ],
      "metadata": {
        "id": "baqEwuB1HJvP"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def tokenize_adjust_labels(all_samples_per_split):\n",
        "  tokenized_samples = tokenizer.batch_encode_plus(all_samples_per_split[\"tokens\"], is_split_into_words=True, max_length=50)\n",
        "\n",
        "  #tokenized_samples is not a datasets object so this alone won't work with Trainer API, hence map is used\n",
        "  #so the new keys [input_ids, labels (after adjustment)]\n",
        "  #can be added to the datasets dict for each train test validation split\n",
        "  total_adjusted_labels = []\n",
        "  print(len(tokenized_samples[\"input_ids\"]))\n",
        "  for k in range(0, len(tokenized_samples[\"input_ids\"])):\n",
        "    prev_wid = -1\n",
        "    word_ids_list = tokenized_samples.word_ids(batch_index=k)\n",
        "    existing_label_ids = all_samples_per_split[\"ner_tags\"][k]\n",
        "    i = -1\n",
        "    adjusted_label_ids = []\n",
        "\n",
        "    for wid in word_ids_list:\n",
        "      if(wid is None):\n",
        "        adjusted_label_ids.append(-100)\n",
        "      elif(wid!=prev_wid):\n",
        "        i = i + 1\n",
        "        adjusted_label_ids.append(existing_label_ids[i])\n",
        "        prev_wid = wid\n",
        "      else:\n",
        "        label_name = label_names[existing_label_ids[i]]\n",
        "        adjusted_label_ids.append(existing_label_ids[i])\n",
        "\n",
        "    total_adjusted_labels.append(adjusted_label_ids)\n",
        "  tokenized_samples[\"labels\"] = total_adjusted_labels\n",
        "  return tokenized_samples\n",
        "\n",
        "tokenized_dataset = dataset.map(tokenize_adjust_labels, batched=True)"
      ],
      "metadata": {
        "id": "RfnzuDNOh_5y"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The method above adds `input_ids`, `token_type_ids` and `attention_mask` from the tokenizer to the dataset. These will be used later by the trainer. In contrast, `ner_tags`, `langs`, `tokens`, and `spans` won't be used."
      ],
      "metadata": {
        "id": "Br3CIPkbIVtJ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "example_no = 1\n",
        "dataset['train'][example_no].keys(), tokenized_dataset['train'][example_no].keys()"
      ],
      "metadata": {
        "id": "_8U7Z6KliXqb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "For training, we need the samples to be the same length. A simple padding can be done by the tokenizer itself. Example code:"
      ],
      "metadata": {
        "id": "mKexuW4kH1EN"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "batch_sentences = [\n",
        "    \"But what about second breakfast?\",\n",
        "    \"Don't think he knows about second breakfast, Pip.\",\n",
        "    \"What about elevensies?\",\n",
        "]\n",
        "encoded_input = tokenizer(batch_sentences, padding=True)\n",
        "print(encoded_input)"
      ],
      "metadata": {
        "id": "SEmASRFzIA8U"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Data Collator preprocessess the input data into batches, optionally using methods such as padding and truncation. By default the `DataCollatorForTokenClassification` pads samples to `max_length`. Depending on your data and available memory, it may be useful to set `max_length` as a parameter for the `DataCollatorForTokenClassification`. See documentation for more details: https://huggingface.co/docs/transformers/main_classes/data_collator#transformers.DataCollatorForTokenClassification"
      ],
      "metadata": {
        "id": "E7FEvqXRDl92"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from transformers import DataCollatorForTokenClassification\n",
        "\n",
        "data_collator = DataCollatorForTokenClassification(tokenizer)"
      ],
      "metadata": {
        "id": "r5CqBtRfizti"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Next, we load the BERT model. For languages different from English, the multilingual model is most suitable. For scripts that distinguish upper case and lower case, the cased model is more suitable. We will load model weights, so that further training does not require much data. Moreover, even though the base model does not contain NER classification, we can benefit from the weights for the NER classification task."
      ],
      "metadata": {
        "id": "TB-BAykhOTEw"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from transformers import AutoModelForTokenClassification, TrainingArguments, Trainer"
      ],
      "metadata": {
        "id": "-4C3DCpYCADW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model = AutoModelForTokenClassification.from_pretrained(\"bert-base-multilingual-cased\",\n",
        "                                                        num_labels=len(label_names))"
      ],
      "metadata": {
        "id": "_9Ej6579C7kd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Next step is to set the metric. Since NER is a task on sequences, we use `sequeval` for evaluation. The `sequeval` is able to work with different IOB schemata. By default, it evaluates incorrect B-tags and I-tags as true positives (e.g. York is predicted as I-LOC but it should be B-LOC - still it is considered correct). More information is available here: https://huggingface.co/spaces/evaluate-metric/seqeval\n",
        "\n",
        "The method below considers the boundary tags (-100), plus it flattens the results to make them more readable. The code is copied from https://www.freecodecamp.org/news/getting-started-with-ner-models-using-huggingface/"
      ],
      "metadata": {
        "id": "MS5lkLvHJAq8"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "from datasets import load_metric\n",
        "metric = load_metric(\"seqeval\")\n",
        "def compute_metrics(p):\n",
        "    predictions, labels = p\n",
        "    predictions = np.argmax(predictions, axis=2)\n",
        "\n",
        "    # Remove ignored index (special tokens)\n",
        "    true_predictions = [\n",
        "        [label_names[p] for (p, l) in zip(prediction, label) if l != -100]\n",
        "        for prediction, label in zip(predictions, labels)\n",
        "    ]\n",
        "    true_labels = [\n",
        "        [label_names[l] for (p, l) in zip(prediction, label) if l != -100]\n",
        "        for prediction, label in zip(predictions, labels)\n",
        "    ]\n",
        "\n",
        "    results = metric.compute(predictions=true_predictions, references=true_labels)\n",
        "    flattened_results = {\n",
        "        \"overall_precision\": results[\"overall_precision\"],\n",
        "        \"overall_recall\": results[\"overall_recall\"],\n",
        "        \"overall_f1\": results[\"overall_f1\"],\n",
        "        \"overall_accuracy\": results[\"overall_accuracy\"],\n",
        "    }\n",
        "    for k in results.keys():\n",
        "      if(k not in flattened_results.keys()):\n",
        "        flattened_results[k+\"_f1\"]=results[k][\"f1\"]\n",
        "\n",
        "    return flattened_results"
      ],
      "metadata": {
        "id": "BzdRelLblvM0"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "##Training\n",
        "Run the training with the below parameters.\n",
        "**TASK 3**: Observe the training step results. Does the model improve? We train only in one epoch. What would you expect if the number of epochs increases?"
      ],
      "metadata": {
        "id": "t3gcntinPolE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "training_args = TrainingArguments(\n",
        "    output_dir=\"./fine_tune_bert_output\",\n",
        "    evaluation_strategy=\"steps\",\n",
        "    learning_rate=2e-5,\n",
        "    per_device_train_batch_size=2,\n",
        "    per_device_eval_batch_size=2,\n",
        "    num_train_epochs=1,\n",
        "    weight_decay=0.01,\n",
        "    logging_steps=1000,\n",
        ")"
      ],
      "metadata": {
        "id": "e-YvBlyxgUDN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "trainer = Trainer(\n",
        "    model=model,\n",
        "    args=training_args,\n",
        "    train_dataset=tokenized_dataset[\"train\"],\n",
        "    eval_dataset=tokenized_dataset[\"validation\"],\n",
        "    compute_metrics=compute_metrics,\n",
        "    data_collator=data_collator,\n",
        "    tokenizer=tokenizer,\n",
        ")"
      ],
      "metadata": {
        "id": "_laBeyBngXSJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "trainer.train()"
      ],
      "metadata": {
        "id": "9QJK2WeFhMK_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "out_dir = './bert_ner'\n",
        "trainer.save_model(out_dir)"
      ],
      "metadata": {
        "id": "E-1VaguisuyM"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "##Predictions\n",
        "Load the fine-tuned model into the pipeline and run prediction on sentences in your language.\n",
        "\n",
        "**TASK 4**: Put down some observations. Where does the model perform well? How does it deal with rare or OOV words?"
      ],
      "metadata": {
        "id": "P56tSzwaPTId"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import transformers"
      ],
      "metadata": {
        "id": "aV8XGEAp0moz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "token_classifier = transformers.pipeline(\n",
        "    \"token-classification\", model=\"./bert_ner\", aggregation_strategy=\"first\"\n",
        ")"
      ],
      "metadata": {
        "id": "YsFgiPdC1k4Q"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "label_names"
      ],
      "metadata": {
        "id": "Vv6Qu29rJ3ZK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "text = \"Ja som Katka z Blavy. Bola som na Slovensku, no teraz som v Brne.\"\n",
        "token_classifier(text)"
      ],
      "metadata": {
        "id": "ff-ormyL2ggR"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}