q2.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import modules\n",
    "import pandas as pd\n",
    "from datasets import Dataset, load_metric, ClassLabel\n",
    "import numpy as np\n",
    "from transformers import AutoTokenizer, AutoModelForTokenClassification, TrainingArguments, Trainer, DataCollatorForTokenClassification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set model constants and all that\n",
    "batch_size=16\n",
    "label_list=['O', 'B-CW', 'I-CW', 'B-PROD', 'I-PROD', 'B-PER', 'I-PER', 'B-LOC', 'I-LOC', 'B-CORP', 'I-CORP', 'B-GRP', 'I-GRP']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [],
   "source": [
    "# also could create manual dict from labels but i have used `ClassLabel` \n",
    "CLASS_LABELS=ClassLabel(names=label_list,num_classes = len(label_list))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# on indic-bert model page this is explained\n",
    "tokenizer=AutoTokenizer.from_pretrained('ai4bharat/indic-bert')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [],
   "source": [
    "# first task is to pre-process and create Dataset that can be stremed into the model\n",
    "train_file_path='data/hi_train.conll'\n",
    "dev_file_path='data/hi_dev.conll'\n",
    "\n",
    "def get_dataset(file_path):\n",
    "    \"\"\"\n",
    "    simple function to get ner data into nice dataset that can be used with indic-bert for tuning purposes\n",
    "    \"\"\"\n",
    "    split_mark=' _ _ '\n",
    "    tokens=[]\n",
    "    ner_tags=[]\n",
    "    with open(file_path,'r') as file:\n",
    "        lines=file.readlines()\n",
    "        tmp_tokens=[]\n",
    "        tmp_ner_tags=[]\n",
    "        for i,line in enumerate(lines):\n",
    "            split_=line.split(split_mark)\n",
    "            if len(split_)>1:\n",
    "                tmp_tokens.append(split_[0])\n",
    "                tmp_ner_tags.append(split_[1].strip())\n",
    "            if line=='\\n':\n",
    "                if len(tmp_tokens)>0:\n",
    "                    tokens.append(tmp_tokens)\n",
    "                    ner_tags.append(tmp_ner_tags)\n",
    "                tmp_tokens=[]\n",
    "                tmp_ner_tags=[]\n",
    "    \n",
    "    df=pd.DataFrame({'tokens':tokens,'ner_tags':ner_tags})\n",
    "    return Dataset.from_pandas(df)\n",
    "\n",
    "train_dataset=get_dataset(train_file_path)\n",
    "dev_dataset=get_dataset(dev_file_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# need to tokenize these datasets to convert the tokens and labels into numeric form for our BERT model\n",
    "def tokenize_and_align_labels(examples):\n",
    "    label_all_tokens = True\n",
    "    tokenized_inputs = tokenizer(list(examples[\"tokens\"]), truncation=True, is_split_into_words=True)\n",
    "\n",
    "    labels = []\n",
    "    for i, label in enumerate(examples['ner_tags']):\n",
    "        word_ids = tokenized_inputs.word_ids(batch_index=i)\n",
    "        previous_word_idx = None\n",
    "        label_ids = []\n",
    "        for word_idx in word_ids:\n",
    "            if word_idx is None:\n",
    "                label_ids.append(-100)\n",
    "            elif label[word_idx] == '0':\n",
    "                label_ids.append(0)\n",
    "            elif word_idx != previous_word_idx:\n",
    "                label_ids.append(CLASS_LABELS.str2int(label[word_idx]))\n",
    "            else:\n",
    "                label_ids.append(CLASS_LABELS.str2int(label[word_idx]) if label_all_tokens else -100)\n",
    "            previous_word_idx = word_idx\n",
    "        labels.append(label_ids)\n",
    "        \n",
    "    tokenized_inputs[\"labels\"] = labels\n",
    "    return tokenized_inputs\n",
    "\n",
    "\n",
    "train_tokenized_datasets = train_dataset.map(tokenize_and_align_labels, batched=True)\n",
    "dev_tokenized_datasets = dev_dataset.map(tokenize_and_align_labels, batched=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fine tune the model for ner task\n",
    "model = AutoModelForTokenClassification.from_pretrained('ai4bharat/indic-bert', num_labels=len(label_list))\n",
    "\n",
    "training_args = TrainingArguments(\n",
    "    \"test-ner\",\n",
    "    evaluation_strategy = \"epoch\",\n",
    "    learning_rate=2e-5,\n",
    "    per_device_train_batch_size=batch_size,\n",
    "    per_device_eval_batch_size=batch_size,\n",
    "    num_train_epochs=2,\n",
    "    weight_decay=0.01,\n",
    ")\n",
    "\n",
    "data_collator = DataCollatorForTokenClassification(tokenizer)\n",
    "metric = load_metric(\"seqeval\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "def compute_metrics(p):\n",
    "    predictions, labels = p\n",
    "    predictions = np.argmax(predictions, axis=2)\n",
    "\n",
    "    true_predictions = [[label_list[p] for (p, l) in zip(prediction, label) if l != -100] for prediction, label in zip(predictions, labels)]\n",
    "    true_labels = [[label_list[l] for (p, l) in zip(prediction, label) if l != -100] for prediction, label in zip(predictions, labels)]\n",
    "\n",
    "    results = metric.compute(predictions=true_predictions, references=true_labels)\n",
    "    return results[\"overall_f1\"]\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=train_tokenized_datasets,\n",
    "    eval_dataset=dev_tokenized_datasets,\n",
    "    data_collator=data_collator,\n",
    "    tokenizer=tokenizer,\n",
    "    compute_metrics=compute_metrics\n",
    ")\n",
    "\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "trainer.evaluate()\n",
    "trainer.save_model('hindi-ner.model')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "253"
      ]
     },
     "execution_count": 78,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import gc\n",
    "gc.collect()"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "365d70965140afb04a698773bfdd31483bc82432b779112c2a78b5de7c16d125"
  },
  "kernelspec": {
   "display_name": "Python 3.8.5 ('base')",
   "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.8.5"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}