train.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import evaluate\n",
    "import datetime\n",
    "from transformers import AutoTokenizer, DataCollatorForTokenClassification, TrainingArguments, Trainer, AutoModelForTokenClassification, pipeline, EarlyStoppingCallback\n",
    "\n",
    "from tools.labeling_tool import LABELS\n",
    "from utils.train_utils import load_data, create_dataset_dict, generate_tokenized_datasets, compute_metrics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Configs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "BASE_MODEL_NAME = \"neuralmind/bert-base-portuguese-cased\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The file must be in format of the output of the labelling tool\n",
    "a txt file with a list of json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = load_data(\"train_data.txt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "504"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'name': 'Kit Shampoo Revitay Óleo de Coco Novex 300ml e Condicionador Revitay Óleo de Coco Novex 300ml',\n",
       " 'tags': [['O', 0, 3],\n",
       "  ['B-PRO', 4, 11],\n",
       "  ['B-ESP', 12, 19],\n",
       "  ['B-ESP', 20, 24],\n",
       "  ['I-ESP', 25, 27],\n",
       "  ['I-ESP', 28, 32],\n",
       "  ['B-MAR', 33, 38],\n",
       "  ['B-TAM', 39, 44],\n",
       "  ['O', 45, 46],\n",
       "  ['B-PRO', 47, 60],\n",
       "  ['B-ESP', 61, 68],\n",
       "  ['B-ESP', 69, 73],\n",
       "  ['I-ESP', 74, 76],\n",
       "  ['I-ESP', 77, 81],\n",
       "  ['B-MAR', 82, 87],\n",
       "  ['B-TAM', 88, 93]]}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create Datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['id', 'ner_tags', 'tokens'],\n",
       "        num_rows: 403\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['id', 'ner_tags', 'tokens'],\n",
       "        num_rows: 101\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset = create_dataset_dict(data, test_size=0.2)\n",
    "dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of BertForTokenClassification were not initialized from the model checkpoint at neuralmind/bert-base-portuguese-cased 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"
     ]
    }
   ],
   "source": [
    "id2label = {i: label for i, label in enumerate(LABELS)}\n",
    "label2id = {v: k for k, v in id2label.items()}\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_NAME, max_length=512, truncation=True)\n",
    "model = AutoModelForTokenClassification.from_pretrained(\n",
    "    BASE_MODEL_NAME,\n",
    "    id2label=id2label,\n",
    "    label2id=label2id,\n",
    ")\n",
    "data_collator = DataCollatorForTokenClassification(tokenizer=tokenizer)\n",
    "metric = evaluate.load(\"seqeval\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b4483f8cb2b9453c942bfcf7a1bedb36",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/403 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Asking to truncate to max_length but no maximum length is provided and the model has no predefined maximum length. Default to no truncation.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4d53aac41c544b8fb4535a2fd1d56134",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/101 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "tokenized_datasets = generate_tokenized_datasets(tokenizer, dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_model_dir = f\"bert_finetuned_ner_{int(datetime.datetime.now().timestamp())}\"\n",
    "args = TrainingArguments(\n",
    "    output_dir=output_model_dir,\n",
    "    evaluation_strategy=\"epoch\",\n",
    "    save_strategy=\"epoch\",\n",
    "    learning_rate=2e-5,\n",
    "    num_train_epochs=10,\n",
    "    weight_decay=0.01,\n",
    "    push_to_hub=False,\n",
    "    metric_for_best_model=\"precision\",\n",
    "    load_best_model_at_end=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c585f1125d994579af97d71133badd68",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/510 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "You're using a BertTokenizerFast tokenizer. Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the text followed by a call to the `pad` method to get a padded encoding.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f0df2dad33234185820ec16767fe7eb2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/13 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/stelios/.virtualenvs/ita/lib/python3.9/site-packages/seqeval/metrics/v1.py:57: UndefinedMetricWarning: Precision and F-score are ill-defined and being set to 0.0 in labels with no predicted samples. Use `zero_division` parameter to control this behavior.\n",
      "  _warn_prf(average, modifier, msg_start, len(result))\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.662858247756958, 'eval_precision': 0.5891472868217055, 'eval_recall': 0.7364341085271318, 'eval_f1': 0.6546080964685616, 'eval_accuracy': 0.7921122994652406, 'eval_runtime': 2.4376, 'eval_samples_per_second': 41.434, 'eval_steps_per_second': 5.333, 'epoch': 1.0}\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5a618a8505334d07a91eac13e16d35cb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/13 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.48267507553100586, 'eval_precision': 0.7084048027444254, 'eval_recall': 0.8003875968992248, 'eval_f1': 0.751592356687898, 'eval_accuracy': 0.8348930481283422, 'eval_runtime': 4.3304, 'eval_samples_per_second': 23.323, 'eval_steps_per_second': 3.002, 'epoch': 2.0}\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9097d698d72143ef92c02b6afbc4b6af",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/13 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.4402821362018585, 'eval_precision': 0.7574692442882249, 'eval_recall': 0.8352713178294574, 'eval_f1': 0.7944700460829494, 'eval_accuracy': 0.8542780748663101, 'eval_runtime': 1.2369, 'eval_samples_per_second': 81.653, 'eval_steps_per_second': 10.51, 'epoch': 3.0}\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7f69bdd4bd1f490482ed73f07eb035f1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/13 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.4186031222343445, 'eval_precision': 0.7851985559566786, 'eval_recall': 0.8430232558139535, 'eval_f1': 0.8130841121495327, 'eval_accuracy': 0.8656417112299465, 'eval_runtime': 1.2529, 'eval_samples_per_second': 80.612, 'eval_steps_per_second': 10.376, 'epoch': 4.0}\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f745316d71254f73af89fedb9a6754e9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/13 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.43677565455436707, 'eval_precision': 0.7938517179023508, 'eval_recall': 0.8507751937984496, 'eval_f1': 0.8213283442469597, 'eval_accuracy': 0.8676470588235294, 'eval_runtime': 1.9076, 'eval_samples_per_second': 52.946, 'eval_steps_per_second': 6.815, 'epoch': 5.0}\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "46f03831dc434674ac5f186423a7ddef",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/13 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.4555143713951111, 'eval_precision': 0.8055045871559633, 'eval_recall': 0.8507751937984496, 'eval_f1': 0.827521206409048, 'eval_accuracy': 0.8676470588235294, 'eval_runtime': 1.5544, 'eval_samples_per_second': 64.977, 'eval_steps_per_second': 8.363, 'epoch': 6.0}\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "63e22146741d4e58ac06722015ca706e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/13 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.4986459016799927, 'eval_precision': 0.8022181146025879, 'eval_recall': 0.8410852713178295, 'eval_f1': 0.8211920529801325, 'eval_accuracy': 0.8582887700534759, 'eval_runtime': 1.2289, 'eval_samples_per_second': 82.188, 'eval_steps_per_second': 10.579, 'epoch': 7.0}\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d903be78bd004c788a9b8e4d649f2980",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/13 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.48684245347976685, 'eval_precision': 0.8264925373134329, 'eval_recall': 0.8585271317829457, 'eval_f1': 0.8422053231939164, 'eval_accuracy': 0.8709893048128342, 'eval_runtime': 1.2756, 'eval_samples_per_second': 79.179, 'eval_steps_per_second': 10.191, 'epoch': 8.0}\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ef85f9f93e2e4ee2aed279732d01fe80",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/13 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.5034500360488892, 'eval_precision': 0.8114602587800369, 'eval_recall': 0.8507751937984496, 'eval_f1': 0.8306527909176915, 'eval_accuracy': 0.8643048128342246, 'eval_runtime': 2.2729, 'eval_samples_per_second': 44.436, 'eval_steps_per_second': 5.719, 'epoch': 9.0}\n",
      "{'loss': 0.4004, 'learning_rate': 3.921568627450981e-07, 'epoch': 9.8}\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a0045c2efb754b98aafaa4b23329e845",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/13 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.49399346113204956, 'eval_precision': 0.8121546961325967, 'eval_recall': 0.8546511627906976, 'eval_f1': 0.8328611898016998, 'eval_accuracy': 0.8716577540106952, 'eval_runtime': 1.2592, 'eval_samples_per_second': 80.211, 'eval_steps_per_second': 10.324, 'epoch': 10.0}\n",
      "{'train_runtime': 76.0616, 'train_samples_per_second': 52.983, 'train_steps_per_second': 6.705, 'train_loss': 0.3945703628016453, 'epoch': 10.0}\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=510, training_loss=0.3945703628016453, metrics={'train_runtime': 76.0616, 'train_samples_per_second': 52.983, 'train_steps_per_second': 6.705, 'train_loss': 0.3945703628016453, 'epoch': 10.0})"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=args,\n",
    "    train_dataset=tokenized_datasets[\"train\"],\n",
    "    eval_dataset=tokenized_datasets[\"test\"],\n",
    "    data_collator=data_collator,\n",
    "    compute_metrics=compute_metrics,\n",
    "    tokenizer=tokenizer,\n",
    "    callbacks=[EarlyStoppingCallback(early_stopping_patience=2)]\n",
    ")\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ca6d3547882d4a8782591558ba5d9e9b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/13 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "{'eval_loss': 0.48684245347976685,\n",
       " 'eval_precision': 0.8264925373134329,\n",
       " 'eval_recall': 0.8585271317829457,\n",
       " 'eval_f1': 0.8422053231939164,\n",
       " 'eval_accuracy': 0.8709893048128342,\n",
       " 'eval_runtime': 1.8056,\n",
       " 'eval_samples_per_second': 55.936,\n",
       " 'eval_steps_per_second': 7.2,\n",
       " 'epoch': 10.0}"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.evaluate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "fine_tunned_model_path = \"best_bert_finetuned_ner\"\n",
    "trainer.save_model(fine_tunned_model_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Validation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    }
   ],
   "source": [
    "token_classifier = pipeline(\n",
    "    \"token-classification\", model=fine_tunned_model_path, aggregation_strategy=\"simple\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "transformers.pipelines.token_classification.TokenClassificationPipeline"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(token_classifier)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_predict(text):\n",
    "    tokens = token_classifier(text)\n",
    "    print(text)\n",
    "    for token in tokens:\n",
    "        if token:\n",
    "            print(f\"{token['word']} : {token['entity_group']}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Asking to truncate to max_length but no maximum length is provided and the model has no predefined maximum length. Default to no truncation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Coca Cola 1L\n",
      "Coca Cola : PRO\n",
      "1L : TAM\n"
     ]
    }
   ],
   "source": [
    "show_predict(\"Coca Cola 1L\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'entity_group': 'PRO',\n",
       "  'score': 0.9833276,\n",
       "  'word': 'Coca Cola',\n",
       "  'start': 0,\n",
       "  'end': 9},\n",
       " {'entity_group': 'TAM',\n",
       "  'score': 0.98347676,\n",
       "  'word': '1L',\n",
       "  'start': 10,\n",
       "  'end': 12}]"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "token_classifier(\"Coca Cola 1L\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ita",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}