experiment_ner_knowledge_transfer.ipynb

{
 "cells": [
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-07T16:44:16.164093Z",
     "start_time": "2024-12-07T16:44:11.646501Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from datetime import datetime\n",
    "\n",
    "from datasets import load_dataset, load_from_disk\n",
    "from transformers import BatchEncoding, PreTrainedTokenizer, AutoTokenizer, Trainer, TrainingArguments\n",
    "from transformers.data import data_collator\n",
    "\n",
    "from modelling_xlm_roberta import XLMRobertaForTokenClassification\n",
    "import nervaluate\n",
    "\n",
    "from functools import partial\n",
    "import torch\n",
    "\n",
    "from typing import Iterable\n",
    "from torch import Tensor\n",
    "from torch.nn.utils.rnn import pad_sequence\n",
    "\n",
    "from sklearn.metrics import precision_score, recall_score, f1_score\n",
    "\n",
    "import numpy as np\n",
    "import wandb\n",
    "\n",
    "model_name = 'facebook/xlm-v-base' # 'FacebookAI/xlm-roberta-base' \n",
    "device = 'cuda'\n",
    "model_dtype = torch.bfloat16\n",
    "torch.cuda.get_device_name(0)\n",
    "\n",
    "lr = 2e-5\n",
    "steps = 2500\n",
    "batch_size = 16"
   ],
   "id": "164bfe6522ca666c",
   "outputs": [],
   "execution_count": 1
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "# 1. Test that layer cutting works",
   "id": "bd4e31a4b6b090d1"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-07T16:44:16.639988Z",
     "start_time": "2024-12-07T16:44:16.165126Z"
    }
   },
   "cell_type": "code",
   "source": [
    "model_test = XLMRobertaForTokenClassification.from_pretrained(model_name)\n",
    "model_test"
   ],
   "id": "50cd8944790c35f2",
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of XLMRobertaForTokenClassification were not initialized from the model checkpoint at facebook/xlm-v-base and are newly initialized: ['classifier.bias', 'classifier.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "XLMRobertaForTokenClassification(\n",
       "  (roberta): XLMRobertaModel(\n",
       "    (embeddings): XLMRobertaEmbeddings(\n",
       "      (word_embeddings): Embedding(901629, 768, padding_idx=1)\n",
       "      (position_embeddings): Embedding(514, 768, padding_idx=1)\n",
       "      (token_type_embeddings): Embedding(1, 768)\n",
       "      (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (dropout): Dropout(p=0.1, inplace=False)\n",
       "    )\n",
       "    (encoder): XLMRobertaEncoder(\n",
       "      (layer): ModuleList(\n",
       "        (0-11): 12 x XLMRobertaLayer(\n",
       "          (attention): XLMRobertaAttention(\n",
       "            (self): XLMRobertaSdpaSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): XLMRobertaSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): XLMRobertaIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            (intermediate_act_fn): GELUActivation()\n",
       "          )\n",
       "          (output): XLMRobertaOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (dropout): Dropout(p=0.1, inplace=False)\n",
       "  (classifier): Linear(in_features=768, out_features=2, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 2
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-07T16:44:17.060793Z",
     "start_time": "2024-12-07T16:44:16.640723Z"
    }
   },
   "cell_type": "code",
   "source": [
    "model_test = XLMRobertaForTokenClassification.from_pretrained(model_name, skip_last_layer=True)\n",
    "model_test"
   ],
   "id": "45dbf4a4718b4f0a",
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of XLMRobertaForTokenClassification were not initialized from the model checkpoint at facebook/xlm-v-base and are newly initialized: ['classifier.bias', 'classifier.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "XLMRobertaForTokenClassification(\n",
       "  (roberta): XLMRobertaModel(\n",
       "    (embeddings): XLMRobertaEmbeddings(\n",
       "      (word_embeddings): Embedding(901629, 768, padding_idx=1)\n",
       "      (position_embeddings): Embedding(514, 768, padding_idx=1)\n",
       "      (token_type_embeddings): Embedding(1, 768)\n",
       "      (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (dropout): Dropout(p=0.1, inplace=False)\n",
       "    )\n",
       "    (encoder): XLMRobertaEncoder(\n",
       "      (layer): ModuleList(\n",
       "        (0-10): 11 x XLMRobertaLayer(\n",
       "          (attention): XLMRobertaAttention(\n",
       "            (self): XLMRobertaSdpaSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): XLMRobertaSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): XLMRobertaIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            (intermediate_act_fn): GELUActivation()\n",
       "          )\n",
       "          (output): XLMRobertaOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (dropout): Dropout(p=0.1, inplace=False)\n",
       "  (classifier): Linear(in_features=768, out_features=2, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 3
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Works! Passing `skip_last_layer=True` removes the last layer in the transformer stack (11 x XLMRobertaLayer instead of 12 x XLMRobertaLayer)",
   "id": "d91e348b9d48ece2"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "# 2. Train models on the downstream tagging task and evaluate the knowledge transfer to a different language\n",
    "For this we will use CoNLL 2003 corpus (`eriktks/conll2003`, 14k examples) to train the model and Afrikaans NER Corpus (`nwu-ctext/afrikaans_ner_corpus`, 9k examples) to test the model. The validation is done over CoNLL 2003, only the final scores for Afrikaans are reported."
   ],
   "id": "9af1cd678fa696d8"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-07T16:44:21.919723Z",
     "start_time": "2024-12-07T16:44:17.061546Z"
    }
   },
   "cell_type": "code",
   "source": [
    "train_dataset = load_dataset('eriktks/conll2003', split='train')\n",
    "valid_dataset = load_dataset('eriktks/conll2003', split='validation')\n",
    "test_dataset = load_dataset('nwu-ctext/afrikaans_ner_corpus', split='train')"
   ],
   "id": "19ee1cb470c3ca97",
   "outputs": [],
   "execution_count": 4
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Make sure that the labelling scheme is identical across datasets",
   "id": "ecafcb621612d84f"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-07T16:44:21.927236Z",
     "start_time": "2024-12-07T16:44:21.922484Z"
    }
   },
   "cell_type": "code",
   "source": "train_dataset.features['ner_tags']",
   "id": "668a8c6cfaac7811",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequence(feature=ClassLabel(names=['O', 'B-PER', 'I-PER', 'B-ORG', 'I-ORG', 'B-LOC', 'I-LOC', 'B-MISC', 'I-MISC'], id=None), length=-1, id=None)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 5
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-07T16:44:21.932457Z",
     "start_time": "2024-12-07T16:44:21.928530Z"
    }
   },
   "cell_type": "code",
   "source": "valid_dataset.features['ner_tags']",
   "id": "432d9997a6a8bcf1",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequence(feature=ClassLabel(names=['O', 'B-PER', 'I-PER', 'B-ORG', 'I-ORG', 'B-LOC', 'I-LOC', 'B-MISC', 'I-MISC'], id=None), length=-1, id=None)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 6
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-07T16:44:21.937232Z",
     "start_time": "2024-12-07T16:44:21.933576Z"
    }
   },
   "cell_type": "code",
   "source": "test_dataset.features['ner_tags']",
   "id": "8c160d6214600da3",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequence(feature=ClassLabel(names=['OUT', 'B-PERS', 'I-PERS', 'B-ORG', 'I-ORG', 'B-LOC', 'I-LOC', 'B-MISC', 'I-MISC'], id=None), length=-1, id=None)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 7
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "The names are a bit different, but otherwise the schemes are identical",
   "id": "d61fbac113d2c666"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "## 2.1 Convert word-level tags to subtoken-level tags",
   "id": "2e089cee56f48462"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-07T16:44:25.752172Z",
     "start_time": "2024-12-07T16:44:21.938505Z"
    }
   },
   "cell_type": "code",
   "source": [
    "tok = AutoTokenizer.from_pretrained(model_name)\n",
    "tok_name = model_name.replace('/', '__')\n",
    "\n",
    "tok('test <mask> test', return_offsets_mapping=True)"
   ],
   "id": "7c3f74ba028a74e7",
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/mloscratch/homes/shcherba/conda/envs/char-llm/lib/python3.11/site-packages/transformers/tokenization_utils_base.py:1601: FutureWarning: `clean_up_tokenization_spaces` was not set. It will be set to `True` by default. This behavior will be depracted in transformers v4.45, and will be then set to `False` by default. For more details check this issue: https://github.com/huggingface/transformers/issues/31884\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'input_ids': [0, 1340, 901628, 1340, 2], 'attention_mask': [1, 1, 1, 1, 1], 'offset_mapping': [(0, 0), (0, 4), (4, 11), (11, 16), (0, 0)]}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 8
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-07T16:44:25.755644Z",
     "start_time": "2024-12-07T16:44:25.753131Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# for reference\n",
    "ner_tags_scheme = np.array(['O', 'B-PER', 'I-PER', 'B-ORG', 'I-ORG', 'B-LOC', 'I-LOC', 'B-MISC', 'I-MISC'])\n",
    "ner_tags_ext    =          [  0,       2,       2,       4,       4,       6,       6,        8,        8]\n",
    "# the ext is used when we need to split one word into multiple sub tokens"
   ],
   "id": "cfb59f07b44af84f",
   "outputs": [],
   "execution_count": 9
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-07T16:44:25.940260Z",
     "start_time": "2024-12-07T16:44:25.756287Z"
    }
   },
   "cell_type": "code",
   "source": [
    "def tokenize(example: dict, tokenizer: PreTrainedTokenizer, tokenizer_name: str, max_length: int = 512) -> dict:\n",
    "    ner_tags: list[int] = example['ner_tags']\n",
    "    example_words: list[str] = example['tokens']\n",
    "    text = ' '.join(example_words)\n",
    "    \n",
    "    # map words to positions in text\n",
    "    word_positions: list[int] = example.get('word_positions', [])\n",
    "    \n",
    "    if len(word_positions) != len(example_words):\n",
    "        text_iterator = 0\n",
    "        for word in example_words:\n",
    "            while text[text_iterator:text_iterator + len(word)] != word:\n",
    "                text_iterator += 1\n",
    "                assert text_iterator < len(text)\n",
    "            \n",
    "            word_positions.append(text_iterator)\n",
    "    \n",
    "    encoding: BatchEncoding = tokenizer(text, return_offsets_mapping=True, truncation=True, max_length=max_length)\n",
    "    num_sub_tokens = len(encoding.offset_mapping)\n",
    "    \n",
    "    sub_token_iterator = 0\n",
    "    sub_token_ner_tags: list[int] = []\n",
    "    for word_id, ner_tag in enumerate(ner_tags):\n",
    "        word_start = word_positions[word_id]\n",
    "        word_end = word_start + len(example_words[word_id])\n",
    "        \n",
    "        # there may be some empty space between words. the sub tokens that include this empty space receive O label\n",
    "        # we compare with the end ([1]) to ensure that 0-length tokens are labelled as -100 (for example <CLS>)\n",
    "        while sub_token_iterator < num_sub_tokens and  encoding.offset_mapping[sub_token_iterator][1] <= word_start:\n",
    "            if encoding.offset_mapping[sub_token_iterator][1] - encoding.offset_mapping[sub_token_iterator][0] == 0:\n",
    "                # set to -100 for special tokens like <CLS>\n",
    "                sub_token_ner_tags.append(-100)\n",
    "            else:\n",
    "                sub_token_ner_tags.append(0)  # 0 = O\n",
    "            sub_token_iterator += 1\n",
    "            \n",
    "        ext_tag = ner_tags_ext[ner_tag]\n",
    "        \n",
    "        if sub_token_iterator < num_sub_tokens:\n",
    "            # the first sub token of a word receives original label, the rest receive extended label\n",
    "            sub_token_ner_tags.append(ner_tag)\n",
    "            sub_token_iterator += 1\n",
    "        \n",
    "        # again, we need to be careful about 0-length tokens, so we compare start ([0]) with the word end\n",
    "        while sub_token_iterator < num_sub_tokens and encoding.offset_mapping[sub_token_iterator][0] < word_end:\n",
    "            \n",
    "            # there is a weird quirk with transformers tokenizers: <SEP> token has (0, 0) offset \n",
    "            #   regardless of its real position, see https://github.com/huggingface/transformers/issues/35125\n",
    "            if encoding.offset_mapping[sub_token_iterator][1] - encoding.offset_mapping[sub_token_iterator][0] == 0:\n",
    "                sub_token_ner_tags.append(-100)\n",
    "            else:\n",
    "                sub_token_ner_tags.append(ext_tag)\n",
    "                \n",
    "            sub_token_iterator += 1\n",
    "    \n",
    "    # any tokens at the end (like <SEP>) receive O tokens\n",
    "    while sub_token_iterator < num_sub_tokens:\n",
    "        sub_token_iterator += 1\n",
    "        sub_token_ner_tags.append(0)\n",
    "        \n",
    "    return {\n",
    "        'word_positions': word_positions,\n",
    "        f'{tokenizer_name}_sub_tokens': encoding.input_ids,\n",
    "        f'{tokenizer_name}_sub_token_offsets': encoding.offset_mapping,\n",
    "        f'{tokenizer_name}_sub_token_ner_tags': sub_token_ner_tags,\n",
    "        'length': len(encoding.offset_mapping)\n",
    "    }\n",
    "\n",
    "tokenize_fn = partial(tokenize, tokenizer=tok, tokenizer_name=tok_name, max_length=512)\n",
    "\n",
    "train_dataset = train_dataset.map(tokenize_fn)\n",
    "valid_dataset = valid_dataset.map(tokenize_fn)\n",
    "test_dataset = test_dataset.map(tokenize_fn)"
   ],
   "id": "21f16632e08f1c6a",
   "outputs": [],
   "execution_count": 10
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-07T16:44:25.957711Z",
     "start_time": "2024-12-07T16:44:25.941298Z"
    }
   },
   "cell_type": "code",
   "source": [
    "for test_idx in range(25):\n",
    "    ner_tags = torch.as_tensor(train_dataset[test_idx][f'{tok_name}_sub_token_ner_tags'])\n",
    "    tokens = torch.as_tensor(train_dataset[test_idx][f'{tok_name}_sub_tokens'])\n",
    "    print('Text:', ' '.join(train_dataset[test_idx]['tokens']))\n",
    "    print('Ents:', tok.decode(tokens[ner_tags > 0]))\n",
    "    print()"
   ],
   "id": "d1bc289c17bf712c",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Text: EU rejects German call to boycott British lamb .\n",
      "Ents: EU German British\n",
      "\n",
      "Text: Peter Blackburn\n",
      "Ents: Peter Blackburn\n",
      "\n",
      "Text: BRUSSELS 1996-08-22\n",
      "Ents: BRUSSELS\n",
      "\n",
      "Text: The European Commission said on Thursday it disagreed with German advice to consumers to shun British lamb until scientists determine whether mad cow disease can be transmitted to sheep .\n",
      "Ents: European Commission German British\n",
      "\n",
      "Text: Germany 's representative to the European Union 's veterinary committee Werner Zwingmann said on Wednesday consumers should buy sheepmeat from countries other than Britain until the scientific advice was clearer .\n",
      "Ents: Germany European Union Werner Zwingmann Britain\n",
      "\n",
      "Text: \" We do n't support any such recommendation because we do n't see any grounds for it , \" the Commission 's chief spokesman Nikolaus van der Pas told a news briefing .\n",
      "Ents: Commission Nikolaus van der Pas\n",
      "\n",
      "Text: He said further scientific study was required and if it was found that action was needed it should be taken by the European Union .\n",
      "Ents: European Union\n",
      "\n",
      "Text: He said a proposal last month by EU Farm Commissioner Franz Fischler to ban sheep brains , spleens and spinal cords from the human and animal food chains was a highly specific and precautionary move to protect human health .\n",
      "Ents: EU Franz Fischler\n",
      "\n",
      "Text: Fischler proposed EU-wide measures after reports from Britain and France that under laboratory conditions sheep could contract Bovine Spongiform Encephalopathy ( BSE ) -- mad cow disease .\n",
      "Ents: Fischler EU-wide Britain France Bovine Spongiform Encephalopathy BSE\n",
      "\n",
      "Text: But Fischler agreed to review his proposal after the EU 's standing veterinary committee , mational animal health officials , questioned if such action was justified as there was only a slight risk to human health .\n",
      "Ents: Fischler EU\n",
      "\n",
      "Text: Spanish Farm Minister Loyola de Palacio had earlier accused Fischler at an EU farm ministers ' meeting of causing unjustified alarm through \" dangerous generalisation . \"\n",
      "Ents: Spanish Loyola de Palacio Fischler EU\n",
      "\n",
      "Text: .\n",
      "Ents: \n",
      "\n",
      "Text: Only France and Britain backed Fischler 's proposal .\n",
      "Ents: France Britain Fischler\n",
      "\n",
      "Text: The EU 's scientific veterinary and multidisciplinary committees are due to re-examine the issue early next month and make recommendations to the senior veterinary officials .\n",
      "Ents: EU\n",
      "\n",
      "Text: Sheep have long been known to contract scrapie , a brain-wasting disease similar to BSE which is believed to have been transferred to cattle through feed containing animal waste .\n",
      "Ents: BSE\n",
      "\n",
      "Text: British farmers denied on Thursday there was any danger to human health from their sheep , but expressed concern that German government advice to consumers to avoid British lamb might influence consumers across Europe .\n",
      "Ents: British German British Europe\n",
      "\n",
      "Text: \" What we have to be extremely careful of is how other countries are going to take Germany 's lead , \" Welsh National Farmers ' Union ( NFU ) chairman John Lloyd Jones said on BBC radio .\n",
      "Ents: Germany Welsh National Farmers'Union NFU John Lloyd Jones BBC radio\n",
      "\n",
      "Text: Bonn has led efforts to protect public health after consumer confidence collapsed in March after a British report suggested humans could contract an illness similar to mad cow disease by eating contaminated beef .\n",
      "Ents: Bonn British\n",
      "\n",
      "Text: Germany imported 47,600 sheep from Britain last year , nearly half of total imports .\n",
      "Ents: Germany Britain\n",
      "\n",
      "Text: It brought in 4,275 tonnes of British mutton , some 10 percent of overall imports .\n",
      "Ents: British\n",
      "\n",
      "Text: Rare Hendrix song draft sells for almost $ 17,000 .\n",
      "Ents: Hendrix\n",
      "\n",
      "Text: LONDON 1996-08-22\n",
      "Ents: LONDON\n",
      "\n",
      "Text: A rare early handwritten draft of a song by U.S. guitar legend Jimi Hendrix was sold for almost $ 17,000 on Thursday at an auction of some of the late musician 's favourite possessions .\n",
      "Ents: U.S. Jimi Hendrix\n",
      "\n",
      "Text: A Florida restaurant paid 10,925 pounds ( $ 16,935 ) for the draft of \" Ai n't no telling \" , which Hendrix penned on a piece of London hotel stationery in late 1966 .\n",
      "Ents: Florida Ain't no telling Hendrix London\n",
      "\n",
      "Text: At the end of a January 1967 concert in the English city of Nottingham he threw the sheet of paper into the audience , where it was retrieved by a fan .\n",
      "Ents: English Nottingham\n",
      "\n"
     ]
    }
   ],
   "execution_count": 11
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Looks nice!",
   "id": "9234e8f0ab28fb5"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-07T16:44:26.158547Z",
     "start_time": "2024-12-07T16:44:25.958585Z"
    }
   },
   "cell_type": "code",
   "source": [
    "train_dataset.save_to_disk('data/train')\n",
    "valid_dataset.save_to_disk('data/valid')\n",
    "test_dataset.save_to_disk('data/test')"
   ],
   "id": "8d7920ea056c2596",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Saving the dataset (0/1 shards):   0%|          | 0/14041 [00:00<?, ? examples/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "b5fd769451264c6a9c55b215b6f8189e"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "Saving the dataset (0/1 shards):   0%|          | 0/3250 [00:00<?, ? examples/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "f69f2ddb7e3444ac8c2ae14feba0a23b"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "Saving the dataset (0/1 shards):   0%|          | 0/8962 [00:00<?, ? examples/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "5eca764b2f9d4e2ca9f07972319ecf7f"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "execution_count": 12
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-07T16:44:26.163399Z",
     "start_time": "2024-12-07T16:44:26.159588Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class Dataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, examples: Iterable[dict], tokenizer_name: str):\n",
    "        self.input_ids = []\n",
    "        self.labels = []\n",
    "        \n",
    "        for example in examples:\n",
    "            self.input_ids.append(torch.as_tensor(example[f'{tokenizer_name}_sub_tokens']))\n",
    "            self.labels.append(torch.as_tensor(example[f'{tokenizer_name}_sub_token_ner_tags']))\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.input_ids)\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        return self.input_ids[idx], self.labels[idx]\n",
    "    \n",
    "\n",
    "def collate_fn(inputs: list[(Tensor, Tensor)], *, pad_token: int) -> dict:\n",
    "    all_input_ids = []\n",
    "    all_labels = []\n",
    "    for input_ids, labels in inputs:\n",
    "        all_input_ids.append(input_ids)\n",
    "        all_labels.append(labels)\n",
    "    \n",
    "    input_ids = pad_sequence(all_input_ids, batch_first=True, padding_value=pad_token)\n",
    "    \n",
    "    batch_size, seq_length = input_ids.shape\n",
    "\n",
    "    # do not attend to pad and pad does not attend to anything\n",
    "    pad_mask = (input_ids != pad_token).long()\n",
    "    return {\n",
    "        'input_ids': input_ids,\n",
    "        'labels': pad_sequence(all_labels, batch_first=True, padding_value=-100),\n",
    "        'attention_mask': pad_mask\n",
    "    }"
   ],
   "id": "2bbb670d6122a3f5",
   "outputs": [],
   "execution_count": 13
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-07T16:44:26.167670Z",
     "start_time": "2024-12-07T16:44:26.164217Z"
    }
   },
   "cell_type": "code",
   "source": [
    "def compute_ner_metrics(eval_pred) -> dict:\n",
    "    predictions, labels = eval_pred\n",
    "    \n",
    "    predictions = np.argmax(predictions, axis=-1)\n",
    "    padding = (labels < 0)\n",
    "    \n",
    "    predictions = predictions[~padding]\n",
    "    labels = labels[~padding]\n",
    "\n",
    "    predictions = ner_tags_scheme[predictions]\n",
    "    labels = ner_tags_scheme[labels]\n",
    "    \n",
    "    token_precision = precision_score(labels, predictions, average='macro', zero_division=0)\n",
    "    token_recall = recall_score(labels, predictions, average='macro', zero_division=0)\n",
    "    token_f1 = f1_score(labels, predictions, average='macro', zero_division=0)\n",
    "\n",
    "    evaluator = nervaluate.Evaluator([labels], [predictions], tags=['PER', 'LOC', 'ORG', 'MISC'], loader='list')\n",
    "    results, results_per_tag, _, _ = evaluator.evaluate()\n",
    "\n",
    "    overall_metrics = results['strict']\n",
    "    \n",
    "    metrics = {\n",
    "        'token_precision_macro': token_precision,\n",
    "        'token_recall_macro': token_recall,\n",
    "        'token_f1_macro': token_f1,\n",
    "        'overall_precision': overall_metrics['precision'],\n",
    "        'overall_recall': overall_metrics['recall'],\n",
    "        'overall_f1': overall_metrics['f1'],\n",
    "    }\n",
    "    \n",
    "    for tag, tag_metrics in results_per_tag.items():\n",
    "        metrics[f'{tag}_precision'] = tag_metrics['strict']['precision']\n",
    "        metrics[f'{tag}_recall'] = tag_metrics['strict']['recall']\n",
    "        metrics[f'{tag}_f1'] = tag_metrics['strict']['f1']\n",
    "\n",
    "    # Return desired metrics\n",
    "    return metrics"
   ],
   "id": "3da5ff8b55ce26e",
   "outputs": [],
   "execution_count": 14
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "## 2.2 Train a conventional model",
   "id": "c9559c9f274e9c04"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-07T16:44:26.170227Z",
     "start_time": "2024-12-07T16:44:26.168481Z"
    }
   },
   "cell_type": "code",
   "source": "n_run = 0",
   "id": "4a4988357b4e9d66",
   "outputs": [],
   "execution_count": 15
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-07T16:47:22.489711Z",
     "start_time": "2024-12-07T16:44:26.170931Z"
    }
   },
   "cell_type": "code",
   "source": [
    "run_name = f'{tok_name}-finetuned-l12-conll03/{datetime.now().strftime(\"%m-%d\")}/{n_run}'\n",
    "wandb.init(\n",
    "    project='ner-alignment',\n",
    "    name=run_name,\n",
    "    dir=run_name,\n",
    "    resume=False\n",
    ")\n",
    "n_run += 1\n",
    "\n",
    "model = XLMRobertaForTokenClassification.from_pretrained(\n",
    "    model_name, \n",
    "    num_labels=9,\n",
    ")\n",
    "model.roberta.embeddings.requires_grad_(False)  # freeze input embeddings to avoid parameter shift (training on english and inferencing on africaans -> different tokens are activated)\n",
    "print(f\"Percentage of frozen modules: {100 * sum(1 for module in model.modules() if not any(p.requires_grad for p in module.parameters())) / sum(1 for module in model.modules()):.2f}%\")\n",
    "print(f\"Percentage of frozen parameters: {100 * sum(p.numel() for p in model.parameters() if not p.requires_grad) / sum(p.numel() for p in model.parameters()):.2f}%\")\n",
    "\n",
    "\n",
    "tok = AutoTokenizer.from_pretrained(model_name)\n",
    "\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=TrainingArguments(\n",
    "        output_dir=run_name,\n",
    "        overwrite_output_dir=True,\n",
    "        eval_strategy='steps',\n",
    "        per_device_train_batch_size=batch_size,\n",
    "        per_device_eval_batch_size=128,\n",
    "        learning_rate=lr,\n",
    "        max_steps=steps,\n",
    "        lr_scheduler_type='cosine_with_min_lr',\n",
    "        lr_scheduler_kwargs={ 'num_cycles': 0.5, 'min_lr_rate': 0.01 },\n",
    "        warmup_ratio=0.1,\n",
    "        adam_epsilon=1e-8,\n",
    "        adam_beta1=0.9,\n",
    "        adam_beta2=0.999,\n",
    "        weight_decay=0.0,\n",
    "        logging_steps=100,\n",
    "        eval_steps=200,\n",
    "        bf16=True,\n",
    "        torch_compile=False,\n",
    "        include_num_input_tokens_seen=True,\n",
    "        disable_tqdm=True,\n",
    "        report_to='wandb'\n",
    "    ),\n",
    "    data_collator=partial(collate_fn, pad_token=tok.pad_token_id),\n",
    "    train_dataset=Dataset(load_from_disk('data/train'), tokenizer_name=tok_name),\n",
    "    eval_dataset=Dataset(load_from_disk('data/valid'), tokenizer_name=tok_name),\n",
    "    compute_metrics=compute_ner_metrics\n",
    ")\n",
    "trainer.train()"
   ],
   "id": "c13273fe9f7aff09",
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001B[34m\u001B[1mwandb\u001B[0m: Currently logged in as: \u001B[33mviktoroo-sch\u001B[0m (\u001B[33mviktoroo-sch-epfl\u001B[0m). Use \u001B[1m`wandb login --relogin`\u001B[0m to force relogin\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ],
      "text/html": [
       "Tracking run with wandb version 0.17.9"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ],
      "text/html": [
       "Run data is saved locally in <code>facebook__xlm-v-base-finetuned-l12-conll03/12-07/0/wandb/run-20241207_164430-yaiig3k9</code>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ],
      "text/html": [
       "Syncing run <strong><a href='https://wandb.ai/viktoroo-sch-epfl/ner-alignment/runs/yaiig3k9' target=\"_blank\">facebook__xlm-v-base-finetuned-l12-conll03/12-07/0</a></strong> to <a href='https://wandb.ai/viktoroo-sch-epfl/ner-alignment' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ],
      "text/html": [
       " View project at <a href='https://wandb.ai/viktoroo-sch-epfl/ner-alignment' target=\"_blank\">https://wandb.ai/viktoroo-sch-epfl/ner-alignment</a>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ],
      "text/html": [
       " View run at <a href='https://wandb.ai/viktoroo-sch-epfl/ner-alignment/runs/yaiig3k9' target=\"_blank\">https://wandb.ai/viktoroo-sch-epfl/ner-alignment/runs/yaiig3k9</a>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of XLMRobertaForTokenClassification were not initialized from the model checkpoint at facebook/xlm-v-base and are newly initialized: ['classifier.bias', 'classifier.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Percentage of frozen modules: 24.12%\n",
      "Percentage of frozen parameters: 89.07%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/mloscratch/homes/shcherba/conda/envs/char-llm/lib/python3.11/site-packages/transformers/tokenization_utils_base.py:1601: FutureWarning: `clean_up_tokenization_spaces` was not set. It will be set to `True` by default. This behavior will be depracted in transformers v4.45, and will be then set to `False` by default. For more details check this issue: https://github.com/huggingface/transformers/issues/31884\n",
      "  warnings.warn(\n",
      "max_steps is given, it will override any value given in num_train_epochs\n",
      "\u001B[34m\u001B[1mwandb\u001B[0m: \u001B[33mWARNING\u001B[0m The `run_name` is currently set to the same value as `TrainingArguments.output_dir`. If this was not intended, please specify a different run name by setting the `TrainingArguments.run_name` parameter.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 1.6795, 'grad_norm': 8.110573768615723, 'learning_rate': 8.000000000000001e-06, 'epoch': 0.11389521640091116, 'num_input_tokens_seen': 76672}\n",
      "{'loss': 0.7058, 'grad_norm': 2.834045171737671, 'learning_rate': 1.6000000000000003e-05, 'epoch': 0.22779043280182232, 'num_input_tokens_seen': 158272}\n",
      "{'eval_loss': 0.5673161149024963, 'eval_token_precision_macro': 0.08892700414117066, 'eval_token_recall_macro': 0.1111111111111111, 'eval_token_f1_macro': 0.09878895554926144, 'eval_overall_precision': 0.0, 'eval_overall_recall': 0.0, 'eval_overall_f1': 0, 'eval_PER_precision': 0.0, 'eval_PER_recall': 0.0, 'eval_PER_f1': 0, 'eval_LOC_precision': 0, 'eval_LOC_recall': 0.0, 'eval_LOC_f1': 0, 'eval_ORG_precision': 0, 'eval_ORG_recall': 0.0, 'eval_ORG_f1': 0, 'eval_MISC_precision': 0, 'eval_MISC_recall': 0.0, 'eval_MISC_f1': 0, 'eval_runtime': 0.5682, 'eval_samples_per_second': 5719.996, 'eval_steps_per_second': 45.76, 'epoch': 0.22779043280182232, 'num_input_tokens_seen': 158272}\n",
      "{'loss': 0.5139, 'grad_norm': 0.7886868119239807, 'learning_rate': 1.997588409757226e-05, 'epoch': 0.3416856492027335, 'num_input_tokens_seen': 236576}\n",
      "{'loss': 0.3832, 'grad_norm': 2.3624935150146484, 'learning_rate': 1.9783661247264677e-05, 'epoch': 0.45558086560364464, 'num_input_tokens_seen': 313200}\n",
      "{'eval_loss': 0.2675582468509674, 'eval_token_precision_macro': 0.4717206185879596, 'eval_token_recall_macro': 0.5440138825078188, 'eval_token_f1_macro': 0.4835382908445392, 'eval_overall_precision': 0.3770282355947362, 'eval_overall_recall': 0.49604975626155656, 'eval_overall_f1': 0.4284262485481998, 'eval_PER_precision': 0.9072386058981233, 'eval_PER_recall': 0.9185667752442996, 'eval_PER_f1': 0.9128675478823848, 'eval_LOC_precision': 0.5461065573770492, 'eval_LOC_recall': 0.5802939575394666, 'eval_LOC_f1': 0.5626814462918976, 'eval_ORG_precision': 0.06073001887979861, 'eval_ORG_recall': 0.14360119047619047, 'eval_ORG_f1': 0.08536045997346306, 'eval_MISC_precision': 0.0, 'eval_MISC_recall': 0.0, 'eval_MISC_f1': 0, 'eval_runtime': 7.7884, 'eval_samples_per_second': 417.287, 'eval_steps_per_second': 3.338, 'epoch': 0.45558086560364464, 'num_input_tokens_seen': 313200}\n",
      "{'loss': 0.251, 'grad_norm': 1.4386305809020996, 'learning_rate': 1.9402956945780496e-05, 'epoch': 0.5694760820045558, 'num_input_tokens_seen': 390432}\n",
      "{'loss': 0.1878, 'grad_norm': 3.3927533626556396, 'learning_rate': 1.884118116930338e-05, 'epoch': 0.683371298405467, 'num_input_tokens_seen': 468592}\n",
      "{'eval_loss': 0.13975264132022858, 'eval_token_precision_macro': 0.8731010962483373, 'eval_token_recall_macro': 0.8242597065481019, 'eval_token_f1_macro': 0.8303982821329423, 'eval_overall_precision': 0.8033459595959596, 'eval_overall_recall': 0.855605984199025, 'eval_overall_f1': 0.8286528286528286, 'eval_PER_precision': 0.9377013963480129, 'eval_PER_recall': 0.9478827361563518, 'eval_PER_f1': 0.9427645788336934, 'eval_LOC_precision': 0.8471241170534813, 'eval_LOC_recall': 0.9139902014153511, 'eval_LOC_f1': 0.8792877716679759, 'eval_ORG_precision': 0.7669773635153129, 'eval_ORG_recall': 0.8571428571428571, 'eval_ORG_f1': 0.8095572733661278, 'eval_MISC_precision': 0.5181818181818182, 'eval_MISC_recall': 0.5539956803455723, 'eval_MISC_f1': 0.535490605427975, 'eval_runtime': 3.0171, 'eval_samples_per_second': 1077.2, 'eval_steps_per_second': 8.618, 'epoch': 0.683371298405467, 'num_input_tokens_seen': 468592}\n",
      "{'loss': 0.1468, 'grad_norm': 7.070267200469971, 'learning_rate': 1.8109268244311982e-05, 'epoch': 0.7972665148063781, 'num_input_tokens_seen': 549056}\n",
      "{'loss': 0.1209, 'grad_norm': 3.616713523864746, 'learning_rate': 1.7221464023352647e-05, 'epoch': 0.9111617312072893, 'num_input_tokens_seen': 626848}\n",
      "{'eval_loss': 0.09430964291095734, 'eval_token_precision_macro': 0.9116448156659415, 'eval_token_recall_macro': 0.9048737317819852, 'eval_token_f1_macro': 0.9066981420584929, 'eval_overall_precision': 0.879259501965924, 'eval_overall_recall': 0.9021684316691881, 'eval_overall_f1': 0.8905666639011034, 'eval_PER_precision': 0.9725067385444744, 'eval_PER_recall': 0.9793702497285559, 'eval_PER_f1': 0.9759264268325669, 'eval_LOC_precision': 0.8610515021459227, 'eval_LOC_recall': 0.8737071311921611, 'eval_LOC_f1': 0.8673331532018373, 'eval_ORG_precision': 0.8513037350246653, 'eval_ORG_recall': 0.8988095238095238, 'eval_ORG_f1': 0.8744118711545422, 'eval_MISC_precision': 0.7763975155279503, 'eval_MISC_recall': 0.8099352051835853, 'eval_MISC_f1': 0.7928118393234673, 'eval_runtime': 2.2401, 'eval_samples_per_second': 1450.846, 'eval_steps_per_second': 11.607, 'epoch': 0.9111617312072893, 'num_input_tokens_seen': 626848}\n",
      "{'loss': 0.1023, 'grad_norm': 2.4014155864715576, 'learning_rate': 1.619504860572402e-05, 'epoch': 1.0250569476082005, 'num_input_tokens_seen': 705647}\n",
      "{'loss': 0.084, 'grad_norm': 5.686334133148193, 'learning_rate': 1.505e-05, 'epoch': 1.1389521640091116, 'num_input_tokens_seen': 784495}\n",
      "{'eval_loss': 0.0750492587685585, 'eval_token_precision_macro': 0.9265061229662285, 'eval_token_recall_macro': 0.9306301160005456, 'eval_token_f1_macro': 0.9283809847257317, 'eval_overall_precision': 0.9047929572872514, 'eval_overall_recall': 0.9329299041855774, 'eval_overall_f1': 0.9186460316146652, 'eval_PER_precision': 0.9692224622030238, 'eval_PER_recall': 0.9744842562432139, 'eval_PER_f1': 0.9718462371413101, 'eval_LOC_precision': 0.9462480042575838, 'eval_LOC_recall': 0.9678824169842134, 'eval_LOC_f1': 0.9569429494079656, 'eval_ORG_precision': 0.8618980169971672, 'eval_ORG_recall': 0.9055059523809523, 'eval_ORG_f1': 0.8831640058055154, 'eval_MISC_precision': 0.7669021190716448, 'eval_MISC_recall': 0.8207343412526998, 'eval_MISC_f1': 0.792905581637976, 'eval_runtime': 1.9982, 'eval_samples_per_second': 1626.451, 'eval_steps_per_second': 13.012, 'epoch': 1.1389521640091116, 'num_input_tokens_seen': 784495}\n",
      "{'loss': 0.0803, 'grad_norm': 9.59705924987793, 'learning_rate': 1.3808605274817533e-05, 'epoch': 1.2528473804100229, 'num_input_tokens_seen': 861071}\n",
      "{'loss': 0.0747, 'grad_norm': 7.029209613800049, 'learning_rate': 1.2495026766436711e-05, 'epoch': 1.366742596810934, 'num_input_tokens_seen': 936543}\n",
      "{'eval_loss': 0.06733483821153641, 'eval_token_precision_macro': 0.9413803218694229, 'eval_token_recall_macro': 0.9334748044222985, 'eval_token_f1_macro': 0.9365496336950234, 'eval_overall_precision': 0.9235858418789282, 'eval_overall_recall': 0.938645150445453, 'eval_overall_f1': 0.9310546060858692, 'eval_PER_precision': 0.9718918918918918, 'eval_PER_recall': 0.9761129207383279, 'eval_PER_f1': 0.9739978331527628, 'eval_LOC_precision': 0.9477088948787062, 'eval_LOC_recall': 0.9569951007076756, 'eval_LOC_f1': 0.952329360780065, 'eval_ORG_precision': 0.8926270579813886, 'eval_ORG_recall': 0.9278273809523809, 'eval_ORG_f1': 0.9098869025902955, 'eval_MISC_precision': 0.8273305084745762, 'eval_MISC_recall': 0.8434125269978402, 'eval_MISC_f1': 0.8352941176470589, 'eval_runtime': 1.7465, 'eval_samples_per_second': 1860.899, 'eval_steps_per_second': 14.887, 'epoch': 1.366742596810934, 'num_input_tokens_seen': 936543}\n",
      "{'loss': 0.0761, 'grad_norm': 0.6913483738899231, 'learning_rate': 1.113483178634977e-05, 'epoch': 1.4806378132118452, 'num_input_tokens_seen': 1015087}\n",
      "{'loss': 0.0702, 'grad_norm': 4.288684844970703, 'learning_rate': 9.754494982645244e-06, 'epoch': 1.5945330296127562, 'num_input_tokens_seen': 1092351}\n",
      "{'eval_loss': 0.05995682626962662, 'eval_token_precision_macro': 0.9378648661331126, 'eval_token_recall_macro': 0.9392058246046988, 'eval_token_f1_macro': 0.9383677900123042, 'eval_overall_precision': 0.9218284200066247, 'eval_overall_recall': 0.9356194318372836, 'eval_overall_f1': 0.928672728789522, 'eval_PER_precision': 0.97670639219935, 'eval_PER_recall': 0.9788273615635179, 'eval_PER_f1': 0.9777657266811279, 'eval_LOC_precision': 0.95, 'eval_LOC_recall': 0.9618943930321175, 'eval_LOC_f1': 0.9559101974573978, 'eval_ORG_precision': 0.8584974471188913, 'eval_ORG_recall': 0.8757440476190477, 'eval_ORG_f1': 0.8670349907918968, 'eval_MISC_precision': 0.8522372528616025, 'eval_MISC_recall': 0.8844492440604752, 'eval_MISC_f1': 0.8680445151033387, 'eval_runtime': 1.7833, 'eval_samples_per_second': 1822.473, 'eval_steps_per_second': 14.58, 'epoch': 1.5945330296127562, 'num_input_tokens_seen': 1092351}\n",
      "{'loss': 0.0623, 'grad_norm': 2.1501197814941406, 'learning_rate': 8.380883041097391e-06, 'epoch': 1.7084282460136673, 'num_input_tokens_seen': 1172831}\n",
      "{'loss': 0.0618, 'grad_norm': 4.510129451751709, 'learning_rate': 7.0407317556880214e-06, 'epoch': 1.8223234624145785, 'num_input_tokens_seen': 1251359}\n",
      "{'eval_loss': 0.05515379458665848, 'eval_token_precision_macro': 0.9446877897882688, 'eval_token_recall_macro': 0.9460744553895584, 'eval_token_f1_macro': 0.9453217120865318, 'eval_overall_precision': 0.9279457940836225, 'eval_overall_recall': 0.9438561102706338, 'eval_overall_f1': 0.9358333333333333, 'eval_PER_precision': 0.9718462371413102, 'eval_PER_recall': 0.9744842562432139, 'eval_PER_f1': 0.9731634589319599, 'eval_LOC_precision': 0.9576861274772362, 'eval_LOC_recall': 0.9733260751224823, 'eval_LOC_f1': 0.9654427645788337, 'eval_ORG_precision': 0.8827785817655571, 'eval_ORG_recall': 0.9077380952380952, 'eval_ORG_f1': 0.8950843727072633, 'eval_MISC_precision': 0.8502617801047121, 'eval_MISC_recall': 0.8768898488120951, 'eval_MISC_f1': 0.863370547581074, 'eval_runtime': 1.6983, 'eval_samples_per_second': 1913.704, 'eval_steps_per_second': 15.31, 'epoch': 1.8223234624145785, 'num_input_tokens_seen': 1251359}\n",
      "{'loss': 0.0635, 'grad_norm': 4.409112930297852, 'learning_rate': 5.760125646788133e-06, 'epoch': 1.9362186788154898, 'num_input_tokens_seen': 1325695}\n",
      "{'loss': 0.0464, 'grad_norm': 2.0779221057891846, 'learning_rate': 4.563990255639608e-06, 'epoch': 2.050113895216401, 'num_input_tokens_seen': 1404999}\n",
      "{'eval_loss': 0.055309880524873734, 'eval_token_precision_macro': 0.9445593958027391, 'eval_token_recall_macro': 0.9479845214220121, 'eval_token_f1_macro': 0.9462424206348479, 'eval_overall_precision': 0.9317693705600528, 'eval_overall_recall': 0.9480584972264247, 'eval_overall_f1': 0.9398433594400933, 'eval_PER_precision': 0.9730167296276309, 'eval_PER_recall': 0.9788273615635179, 'eval_PER_f1': 0.975913396481732, 'eval_LOC_precision': 0.953862660944206, 'eval_LOC_recall': 0.9678824169842134, 'eval_LOC_f1': 0.9608213996217239, 'eval_ORG_precision': 0.8990559186637618, 'eval_ORG_recall': 0.9211309523809523, 'eval_ORG_f1': 0.9099595736861447, 'eval_MISC_precision': 0.8561001042752867, 'eval_MISC_recall': 0.8866090712742981, 'eval_MISC_f1': 0.8710875331564988, 'eval_runtime': 1.7332, 'eval_samples_per_second': 1875.141, 'eval_steps_per_second': 15.001, 'epoch': 2.050113895216401, 'num_input_tokens_seen': 1404999}\n",
      "{'loss': 0.0429, 'grad_norm': 5.832582473754883, 'learning_rate': 3.4756069970473073e-06, 'epoch': 2.164009111617312, 'num_input_tokens_seen': 1482071}\n",
      "{'loss': 0.0452, 'grad_norm': 8.148289680480957, 'learning_rate': 2.516160013122119e-06, 'epoch': 2.277904328018223, 'num_input_tokens_seen': 1561335}\n",
      "{'eval_loss': 0.057573363184928894, 'eval_token_precision_macro': 0.9427177476113583, 'eval_token_recall_macro': 0.9451433983292427, 'eval_token_f1_macro': 0.9437362650363578, 'eval_overall_precision': 0.9270181219110379, 'eval_overall_recall': 0.9458732560094133, 'eval_overall_f1': 0.9363507779349364, 'eval_PER_precision': 0.978953049109552, 'eval_PER_recall': 0.9847991313789359, 'eval_PER_f1': 0.9818673883626522, 'eval_LOC_precision': 0.9561497326203209, 'eval_LOC_recall': 0.9733260751224823, 'eval_LOC_f1': 0.9646614513083356, 'eval_ORG_precision': 0.8640636297903109, 'eval_ORG_recall': 0.8891369047619048, 'eval_ORG_f1': 0.8764209754308764, 'eval_MISC_precision': 0.8609958506224067, 'eval_MISC_recall': 0.896328293736501, 'eval_MISC_f1': 0.8783068783068783, 'eval_runtime': 1.7738, 'eval_samples_per_second': 1832.205, 'eval_steps_per_second': 14.658, 'epoch': 2.277904328018223, 'num_input_tokens_seen': 1561335}\n",
      "{'loss': 0.045, 'grad_norm': 0.6432189345359802, 'learning_rate': 1.704323848051383e-06, 'epoch': 2.3917995444191344, 'num_input_tokens_seen': 1637959}\n",
      "{'loss': 0.0417, 'grad_norm': 4.28475284576416, 'learning_rate': 1.0558999693382504e-06, 'epoch': 2.5056947608200457, 'num_input_tokens_seen': 1713687}\n",
      "{'eval_loss': 0.053037289530038834, 'eval_token_precision_macro': 0.9477603124517248, 'eval_token_recall_macro': 0.9509462144868546, 'eval_token_f1_macro': 0.9493327338544889, 'eval_overall_precision': 0.9377070907886017, 'eval_overall_recall': 0.9514204067910573, 'eval_overall_f1': 0.9445139758030873, 'eval_PER_precision': 0.9740820734341252, 'eval_PER_recall': 0.9793702497285559, 'eval_PER_f1': 0.9767190037899296, 'eval_LOC_precision': 0.9585575888051668, 'eval_LOC_recall': 0.9695155144256941, 'eval_LOC_f1': 0.9640054127198918, 'eval_ORG_precision': 0.9031318281136198, 'eval_ORG_recall': 0.9226190476190477, 'eval_ORG_f1': 0.9127714390872286, 'eval_MISC_precision': 0.8761804826862539, 'eval_MISC_recall': 0.9017278617710583, 'eval_MISC_f1': 0.8887706226716339, 'eval_runtime': 1.689, 'eval_samples_per_second': 1924.164, 'eval_steps_per_second': 15.393, 'epoch': 2.5056947608200457, 'num_input_tokens_seen': 1713687}\n",
      "{'loss': 0.038, 'grad_norm': 3.7985987663269043, 'learning_rate': 5.83509210210643e-07, 'epoch': 2.619589977220957, 'num_input_tokens_seen': 1791719}\n",
      "{'loss': 0.0419, 'grad_norm': 0.3037065267562866, 'learning_rate': 2.9634611945845344e-07, 'epoch': 2.733485193621868, 'num_input_tokens_seen': 1868679}\n",
      "{'eval_loss': 0.0531209334731102, 'eval_token_precision_macro': 0.9468380761059451, 'eval_token_recall_macro': 0.9505037048379299, 'eval_token_f1_macro': 0.9486537410301908, 'eval_overall_precision': 0.9347323198942499, 'eval_overall_recall': 0.9509161203563624, 'eval_overall_f1': 0.942754770435797, 'eval_PER_precision': 0.9719676549865229, 'eval_PER_recall': 0.9788273615635179, 'eval_PER_f1': 0.9753854476602651, 'eval_LOC_precision': 0.9586243954862977, 'eval_LOC_recall': 0.9711486118671747, 'eval_LOC_f1': 0.9648458626284478, 'eval_ORG_precision': 0.8976034858387799, 'eval_ORG_recall': 0.9196428571428571, 'eval_ORG_f1': 0.908489525909592, 'eval_MISC_precision': 0.8696558915537018, 'eval_MISC_recall': 0.9006479481641468, 'eval_MISC_f1': 0.8848806366047746, 'eval_runtime': 1.7119, 'eval_samples_per_second': 1898.452, 'eval_steps_per_second': 15.188, 'epoch': 2.733485193621868, 'num_input_tokens_seen': 1868679}\n",
      "{'loss': 0.0465, 'grad_norm': 0.15346266329288483, 'learning_rate': 2.0000000000000002e-07, 'epoch': 2.847380410022779, 'num_input_tokens_seen': 1950135}\n",
      "{'train_runtime': 160.8835, 'train_samples_per_second': 248.627, 'train_steps_per_second': 15.539, 'train_loss': 0.20046089973449707, 'epoch': 2.847380410022779, 'num_input_tokens_seen': 1950135}\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=2500, training_loss=0.20046089973449707, metrics={'train_runtime': 160.8835, 'train_samples_per_second': 248.627, 'train_steps_per_second': 15.539, 'train_loss': 0.20046089973449707, 'epoch': 2.847380410022779, 'num_input_tokens_seen': 1950135})"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 16
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-07T16:47:50.599662Z",
     "start_time": "2024-12-07T16:47:22.493669Z"
    }
   },
   "cell_type": "code",
   "source": [
    "test_results_l12 = trainer.evaluate(Dataset(load_from_disk('data/test'), tokenizer_name=tok_name), metric_key_prefix='transfer')\n",
    "test_results_l12"
   ],
   "id": "d24056baa716f185",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'transfer_loss': 0.42435064911842346, 'transfer_token_precision_macro': 0.7183066027403103, 'transfer_token_recall_macro': 0.6434193935372868, 'transfer_token_f1_macro': 0.6434778181808135, 'transfer_overall_precision': 0.5156887354879197, 'transfer_overall_recall': 0.4577675649328041, 'transfer_overall_f1': 0.4850049798959755, 'transfer_PER_precision': 0.4577421344848859, 'transfer_PER_recall': 0.3596703829374697, 'transfer_PER_f1': 0.4028230184581976, 'transfer_LOC_precision': 0.7540342298288508, 'transfer_LOC_recall': 0.8449315068493151, 'transfer_LOC_f1': 0.7968992248062016, 'transfer_ORG_precision': 0.6072380106571936, 'transfer_ORG_recall': 0.7629009762900977, 'transfer_ORG_f1': 0.6762269749041909, 'transfer_MISC_precision': 0.3396679772826562, 'transfer_MISC_recall': 0.22575493612078978, 'transfer_MISC_f1': 0.27123669980812837, 'transfer_runtime': 26.7937, 'transfer_samples_per_second': 334.481, 'transfer_steps_per_second': 2.65, 'epoch': 2.847380410022779, 'num_input_tokens_seen': 1950135}\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'transfer_loss': 0.42435064911842346,\n",
       " 'transfer_token_precision_macro': 0.7183066027403103,\n",
       " 'transfer_token_recall_macro': 0.6434193935372868,\n",
       " 'transfer_token_f1_macro': 0.6434778181808135,\n",
       " 'transfer_overall_precision': 0.5156887354879197,\n",
       " 'transfer_overall_recall': 0.4577675649328041,\n",
       " 'transfer_overall_f1': 0.4850049798959755,\n",
       " 'transfer_PER_precision': 0.4577421344848859,\n",
       " 'transfer_PER_recall': 0.3596703829374697,\n",
       " 'transfer_PER_f1': 0.4028230184581976,\n",
       " 'transfer_LOC_precision': 0.7540342298288508,\n",
       " 'transfer_LOC_recall': 0.8449315068493151,\n",
       " 'transfer_LOC_f1': 0.7968992248062016,\n",
       " 'transfer_ORG_precision': 0.6072380106571936,\n",
       " 'transfer_ORG_recall': 0.7629009762900977,\n",
       " 'transfer_ORG_f1': 0.6762269749041909,\n",
       " 'transfer_MISC_precision': 0.3396679772826562,\n",
       " 'transfer_MISC_recall': 0.22575493612078978,\n",
       " 'transfer_MISC_f1': 0.27123669980812837,\n",
       " 'transfer_runtime': 26.7937,\n",
       " 'transfer_samples_per_second': 334.481,\n",
       " 'transfer_steps_per_second': 2.65,\n",
       " 'epoch': 2.847380410022779,\n",
       " 'num_input_tokens_seen': 1950135}"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 17
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-07T16:47:56.129989Z",
     "start_time": "2024-12-07T16:47:50.601019Z"
    }
   },
   "cell_type": "code",
   "source": "wandb.finish()",
   "id": "495a59f2e667a745",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "VBox(children=(Label(value='0.021 MB of 0.021 MB uploaded\\r'), FloatProgress(value=1.0, max=1.0)))"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "e365629ab68549cb9a14fd742085e7cf"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ],
      "text/html": [
       "<style>\n",
       "    table.wandb td:nth-child(1) { padding: 0 10px; text-align: left ; width: auto;} td:nth-child(2) {text-align: left ; width: 100%}\n",
       "    .wandb-row { display: flex; flex-direction: row; flex-wrap: wrap; justify-content: flex-start; width: 100% }\n",
       "    .wandb-col { display: flex; flex-direction: column; flex-basis: 100%; flex: 1; padding: 10px; }\n",
       "    </style>\n",
       "<div class=\"wandb-row\"><div class=\"wandb-col\"><h3>Run history:</h3><br/><table class=\"wandb\"><tr><td>eval/LOC_f1</td><td>▁▅▇▇████████</td></tr><tr><td>eval/LOC_precision</td><td>▁▅▇▇████████</td></tr><tr><td>eval/LOC_recall</td><td>▁▅█▇████████</td></tr><tr><td>eval/MISC_f1</td><td>▁▁▅▇▇███████</td></tr><tr><td>eval/MISC_precision</td><td>▁▁▅▇▇███████</td></tr><tr><td>eval/MISC_recall</td><td>▁▁▅▇▇███████</td></tr><tr><td>eval/ORG_f1</td><td>▁▂▇█████████</td></tr><tr><td>eval/ORG_precision</td><td>▁▁▇█████████</td></tr><tr><td>eval/ORG_recall</td><td>▁▂▇█████████</td></tr><tr><td>eval/PER_f1</td><td>▁███████████</td></tr><tr><td>eval/PER_precision</td><td>▁▇██████████</td></tr><tr><td>eval/PER_recall</td><td>▁███████████</td></tr><tr><td>eval/loss</td><td>█▄▂▂▁▁▁▁▁▁▁▁</td></tr><tr><td>eval/overall_f1</td><td>▁▄▇█████████</td></tr><tr><td>eval/overall_precision</td><td>▁▄▇█████████</td></tr><tr><td>eval/overall_recall</td><td>▁▅▇█████████</td></tr><tr><td>eval/runtime</td><td>▁█▃▃▂▂▂▂▂▂▂▂</td></tr><tr><td>eval/samples_per_second</td><td>█▁▂▂▃▃▃▃▃▃▃▃</td></tr><tr><td>eval/steps_per_second</td><td>█▁▂▂▃▃▃▃▃▃▃▃</td></tr><tr><td>eval/token_f1_macro</td><td>▁▄▇█████████</td></tr><tr><td>eval/token_precision_macro</td><td>▁▄▇█████████</td></tr><tr><td>eval/token_recall_macro</td><td>▁▅▇█████████</td></tr><tr><td>train/epoch</td><td>▁▁▁▂▂▂▂▂▂▃▃▃▃▄▄▄▄▄▄▅▅▅▅▅▆▆▆▆▇▇▇▇▇▇█████</td></tr><tr><td>train/global_step</td><td>▁▁▁▂▂▂▂▂▂▃▃▃▃▄▄▄▄▄▅▅▅▅▅▅▆▆▆▆▇▇▇▇▇▇█████</td></tr><tr><td>train/grad_norm</td><td>▇▃▁▃▂▃▆▄▃▅█▆▁▄▂▄▄▂▅▇▁▄▄▁▁</td></tr><tr><td>train/learning_rate</td><td>▄▇████▇▇▇▆▆▅▅▄▄▃▃▃▂▂▂▁▁▁▁</td></tr><tr><td>train/loss</td><td>█▄▃▂▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁</td></tr><tr><td>train/num_input_tokens_seen</td><td>▁▁▁▂▂▂▂▂▂▃▃▃▃▄▄▄▄▄▅▅▅▅▅▅▆▆▆▆▇▇▇▇▇▇█████</td></tr><tr><td>train/transfer_LOC_f1</td><td>▁</td></tr><tr><td>train/transfer_LOC_precision</td><td>▁</td></tr><tr><td>train/transfer_LOC_recall</td><td>▁</td></tr><tr><td>train/transfer_MISC_f1</td><td>▁</td></tr><tr><td>train/transfer_MISC_precision</td><td>▁</td></tr><tr><td>train/transfer_MISC_recall</td><td>▁</td></tr><tr><td>train/transfer_ORG_f1</td><td>▁</td></tr><tr><td>train/transfer_ORG_precision</td><td>▁</td></tr><tr><td>train/transfer_ORG_recall</td><td>▁</td></tr><tr><td>train/transfer_PER_f1</td><td>▁</td></tr><tr><td>train/transfer_PER_precision</td><td>▁</td></tr><tr><td>train/transfer_PER_recall</td><td>▁</td></tr><tr><td>train/transfer_loss</td><td>▁</td></tr><tr><td>train/transfer_overall_f1</td><td>▁</td></tr><tr><td>train/transfer_overall_precision</td><td>▁</td></tr><tr><td>train/transfer_overall_recall</td><td>▁</td></tr><tr><td>train/transfer_runtime</td><td>▁</td></tr><tr><td>train/transfer_samples_per_second</td><td>▁</td></tr><tr><td>train/transfer_steps_per_second</td><td>▁</td></tr><tr><td>train/transfer_token_f1_macro</td><td>▁</td></tr><tr><td>train/transfer_token_precision_macro</td><td>▁</td></tr><tr><td>train/transfer_token_recall_macro</td><td>▁</td></tr></table><br/></div><div class=\"wandb-col\"><h3>Run summary:</h3><br/><table class=\"wandb\"><tr><td>eval/LOC_f1</td><td>0.96485</td></tr><tr><td>eval/LOC_precision</td><td>0.95862</td></tr><tr><td>eval/LOC_recall</td><td>0.97115</td></tr><tr><td>eval/MISC_f1</td><td>0.88488</td></tr><tr><td>eval/MISC_precision</td><td>0.86966</td></tr><tr><td>eval/MISC_recall</td><td>0.90065</td></tr><tr><td>eval/ORG_f1</td><td>0.90849</td></tr><tr><td>eval/ORG_precision</td><td>0.8976</td></tr><tr><td>eval/ORG_recall</td><td>0.91964</td></tr><tr><td>eval/PER_f1</td><td>0.97539</td></tr><tr><td>eval/PER_precision</td><td>0.97197</td></tr><tr><td>eval/PER_recall</td><td>0.97883</td></tr><tr><td>eval/loss</td><td>0.05312</td></tr><tr><td>eval/overall_f1</td><td>0.94275</td></tr><tr><td>eval/overall_precision</td><td>0.93473</td></tr><tr><td>eval/overall_recall</td><td>0.95092</td></tr><tr><td>eval/runtime</td><td>1.7119</td></tr><tr><td>eval/samples_per_second</td><td>1898.452</td></tr><tr><td>eval/steps_per_second</td><td>15.188</td></tr><tr><td>eval/token_f1_macro</td><td>0.94865</td></tr><tr><td>eval/token_precision_macro</td><td>0.94684</td></tr><tr><td>eval/token_recall_macro</td><td>0.9505</td></tr><tr><td>total_flos</td><td>995305076666010.0</td></tr><tr><td>train/epoch</td><td>2.84738</td></tr><tr><td>train/global_step</td><td>2500</td></tr><tr><td>train/grad_norm</td><td>0.15346</td></tr><tr><td>train/learning_rate</td><td>0.0</td></tr><tr><td>train/loss</td><td>0.0465</td></tr><tr><td>train/num_input_tokens_seen</td><td>1950135</td></tr><tr><td>train/transfer_LOC_f1</td><td>0.7969</td></tr><tr><td>train/transfer_LOC_precision</td><td>0.75403</td></tr><tr><td>train/transfer_LOC_recall</td><td>0.84493</td></tr><tr><td>train/transfer_MISC_f1</td><td>0.27124</td></tr><tr><td>train/transfer_MISC_precision</td><td>0.33967</td></tr><tr><td>train/transfer_MISC_recall</td><td>0.22575</td></tr><tr><td>train/transfer_ORG_f1</td><td>0.67623</td></tr><tr><td>train/transfer_ORG_precision</td><td>0.60724</td></tr><tr><td>train/transfer_ORG_recall</td><td>0.7629</td></tr><tr><td>train/transfer_PER_f1</td><td>0.40282</td></tr><tr><td>train/transfer_PER_precision</td><td>0.45774</td></tr><tr><td>train/transfer_PER_recall</td><td>0.35967</td></tr><tr><td>train/transfer_loss</td><td>0.42435</td></tr><tr><td>train/transfer_overall_f1</td><td>0.485</td></tr><tr><td>train/transfer_overall_precision</td><td>0.51569</td></tr><tr><td>train/transfer_overall_recall</td><td>0.45777</td></tr><tr><td>train/transfer_runtime</td><td>26.7937</td></tr><tr><td>train/transfer_samples_per_second</td><td>334.481</td></tr><tr><td>train/transfer_steps_per_second</td><td>2.65</td></tr><tr><td>train/transfer_token_f1_macro</td><td>0.64348</td></tr><tr><td>train/transfer_token_precision_macro</td><td>0.71831</td></tr><tr><td>train/transfer_token_recall_macro</td><td>0.64342</td></tr><tr><td>train_loss</td><td>0.20046</td></tr><tr><td>train_runtime</td><td>160.8835</td></tr><tr><td>train_samples_per_second</td><td>248.627</td></tr><tr><td>train_steps_per_second</td><td>15.539</td></tr></table><br/></div></div>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ],
      "text/html": [
       " View run <strong style=\"color:#cdcd00\">facebook__xlm-v-base-finetuned-l12-conll03/12-07/0</strong> at: <a href='https://wandb.ai/viktoroo-sch-epfl/ner-alignment/runs/yaiig3k9' target=\"_blank\">https://wandb.ai/viktoroo-sch-epfl/ner-alignment/runs/yaiig3k9</a><br/> View project at: <a href='https://wandb.ai/viktoroo-sch-epfl/ner-alignment' target=\"_blank\">https://wandb.ai/viktoroo-sch-epfl/ner-alignment</a><br/>Synced 6 W&B file(s), 0 media file(s), 0 artifact file(s) and 0 other file(s)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ],
      "text/html": [
       "Find logs at: <code>facebook__xlm-v-base-finetuned-l12-conll03/12-07/0/wandb/run-20241207_164430-yaiig3k9/logs</code>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ],
      "text/html": [
       "wandb version 0.19.0 is available!  To upgrade, please run:\n",
       " $ pip install wandb --upgrade"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ],
      "text/html": [
       "The new W&B backend becomes opt-out in version 0.18.0; try it out with `wandb.require(\"core\")`! See https://wandb.me/wandb-core for more information."
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "execution_count": 18
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "## 2.3 Train a truncated model (without last layer)",
   "id": "b2a65b20e3ef2f5d"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-07T16:50:38.786934Z",
     "start_time": "2024-12-07T16:47:56.131737Z"
    }
   },
   "cell_type": "code",
   "source": [
    "run_name = f'{tok_name}-finetuned-l11-conll03/{datetime.now().strftime(\"%m-%d\")}/{n_run}'\n",
    "wandb.init(\n",
    "    project='ner-alignment',\n",
    "    name=run_name,\n",
    "    dir=run_name,\n",
    "    resume=False\n",
    ")\n",
    "n_run += 1\n",
    "\n",
    "model = XLMRobertaForTokenClassification.from_pretrained(\n",
    "    model_name, \n",
    "    num_labels=9, \n",
    "    skip_last_layer=True\n",
    ")\n",
    "model.roberta.embeddings.requires_grad_(False)  # freeze input embeddings to avoid parameter shift (training on english and inferencing on africaans -> different tokens are activated)\n",
    "print(f\"Percentage of frozen modules: {100 * sum(1 for module in model.modules() if not any(p.requires_grad for p in module.parameters())) / sum(1 for module in model.modules()):.2f}%\")\n",
    "print(f\"Percentage of frozen parameters: {100 * sum(p.numel() for p in model.parameters() if not p.requires_grad) / sum(p.numel() for p in model.parameters()):.2f}%\")\n",
    "\n",
    "\n",
    "tok = AutoTokenizer.from_pretrained(model_name)\n",
    "\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=TrainingArguments(\n",
    "        output_dir=run_name,\n",
    "        overwrite_output_dir=True,\n",
    "        eval_strategy='steps',\n",
    "        per_device_train_batch_size=batch_size,\n",
    "        per_device_eval_batch_size=128,\n",
    "        learning_rate=lr,\n",
    "        max_steps=steps,\n",
    "        lr_scheduler_type='cosine_with_min_lr',\n",
    "        lr_scheduler_kwargs={ 'num_cycles': 0.5, 'min_lr_rate': 0.01 },\n",
    "        warmup_ratio=0.1,\n",
    "        adam_epsilon=1e-8,\n",
    "        adam_beta1=0.9,\n",
    "        adam_beta2=0.999,\n",
    "        weight_decay=0.0,\n",
    "        logging_steps=100,\n",
    "        eval_steps=200,\n",
    "        torch_compile=False,\n",
    "        bf16=True,\n",
    "        include_num_input_tokens_seen=True,\n",
    "        disable_tqdm=True,\n",
    "        report_to='wandb'\n",
    "    ),\n",
    "    data_collator=partial(collate_fn, pad_token=tok.pad_token_id),\n",
    "    train_dataset=Dataset(load_from_disk('data/train'), tokenizer_name=tok_name),\n",
    "    eval_dataset=Dataset(load_from_disk('data/valid'), tokenizer_name=tok_name),\n",
    "    compute_metrics=compute_ner_metrics\n",
    ")\n",
    "trainer.train()"
   ],
   "id": "90cb1610c1f20c9a",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "VBox(children=(Label(value='Waiting for wandb.init()...\\r'), FloatProgress(value=0.011112837596899933, max=1.0…"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "bcddc162980c44119e581e58aa0f9aa8"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ],
      "text/html": [
       "Tracking run with wandb version 0.17.9"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ],
      "text/html": [
       "Run data is saved locally in <code>facebook__xlm-v-base-finetuned-l11-conll03/12-07/1/wandb/run-20241207_164756-j955yfyx</code>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ],
      "text/html": [
       "Syncing run <strong><a href='https://wandb.ai/viktoroo-sch-epfl/ner-alignment/runs/j955yfyx' target=\"_blank\">facebook__xlm-v-base-finetuned-l11-conll03/12-07/1</a></strong> to <a href='https://wandb.ai/viktoroo-sch-epfl/ner-alignment' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ],
      "text/html": [
       " View project at <a href='https://wandb.ai/viktoroo-sch-epfl/ner-alignment' target=\"_blank\">https://wandb.ai/viktoroo-sch-epfl/ner-alignment</a>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ],
      "text/html": [
       " View run at <a href='https://wandb.ai/viktoroo-sch-epfl/ner-alignment/runs/j955yfyx' target=\"_blank\">https://wandb.ai/viktoroo-sch-epfl/ner-alignment/runs/j955yfyx</a>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of XLMRobertaForTokenClassification were not initialized from the model checkpoint at facebook/xlm-v-base and are newly initialized: ['classifier.bias', 'classifier.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Percentage of frozen modules: 24.29%\n",
      "Percentage of frozen parameters: 89.88%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/mloscratch/homes/shcherba/conda/envs/char-llm/lib/python3.11/site-packages/transformers/tokenization_utils_base.py:1601: FutureWarning: `clean_up_tokenization_spaces` was not set. It will be set to `True` by default. This behavior will be depracted in transformers v4.45, and will be then set to `False` by default. For more details check this issue: https://github.com/huggingface/transformers/issues/31884\n",
      "  warnings.warn(\n",
      "max_steps is given, it will override any value given in num_train_epochs\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 1.3923, 'grad_norm': 5.5005974769592285, 'learning_rate': 8.000000000000001e-06, 'epoch': 0.11389521640091116, 'num_input_tokens_seen': 76672}\n",
      "{'loss': 0.2907, 'grad_norm': 2.5719411373138428, 'learning_rate': 1.6000000000000003e-05, 'epoch': 0.22779043280182232, 'num_input_tokens_seen': 158272}\n",
      "{'eval_loss': 0.1464698612689972, 'eval_token_precision_macro': 0.8515416048410184, 'eval_token_recall_macro': 0.7551059357327375, 'eval_token_f1_macro': 0.7532563089953259, 'eval_overall_precision': 0.7416327807319362, 'eval_overall_recall': 0.7971087577744159, 'eval_overall_f1': 0.7683707364498097, 'eval_PER_precision': 0.9471153846153846, 'eval_PER_recall': 0.9625407166123778, 'eval_PER_f1': 0.9547657512116317, 'eval_LOC_precision': 0.8446215139442231, 'eval_LOC_recall': 0.9232444202504083, 'eval_LOC_f1': 0.882184655396619, 'eval_ORG_precision': 0.605296343001261, 'eval_ORG_recall': 0.7142857142857143, 'eval_ORG_f1': 0.6552901023890785, 'eval_MISC_precision': 0.33728448275862066, 'eval_MISC_recall': 0.33801295896328293, 'eval_MISC_f1': 0.3376483279395901, 'eval_runtime': 3.2682, 'eval_samples_per_second': 994.443, 'eval_steps_per_second': 7.956, 'epoch': 0.22779043280182232, 'num_input_tokens_seen': 158272}\n",
      "{'loss': 0.1617, 'grad_norm': 3.0946362018585205, 'learning_rate': 1.997588409757226e-05, 'epoch': 0.3416856492027335, 'num_input_tokens_seen': 236576}\n",
      "{'loss': 0.1043, 'grad_norm': 3.585635185241699, 'learning_rate': 1.9783661247264677e-05, 'epoch': 0.45558086560364464, 'num_input_tokens_seen': 313200}\n",
      "{'eval_loss': 0.07360916584730148, 'eval_token_precision_macro': 0.9065100215186147, 'eval_token_recall_macro': 0.9094292606350022, 'eval_token_f1_macro': 0.9072403662143186, 'eval_overall_precision': 0.8720218931101095, 'eval_overall_recall': 0.9105732055807699, 'eval_overall_f1': 0.8908806841542637, 'eval_PER_precision': 0.9576634512325831, 'eval_PER_recall': 0.9701411509229099, 'eval_PER_f1': 0.9638619201725998, 'eval_LOC_precision': 0.9004214963119073, 'eval_LOC_recall': 0.9303211758301578, 'eval_LOC_f1': 0.9151271753681393, 'eval_ORG_precision': 0.8240997229916898, 'eval_ORG_recall': 0.8854166666666666, 'eval_ORG_f1': 0.8536585365853658, 'eval_MISC_precision': 0.7280876494023905, 'eval_MISC_recall': 0.7894168466522679, 'eval_MISC_f1': 0.7575129533678756, 'eval_runtime': 2.2544, 'eval_samples_per_second': 1441.653, 'eval_steps_per_second': 11.533, 'epoch': 0.45558086560364464, 'num_input_tokens_seen': 313200}\n",
      "{'loss': 0.1006, 'grad_norm': 2.031691551208496, 'learning_rate': 1.9402956945780496e-05, 'epoch': 0.5694760820045558, 'num_input_tokens_seen': 390432}\n",
      "{'loss': 0.0778, 'grad_norm': 3.8352091312408447, 'learning_rate': 1.884118116930338e-05, 'epoch': 0.683371298405467, 'num_input_tokens_seen': 468592}\n",
      "{'eval_loss': 0.07182179391384125, 'eval_token_precision_macro': 0.9132353519945036, 'eval_token_recall_macro': 0.9149354907963946, 'eval_token_f1_macro': 0.9130048438968559, 'eval_overall_precision': 0.8836569579288026, 'eval_overall_recall': 0.9179694066229619, 'eval_overall_f1': 0.9004864374639294, 'eval_PER_precision': 0.9617662897145934, 'eval_PER_recall': 0.9695982627578719, 'eval_PER_f1': 0.9656663963233306, 'eval_LOC_precision': 0.9040084388185654, 'eval_LOC_recall': 0.9330430048992924, 'eval_LOC_f1': 0.9182962764532548, 'eval_ORG_precision': 0.8372905027932961, 'eval_ORG_recall': 0.8921130952380952, 'eval_ORG_f1': 0.8638328530259366, 'eval_MISC_precision': 0.7658291457286432, 'eval_MISC_recall': 0.8228941684665226, 'eval_MISC_f1': 0.793336803748048, 'eval_runtime': 2.2299, 'eval_samples_per_second': 1457.472, 'eval_steps_per_second': 11.66, 'epoch': 0.683371298405467, 'num_input_tokens_seen': 468592}\n",
      "{'loss': 0.0766, 'grad_norm': 2.468250274658203, 'learning_rate': 1.8109268244311982e-05, 'epoch': 0.7972665148063781, 'num_input_tokens_seen': 549056}\n",
      "{'loss': 0.0683, 'grad_norm': 4.63446569442749, 'learning_rate': 1.7221464023352647e-05, 'epoch': 0.9111617312072893, 'num_input_tokens_seen': 626848}\n",
      "{'eval_loss': 0.0577804334461689, 'eval_token_precision_macro': 0.9250198332153635, 'eval_token_recall_macro': 0.9345968960882356, 'eval_token_f1_macro': 0.9289966332805961, 'eval_overall_precision': 0.9024709540173458, 'eval_overall_recall': 0.9270465624474702, 'eval_overall_f1': 0.9145936981757877, 'eval_PER_precision': 0.952329360780065, 'eval_PER_recall': 0.9543973941368078, 'eval_PER_f1': 0.9533622559652928, 'eval_LOC_precision': 0.9186170212765957, 'eval_LOC_recall': 0.9401197604790419, 'eval_LOC_f1': 0.9292440139897767, 'eval_ORG_precision': 0.8763326226012793, 'eval_ORG_recall': 0.9174107142857143, 'eval_ORG_f1': 0.8964013086150491, 'eval_MISC_precision': 0.8149284253578732, 'eval_MISC_recall': 0.8606911447084233, 'eval_MISC_f1': 0.8371848739495799, 'eval_runtime': 2.1239, 'eval_samples_per_second': 1530.197, 'eval_steps_per_second': 12.242, 'epoch': 0.9111617312072893, 'num_input_tokens_seen': 626848}\n",
      "{'loss': 0.0653, 'grad_norm': 2.882131576538086, 'learning_rate': 1.619504860572402e-05, 'epoch': 1.0250569476082005, 'num_input_tokens_seen': 705647}\n",
      "{'loss': 0.0473, 'grad_norm': 1.6385911703109741, 'learning_rate': 1.505e-05, 'epoch': 1.1389521640091116, 'num_input_tokens_seen': 784495}\n",
      "{'eval_loss': 0.05213475599884987, 'eval_token_precision_macro': 0.9461058628970466, 'eval_token_recall_macro': 0.9436391206725951, 'eval_token_f1_macro': 0.9445758647398412, 'eval_overall_precision': 0.9199215558097729, 'eval_overall_recall': 0.9462094469658766, 'eval_overall_f1': 0.9328803447132914, 'eval_PER_precision': 0.9688674181427804, 'eval_PER_recall': 0.9799131378935939, 'eval_PER_f1': 0.9743589743589742, 'eval_LOC_precision': 0.9470058293587705, 'eval_LOC_recall': 0.9727817093086554, 'eval_LOC_f1': 0.9597207303974221, 'eval_ORG_precision': 0.873229461756374, 'eval_ORG_recall': 0.9174107142857143, 'eval_ORG_f1': 0.8947750362844703, 'eval_MISC_precision': 0.8401253918495298, 'eval_MISC_recall': 0.8682505399568035, 'eval_MISC_f1': 0.8539564524694637, 'eval_runtime': 1.9171, 'eval_samples_per_second': 1695.234, 'eval_steps_per_second': 13.562, 'epoch': 1.1389521640091116, 'num_input_tokens_seen': 784495}\n",
      "{'loss': 0.0506, 'grad_norm': 0.6027234196662903, 'learning_rate': 1.3808605274817533e-05, 'epoch': 1.2528473804100229, 'num_input_tokens_seen': 861071}\n",
      "{'loss': 0.0477, 'grad_norm': 4.328546047210693, 'learning_rate': 1.2495026766436711e-05, 'epoch': 1.366742596810934, 'num_input_tokens_seen': 936543}\n",
      "{'eval_loss': 0.05099925771355629, 'eval_token_precision_macro': 0.948661449812787, 'eval_token_recall_macro': 0.9427786991947933, 'eval_token_f1_macro': 0.9450514309262985, 'eval_overall_precision': 0.9220288903479974, 'eval_overall_recall': 0.944192301227097, 'eval_overall_f1': 0.9329789884561084, 'eval_PER_precision': 0.9644779332615716, 'eval_PER_recall': 0.9728555917480999, 'eval_PER_f1': 0.9686486486486486, 'eval_LOC_precision': 0.9384288747346072, 'eval_LOC_recall': 0.9624387588459444, 'eval_LOC_f1': 0.9502821822090836, 'eval_ORG_precision': 0.8923734853884533, 'eval_ORG_recall': 0.9315476190476191, 'eval_ORG_f1': 0.9115398616672733, 'eval_MISC_precision': 0.8500527983104541, 'eval_MISC_recall': 0.8693304535637149, 'eval_MISC_f1': 0.8595835557928456, 'eval_runtime': 1.9312, 'eval_samples_per_second': 1682.909, 'eval_steps_per_second': 13.463, 'epoch': 1.366742596810934, 'num_input_tokens_seen': 936543}\n",
      "{'loss': 0.0454, 'grad_norm': 4.917657375335693, 'learning_rate': 1.113483178634977e-05, 'epoch': 1.4806378132118452, 'num_input_tokens_seen': 1015087}\n",
      "{'loss': 0.0499, 'grad_norm': 4.22766637802124, 'learning_rate': 9.754494982645244e-06, 'epoch': 1.5945330296127562, 'num_input_tokens_seen': 1092351}\n",
      "{'eval_loss': 0.04322414845228195, 'eval_token_precision_macro': 0.9453936337820141, 'eval_token_recall_macro': 0.949864385368238, 'eval_token_f1_macro': 0.9475776995032249, 'eval_overall_precision': 0.9309264754368612, 'eval_overall_recall': 0.9492351655740461, 'eval_overall_f1': 0.9399916770703288, 'eval_PER_precision': 0.9725067385444744, 'eval_PER_recall': 0.9793702497285559, 'eval_PER_f1': 0.9759264268325669, 'eval_LOC_precision': 0.9497594869053981, 'eval_LOC_recall': 0.9673380511703865, 'eval_LOC_f1': 0.9584681769147788, 'eval_ORG_precision': 0.9042397660818714, 'eval_ORG_recall': 0.9203869047619048, 'eval_ORG_f1': 0.9122418879056047, 'eval_MISC_precision': 0.852880658436214, 'eval_MISC_recall': 0.8952483801295896, 'eval_MISC_f1': 0.8735511064278189, 'eval_runtime': 1.8279, 'eval_samples_per_second': 1777.953, 'eval_steps_per_second': 14.224, 'epoch': 1.5945330296127562, 'num_input_tokens_seen': 1092351}\n",
      "{'loss': 0.0399, 'grad_norm': 1.7565810680389404, 'learning_rate': 8.380883041097391e-06, 'epoch': 1.7084282460136673, 'num_input_tokens_seen': 1172831}\n",
      "{'loss': 0.0411, 'grad_norm': 2.750422477722168, 'learning_rate': 7.0407317556880214e-06, 'epoch': 1.8223234624145785, 'num_input_tokens_seen': 1251359}\n",
      "{'eval_loss': 0.04341994225978851, 'eval_token_precision_macro': 0.9469388943755743, 'eval_token_recall_macro': 0.9514851617732144, 'eval_token_f1_macro': 0.9491575764555945, 'eval_overall_precision': 0.9290756854375308, 'eval_overall_recall': 0.9512523113128257, 'eval_overall_f1': 0.9400332225913621, 'eval_PER_precision': 0.9692556634304207, 'eval_PER_recall': 0.9755700325732899, 'eval_PER_f1': 0.9724025974025974, 'eval_LOC_precision': 0.9559214020180563, 'eval_LOC_recall': 0.979858464888405, 'eval_LOC_f1': 0.9677419354838709, 'eval_ORG_precision': 0.8880057803468208, 'eval_ORG_recall': 0.9144345238095238, 'eval_ORG_f1': 0.9010263929618768, 'eval_MISC_precision': 0.8587628865979381, 'eval_MISC_recall': 0.8995680345572354, 'eval_MISC_f1': 0.878691983122363, 'eval_runtime': 1.8566, 'eval_samples_per_second': 1750.491, 'eval_steps_per_second': 14.004, 'epoch': 1.8223234624145785, 'num_input_tokens_seen': 1251359}\n",
      "{'loss': 0.0433, 'grad_norm': 0.5333500504493713, 'learning_rate': 5.760125646788133e-06, 'epoch': 1.9362186788154898, 'num_input_tokens_seen': 1325695}\n",
      "{'loss': 0.0291, 'grad_norm': 3.8538239002227783, 'learning_rate': 4.563990255639608e-06, 'epoch': 2.050113895216401, 'num_input_tokens_seen': 1404999}\n",
      "{'eval_loss': 0.04461473599076271, 'eval_token_precision_macro': 0.9502250601013588, 'eval_token_recall_macro': 0.953392367595862, 'eval_token_f1_macro': 0.9517811270818467, 'eval_overall_precision': 0.9345179335307667, 'eval_overall_recall': 0.95478231635569, 'eval_overall_f1': 0.9445414484077492, 'eval_PER_precision': 0.9731471535982814, 'eval_PER_recall': 0.9837133550488599, 'eval_PER_f1': 0.978401727861771, 'eval_LOC_precision': 0.9546666666666667, 'eval_LOC_recall': 0.9744148067501361, 'eval_LOC_f1': 0.9644396551724139, 'eval_ORG_precision': 0.9032023289665211, 'eval_ORG_recall': 0.9233630952380952, 'eval_ORG_f1': 0.9131714495952906, 'eval_MISC_precision': 0.8655635987590486, 'eval_MISC_recall': 0.9038876889848813, 'eval_MISC_f1': 0.884310618066561, 'eval_runtime': 1.7558, 'eval_samples_per_second': 1851.031, 'eval_steps_per_second': 14.808, 'epoch': 2.050113895216401, 'num_input_tokens_seen': 1404999}\n",
      "{'loss': 0.0245, 'grad_norm': 1.7131320238113403, 'learning_rate': 3.4756069970473073e-06, 'epoch': 2.164009111617312, 'num_input_tokens_seen': 1482071}\n",
      "{'loss': 0.0279, 'grad_norm': 4.496512413024902, 'learning_rate': 2.516160013122119e-06, 'epoch': 2.277904328018223, 'num_input_tokens_seen': 1561335}\n",
      "{'eval_loss': 0.0454585887491703, 'eval_token_precision_macro': 0.9543066449922732, 'eval_token_recall_macro': 0.9523017313480034, 'eval_token_f1_macro': 0.9532326882123316, 'eval_overall_precision': 0.9390384933091029, 'eval_overall_recall': 0.9554546982686166, 'eval_overall_f1': 0.9471754707548741, 'eval_PER_precision': 0.9773828756058158, 'eval_PER_recall': 0.9853420195439739, 'eval_PER_f1': 0.9813463098134632, 'eval_LOC_precision': 0.958779443254818, 'eval_LOC_recall': 0.974959172563963, 'eval_LOC_f1': 0.9668016194331984, 'eval_ORG_precision': 0.9045189504373178, 'eval_ORG_recall': 0.9233630952380952, 'eval_ORG_f1': 0.9138438880706922, 'eval_MISC_precision': 0.8755230125523012, 'eval_MISC_recall': 0.9038876889848813, 'eval_MISC_f1': 0.8894792773645058, 'eval_runtime': 1.724, 'eval_samples_per_second': 1885.099, 'eval_steps_per_second': 15.081, 'epoch': 2.277904328018223, 'num_input_tokens_seen': 1561335}\n",
      "{'loss': 0.026, 'grad_norm': 16.95728874206543, 'learning_rate': 1.704323848051383e-06, 'epoch': 2.3917995444191344, 'num_input_tokens_seen': 1637959}\n",
      "{'loss': 0.0255, 'grad_norm': 1.885147213935852, 'learning_rate': 1.0558999693382504e-06, 'epoch': 2.5056947608200457, 'num_input_tokens_seen': 1713687}\n",
      "{'eval_loss': 0.04246268793940544, 'eval_token_precision_macro': 0.9533287286074454, 'eval_token_recall_macro': 0.9543990584426258, 'eval_token_f1_macro': 0.9538286644284916, 'eval_overall_precision': 0.9374380983823044, 'eval_overall_recall': 0.9546142208774584, 'eval_overall_f1': 0.9459481968851504, 'eval_PER_precision': 0.9747039827771797, 'eval_PER_recall': 0.9831704668838219, 'eval_PER_f1': 0.9789189189189189, 'eval_LOC_precision': 0.9566844919786096, 'eval_LOC_recall': 0.9738704409363091, 'eval_LOC_f1': 0.9652009711356891, 'eval_ORG_precision': 0.9098901098901099, 'eval_ORG_recall': 0.9241071428571429, 'eval_ORG_f1': 0.9169435215946844, 'eval_MISC_precision': 0.8673575129533678, 'eval_MISC_recall': 0.9038876889848813, 'eval_MISC_f1': 0.8852459016393442, 'eval_runtime': 1.7363, 'eval_samples_per_second': 1871.794, 'eval_steps_per_second': 14.974, 'epoch': 2.5056947608200457, 'num_input_tokens_seen': 1713687}\n",
      "{'loss': 0.0225, 'grad_norm': 3.2388148307800293, 'learning_rate': 5.83509210210643e-07, 'epoch': 2.619589977220957, 'num_input_tokens_seen': 1791719}\n",
      "{'loss': 0.0264, 'grad_norm': 0.7561604976654053, 'learning_rate': 2.9634611945845344e-07, 'epoch': 2.733485193621868, 'num_input_tokens_seen': 1868679}\n",
      "{'eval_loss': 0.041983623057603836, 'eval_token_precision_macro': 0.9534355259852431, 'eval_token_recall_macro': 0.95519784633267, 'eval_token_f1_macro': 0.9542904612194815, 'eval_overall_precision': 0.9377476882430648, 'eval_overall_recall': 0.9546142208774584, 'eval_overall_f1': 0.9461057892544774, 'eval_PER_precision': 0.9752021563342318, 'eval_PER_recall': 0.9820846905537459, 'eval_PER_f1': 0.9786313226940763, 'eval_LOC_precision': 0.9577540106951872, 'eval_LOC_recall': 0.974959172563963, 'eval_LOC_f1': 0.9662800107903966, 'eval_ORG_precision': 0.9093567251461988, 'eval_ORG_recall': 0.9255952380952381, 'eval_ORG_f1': 0.9174041297935103, 'eval_MISC_precision': 0.8670820353063343, 'eval_MISC_recall': 0.9017278617710583, 'eval_MISC_f1': 0.884065643197459, 'eval_runtime': 1.7286, 'eval_samples_per_second': 1880.171, 'eval_steps_per_second': 15.041, 'epoch': 2.733485193621868, 'num_input_tokens_seen': 1868679}\n",
      "{'loss': 0.0268, 'grad_norm': 0.12387741357088089, 'learning_rate': 2.0000000000000002e-07, 'epoch': 2.847380410022779, 'num_input_tokens_seen': 1950135}\n",
      "{'train_runtime': 150.9191, 'train_samples_per_second': 265.043, 'train_steps_per_second': 16.565, 'train_loss': 0.11646024713516236, 'epoch': 2.847380410022779, 'num_input_tokens_seen': 1950135}\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=2500, training_loss=0.11646024713516236, metrics={'train_runtime': 150.9191, 'train_samples_per_second': 265.043, 'train_steps_per_second': 16.565, 'train_loss': 0.11646024713516236, 'epoch': 2.847380410022779, 'num_input_tokens_seen': 1950135})"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 19
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-07T16:51:07.544327Z",
     "start_time": "2024-12-07T16:50:38.787949Z"
    }
   },
   "cell_type": "code",
   "source": [
    "test_results_l11 = trainer.evaluate(Dataset(load_from_disk('data/test'), tokenizer_name=tok_name), metric_key_prefix='transfer')\n",
    "test_results_l11"
   ],
   "id": "d205289b25ad7bd6",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'transfer_loss': 0.41569140553474426, 'transfer_token_precision_macro': 0.7318125169933103, 'transfer_token_recall_macro': 0.6453559423396728, 'transfer_token_f1_macro': 0.6507366667293208, 'transfer_overall_precision': 0.5066474208007293, 'transfer_overall_recall': 0.46438270315437646, 'transfer_overall_f1': 0.4845952623165238, 'transfer_PER_precision': 0.44285714285714284, 'transfer_PER_recall': 0.36063984488608825, 'transfer_PER_f1': 0.3975420785466204, 'transfer_LOC_precision': 0.7189163038219643, 'transfer_LOC_recall': 0.8142465753424658, 'transfer_LOC_f1': 0.7636176772867421, 'transfer_ORG_precision': 0.6208838821490468, 'transfer_ORG_recall': 0.799442119944212, 'transfer_ORG_f1': 0.6989391537617364, 'transfer_MISC_precision': 0.3277083333333333, 'transfer_MISC_recall': 0.22836817653890826, 'transfer_MISC_f1': 0.2691649555099247, 'transfer_runtime': 27.5571, 'transfer_samples_per_second': 325.215, 'transfer_steps_per_second': 2.576, 'epoch': 2.847380410022779, 'num_input_tokens_seen': 1950135}\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'transfer_loss': 0.41569140553474426,\n",
       " 'transfer_token_precision_macro': 0.7318125169933103,\n",
       " 'transfer_token_recall_macro': 0.6453559423396728,\n",
       " 'transfer_token_f1_macro': 0.6507366667293208,\n",
       " 'transfer_overall_precision': 0.5066474208007293,\n",
       " 'transfer_overall_recall': 0.46438270315437646,\n",
       " 'transfer_overall_f1': 0.4845952623165238,\n",
       " 'transfer_PER_precision': 0.44285714285714284,\n",
       " 'transfer_PER_recall': 0.36063984488608825,\n",
       " 'transfer_PER_f1': 0.3975420785466204,\n",
       " 'transfer_LOC_precision': 0.7189163038219643,\n",
       " 'transfer_LOC_recall': 0.8142465753424658,\n",
       " 'transfer_LOC_f1': 0.7636176772867421,\n",
       " 'transfer_ORG_precision': 0.6208838821490468,\n",
       " 'transfer_ORG_recall': 0.799442119944212,\n",
       " 'transfer_ORG_f1': 0.6989391537617364,\n",
       " 'transfer_MISC_precision': 0.3277083333333333,\n",
       " 'transfer_MISC_recall': 0.22836817653890826,\n",
       " 'transfer_MISC_f1': 0.2691649555099247,\n",
       " 'transfer_runtime': 27.5571,\n",
       " 'transfer_samples_per_second': 325.215,\n",
       " 'transfer_steps_per_second': 2.576,\n",
       " 'epoch': 2.847380410022779,\n",
       " 'num_input_tokens_seen': 1950135}"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 20
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-07T16:51:12.972857Z",
     "start_time": "2024-12-07T16:51:07.545586Z"
    }
   },
   "cell_type": "code",
   "source": "wandb.finish()",
   "id": "f6aeb50a791f0eaa",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "VBox(children=(Label(value='0.021 MB of 0.021 MB uploaded\\r'), FloatProgress(value=1.0, max=1.0)))"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "6a0e4c773ddf428285d59a795fa87569"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ],
      "text/html": [
       "<style>\n",
       "    table.wandb td:nth-child(1) { padding: 0 10px; text-align: left ; width: auto;} td:nth-child(2) {text-align: left ; width: 100%}\n",
       "    .wandb-row { display: flex; flex-direction: row; flex-wrap: wrap; justify-content: flex-start; width: 100% }\n",
       "    .wandb-col { display: flex; flex-direction: column; flex-basis: 100%; flex: 1; padding: 10px; }\n",
       "    </style>\n",
       "<div class=\"wandb-row\"><div class=\"wandb-col\"><h3>Run history:</h3><br/><table class=\"wandb\"><tr><td>eval/LOC_f1</td><td>▁▄▄▅▇▇▇█████</td></tr><tr><td>eval/LOC_precision</td><td>▁▄▅▆▇▇▇█████</td></tr><tr><td>eval/LOC_recall</td><td>▁▂▂▃▇▆▆█▇▇▇▇</td></tr><tr><td>eval/MISC_f1</td><td>▁▆▇▇████████</td></tr><tr><td>eval/MISC_precision</td><td>▁▆▇▇████████</td></tr><tr><td>eval/MISC_recall</td><td>▁▇▇▇████████</td></tr><tr><td>eval/ORG_f1</td><td>▁▆▇▇▇███████</td></tr><tr><td>eval/ORG_precision</td><td>▁▆▆▇▇██▇████</td></tr><tr><td>eval/ORG_recall</td><td>▁▇▇████▇████</td></tr><tr><td>eval/PER_f1</td><td>▁▄▄▁▆▅▇▆▇█▇▇</td></tr><tr><td>eval/PER_precision</td><td>▁▃▄▂▆▅▇▆▇█▇▇</td></tr><tr><td>eval/PER_recall</td><td>▃▅▄▁▇▅▇▆███▇</td></tr><tr><td>eval/loss</td><td>█▃▃▂▂▂▁▁▁▁▁▁</td></tr><tr><td>eval/overall_f1</td><td>▁▆▆▇▇▇██████</td></tr><tr><td>eval/overall_precision</td><td>▁▆▆▇▇▇██████</td></tr><tr><td>eval/overall_recall</td><td>▁▆▆▇████████</td></tr><tr><td>eval/runtime</td><td>█▃▃▃▂▂▁▂▁▁▁▁</td></tr><tr><td>eval/samples_per_second</td><td>▁▅▅▅▇▆▇▇████</td></tr><tr><td>eval/steps_per_second</td><td>▁▅▅▅▇▆▇▇████</td></tr><tr><td>eval/token_f1_macro</td><td>▁▆▇▇████████</td></tr><tr><td>eval/token_precision_macro</td><td>▁▅▅▆▇█▇▇████</td></tr><tr><td>eval/token_recall_macro</td><td>▁▆▇▇████████</td></tr><tr><td>train/epoch</td><td>▁▁▁▂▂▂▂▂▂▃▃▃▃▄▄▄▄▄▄▅▅▅▅▅▆▆▆▆▇▇▇▇▇▇█████</td></tr><tr><td>train/global_step</td><td>▁▁▁▂▂▂▂▂▂▃▃▃▃▄▄▄▄▄▅▅▅▅▅▅▆▆▆▆▇▇▇▇▇▇█████</td></tr><tr><td>train/grad_norm</td><td>▃▂▂▂▂▃▂▃▂▂▁▃▃▃▂▂▁▃▂▃█▂▂▁▁</td></tr><tr><td>train/learning_rate</td><td>▄▇████▇▇▇▆▆▅▅▄▄▃▃▃▂▂▂▁▁▁▁</td></tr><tr><td>train/loss</td><td>█▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁</td></tr><tr><td>train/num_input_tokens_seen</td><td>▁▁▁▂▂▂▂▂▂▃▃▃▃▄▄▄▄▄▅▅▅▅▅▅▆▆▆▆▇▇▇▇▇▇█████</td></tr><tr><td>train/transfer_LOC_f1</td><td>▁</td></tr><tr><td>train/transfer_LOC_precision</td><td>▁</td></tr><tr><td>train/transfer_LOC_recall</td><td>▁</td></tr><tr><td>train/transfer_MISC_f1</td><td>▁</td></tr><tr><td>train/transfer_MISC_precision</td><td>▁</td></tr><tr><td>train/transfer_MISC_recall</td><td>▁</td></tr><tr><td>train/transfer_ORG_f1</td><td>▁</td></tr><tr><td>train/transfer_ORG_precision</td><td>▁</td></tr><tr><td>train/transfer_ORG_recall</td><td>▁</td></tr><tr><td>train/transfer_PER_f1</td><td>▁</td></tr><tr><td>train/transfer_PER_precision</td><td>▁</td></tr><tr><td>train/transfer_PER_recall</td><td>▁</td></tr><tr><td>train/transfer_loss</td><td>▁</td></tr><tr><td>train/transfer_overall_f1</td><td>▁</td></tr><tr><td>train/transfer_overall_precision</td><td>▁</td></tr><tr><td>train/transfer_overall_recall</td><td>▁</td></tr><tr><td>train/transfer_runtime</td><td>▁</td></tr><tr><td>train/transfer_samples_per_second</td><td>▁</td></tr><tr><td>train/transfer_steps_per_second</td><td>▁</td></tr><tr><td>train/transfer_token_f1_macro</td><td>▁</td></tr><tr><td>train/transfer_token_precision_macro</td><td>▁</td></tr><tr><td>train/transfer_token_recall_macro</td><td>▁</td></tr></table><br/></div><div class=\"wandb-col\"><h3>Run summary:</h3><br/><table class=\"wandb\"><tr><td>eval/LOC_f1</td><td>0.96628</td></tr><tr><td>eval/LOC_precision</td><td>0.95775</td></tr><tr><td>eval/LOC_recall</td><td>0.97496</td></tr><tr><td>eval/MISC_f1</td><td>0.88407</td></tr><tr><td>eval/MISC_precision</td><td>0.86708</td></tr><tr><td>eval/MISC_recall</td><td>0.90173</td></tr><tr><td>eval/ORG_f1</td><td>0.9174</td></tr><tr><td>eval/ORG_precision</td><td>0.90936</td></tr><tr><td>eval/ORG_recall</td><td>0.9256</td></tr><tr><td>eval/PER_f1</td><td>0.97863</td></tr><tr><td>eval/PER_precision</td><td>0.9752</td></tr><tr><td>eval/PER_recall</td><td>0.98208</td></tr><tr><td>eval/loss</td><td>0.04198</td></tr><tr><td>eval/overall_f1</td><td>0.94611</td></tr><tr><td>eval/overall_precision</td><td>0.93775</td></tr><tr><td>eval/overall_recall</td><td>0.95461</td></tr><tr><td>eval/runtime</td><td>1.7286</td></tr><tr><td>eval/samples_per_second</td><td>1880.171</td></tr><tr><td>eval/steps_per_second</td><td>15.041</td></tr><tr><td>eval/token_f1_macro</td><td>0.95429</td></tr><tr><td>eval/token_precision_macro</td><td>0.95344</td></tr><tr><td>eval/token_recall_macro</td><td>0.9552</td></tr><tr><td>total_flos</td><td>912371233089690.0</td></tr><tr><td>train/epoch</td><td>2.84738</td></tr><tr><td>train/global_step</td><td>2500</td></tr><tr><td>train/grad_norm</td><td>0.12388</td></tr><tr><td>train/learning_rate</td><td>0.0</td></tr><tr><td>train/loss</td><td>0.0268</td></tr><tr><td>train/num_input_tokens_seen</td><td>1950135</td></tr><tr><td>train/transfer_LOC_f1</td><td>0.76362</td></tr><tr><td>train/transfer_LOC_precision</td><td>0.71892</td></tr><tr><td>train/transfer_LOC_recall</td><td>0.81425</td></tr><tr><td>train/transfer_MISC_f1</td><td>0.26916</td></tr><tr><td>train/transfer_MISC_precision</td><td>0.32771</td></tr><tr><td>train/transfer_MISC_recall</td><td>0.22837</td></tr><tr><td>train/transfer_ORG_f1</td><td>0.69894</td></tr><tr><td>train/transfer_ORG_precision</td><td>0.62088</td></tr><tr><td>train/transfer_ORG_recall</td><td>0.79944</td></tr><tr><td>train/transfer_PER_f1</td><td>0.39754</td></tr><tr><td>train/transfer_PER_precision</td><td>0.44286</td></tr><tr><td>train/transfer_PER_recall</td><td>0.36064</td></tr><tr><td>train/transfer_loss</td><td>0.41569</td></tr><tr><td>train/transfer_overall_f1</td><td>0.4846</td></tr><tr><td>train/transfer_overall_precision</td><td>0.50665</td></tr><tr><td>train/transfer_overall_recall</td><td>0.46438</td></tr><tr><td>train/transfer_runtime</td><td>27.5571</td></tr><tr><td>train/transfer_samples_per_second</td><td>325.215</td></tr><tr><td>train/transfer_steps_per_second</td><td>2.576</td></tr><tr><td>train/transfer_token_f1_macro</td><td>0.65074</td></tr><tr><td>train/transfer_token_precision_macro</td><td>0.73181</td></tr><tr><td>train/transfer_token_recall_macro</td><td>0.64536</td></tr><tr><td>train_loss</td><td>0.11646</td></tr><tr><td>train_runtime</td><td>150.9191</td></tr><tr><td>train_samples_per_second</td><td>265.043</td></tr><tr><td>train_steps_per_second</td><td>16.565</td></tr></table><br/></div></div>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ],
      "text/html": [
       " View run <strong style=\"color:#cdcd00\">facebook__xlm-v-base-finetuned-l11-conll03/12-07/1</strong> at: <a href='https://wandb.ai/viktoroo-sch-epfl/ner-alignment/runs/j955yfyx' target=\"_blank\">https://wandb.ai/viktoroo-sch-epfl/ner-alignment/runs/j955yfyx</a><br/> View project at: <a href='https://wandb.ai/viktoroo-sch-epfl/ner-alignment' target=\"_blank\">https://wandb.ai/viktoroo-sch-epfl/ner-alignment</a><br/>Synced 6 W&B file(s), 0 media file(s), 0 artifact file(s) and 0 other file(s)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ],
      "text/html": [
       "Find logs at: <code>facebook__xlm-v-base-finetuned-l11-conll03/12-07/1/wandb/run-20241207_164756-j955yfyx/logs</code>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ],
      "text/html": [
       "wandb version 0.19.0 is available!  To upgrade, please run:\n",
       " $ pip install wandb --upgrade"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ],
      "text/html": [
       "The new W&B backend becomes opt-out in version 0.18.0; try it out with `wandb.require(\"core\")`! See https://wandb.me/wandb-core for more information."
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "execution_count": 21
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-07T16:51:12.975216Z",
     "start_time": "2024-12-07T16:51:12.973639Z"
    }
   },
   "cell_type": "code",
   "source": "",
   "id": "a4de821b3a38e5c8",
   "outputs": [],
   "execution_count": 21
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}