Model AraBERT.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "08af2a81",
   "metadata": {},
   "source": [
    "# Modeling & Training\n",
    "## AraBERT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a67f484f",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install transformers==4.12.2\n",
    "!pip install farasapy==0.0.14\n",
    "!pip install pyarabic==0.6.14\n",
    "!git clone https://github.com/aub-mind/arabert\n",
    "!pip install emoji==1.6.1\n",
    "!pip install sentencepiece==0.1.96"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "950127aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "import math\n",
    "import copy\n",
    "import random\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from datetime import date\n",
    "from transformers import *\n",
    "from sklearn.metrics import *\n",
    "from tqdm import tqdm_notebook as tqdm\n",
    "\n",
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torch.utils.data.sampler import RandomSampler\n",
    "\n",
    "from arabert.preprocess import ArabertPreprocessor\n",
    "\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.metrics import (accuracy_score, classification_report,\n",
    "                             confusion_matrix, f1_score, precision_score,\n",
    "                             recall_score)\n",
    "from transformers import (AutoConfig, AutoModelForSequenceClassification,\n",
    "                          AutoTokenizer, BertTokenizer, Trainer,\n",
    "                          TrainingArguments)\n",
    "from transformers.data.processors.utils import InputFeatures\n",
    "\n",
    "os.environ['WANDB_DISABLED'] = 'true'\n",
    "os.environ['TOKENIZERS_PARALLELISM'] = 'false'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "647dd096",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Expect an update for supporting TPU training using Pytorch/XLA\n",
    "\n",
    "# If there's a GPU available...\n",
    "if torch.cuda.is_available():\n",
    "\n",
    "    # Tell PyTorch to use the GPU.    \n",
    "    device = torch.device('cuda')\n",
    "\n",
    "    print('There are %d GPU(s) available.' % torch.cuda.device_count())\n",
    "\n",
    "    print('We will use the GPU:', torch.cuda.get_device_name(0))\n",
    "    !nvidia-smi\n",
    "\n",
    "# If not...from sklearn.model_selection import train_test_split\n",
    "else:\n",
    "    print('No GPU available, using the CPU instead.')\n",
    "    device = torch.device('cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1d52bcf",
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomDataset:\n",
    "    def __init__(self, train, test, label_list):\n",
    "        '''Class to hold and structure datasets.\n",
    "\n",
    "        Args:\n",
    "\n",
    "        train (List[pd.DataFrame]): holds training pandas dataframe with 2 columns ['text','label']\n",
    "        test (List[pd.DataFrame]): holds testing pandas dataframe with 2 columns ['text','label']\n",
    "        label_list (List[str]): holds the list  of labels\n",
    "        '''\n",
    "        self.train = train\n",
    "        self.test = test\n",
    "        self.label_list = label_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02136c0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    rawData = pd.read_json('arabicDialects.json', encoding='utf-8',\n",
    "                           typ='series', convert_axes=False)\n",
    "except Exception as e:\n",
    "    %run -i fetch.py\n",
    "    rawData = pd.read_json('arabicDialects.json', encoding='utf-8',\n",
    "                           typ='series', convert_axes=False)\n",
    "\n",
    "rawData.index = rawData.index.astype('int64')\n",
    "\n",
    "datasetIds = pd.read_csv('dialect_dataset.csv', index_col='id')\n",
    "datasetIds = datasetIds.loc[rawData.index]\n",
    "datasetIds['text'] = rawData\n",
    "\n",
    "dialects = sorted(datasetIds.dialect.unique())\n",
    "label_map = dict(zip(dialects, range(len(dialects))))\n",
    "inv_label_map = dict(enumerate(dialects))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd4e16bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_COLUMN, LABEL_COLUMN = 'text', 'dialect'\n",
    "\n",
    "print('Total length: ', len(rawData))\n",
    "print(datasetIds[LABEL_COLUMN].value_counts())\n",
    "\n",
    "train_data, test_data = train_test_split(datasetIds, test_size=.1, random_state=42)\n",
    "print('Training length: ', len(train_data))\n",
    "print('Testing length: ', len(test_data))\n",
    "\n",
    "dialects_dataset = CustomDataset(train_data, test_data, dialects)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a76de6ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "# https://huggingface.co/aubmindlab/bert-base-arabertv02/blob/main/pytorch_model.bin\n",
    "AraBert_v02_Pytorch = 'aubmindlab/bert-base-arabertv02'\n",
    "\n",
    "arabic_prep = ArabertPreprocessor(AraBert_v02_Pytorch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19295d5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "dialects_dataset.train.loc[:,DATA_COLUMN] = \\\n",
    "    dialects_dataset.train.loc[:,DATA_COLUMN].apply(lambda x: arabic_prep.preprocess(x))\n",
    "dialects_dataset.test.loc[:,DATA_COLUMN] = \\\n",
    "    dialects_dataset.test.loc[:,DATA_COLUMN].apply(lambda x: arabic_prep.preprocess(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff70d3d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "max_len = 90\n",
    "tok = AutoTokenizer.from_pretrained(AraBert_v02_Pytorch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9431f61e",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_len_hist = [len(tok.tokenize(sentence)) for sentence in \\\n",
    "                  dialects_dataset.train[DATA_COLUMN]]\n",
    "test_len_hist = [len(tok.tokenize(sentence)) for sentence in \\\n",
    "                 dialects_dataset.test[DATA_COLUMN]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0aab4264",
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Training Sentence Lengths: ')\n",
    "plt.hist(train_len_hist,bins=range(0,max_len,2))\n",
    "ax = plt.gca()\n",
    "plt.vlines(max_len, *ax.get_ylim(), colors='red')\n",
    "plt.show()\n",
    "\n",
    "print('Testing Sentence Lengths: ')\n",
    "plt.hist(test_len_hist,bins=range(0,max_len,2))\n",
    "ax = plt.gca()\n",
    "plt.vlines(max_len, *ax.get_ylim(), colors='red')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d86dc92b",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f'At max len of {max_len}, there are:')\n",
    "\n",
    "trunc_train_seq = len(list(filter(lambda x:x>max_len, train_len_hist)))\n",
    "print(f'Truncated training sequences: {trunc_train_seq} / {len(train_data)} ~ '\n",
    "      f'{trunc_train_seq/len(train_data):.2f}%')\n",
    "\n",
    "trunc_test_seq = len(list(filter(lambda x:x>max_len, test_len_hist)))\n",
    "print(f'Truncated testing sequences: {trunc_test_seq} / {len(test_data)} ~ '\n",
    "      f'{trunc_test_seq/len(test_data):.2f}%')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2821f8a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ClassificationDataset(Dataset):\n",
    "    def __init__(self, text, target, model_name, max_len, label_map):\n",
    "        super(ClassificationDataset).__init__()\n",
    "        '''\n",
    "        Args:\n",
    "        text (List[str]): List of the training text\n",
    "        target (List[str]): List of the training labels\n",
    "        tokenizer_name (str): The tokenizer name (same as model_name).\n",
    "        max_len (int): Maximum sentence length\n",
    "        label_map (Dict[str,int]): A dictionary that maps the class labels to integer\n",
    "        '''\n",
    "        self.text = text\n",
    "        self.target = target\n",
    "        self.tokenizer_name = model_name\n",
    "        self.tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "        self.max_len = max_len\n",
    "        self.label_map = label_map\n",
    "\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.text)\n",
    "\n",
    "    def __getitem__(self,item):\n",
    "        text = str(self.text[item])\n",
    "        text = ' '.join(text.split())\n",
    "\n",
    "        inputs = self.tokenizer(\n",
    "            text,\n",
    "            max_length=self.max_len,\n",
    "            padding='max_length',\n",
    "            truncation=True\n",
    "        )      \n",
    "        return InputFeatures(**inputs,label=self.label_map[self.target[item]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0cef3784",
   "metadata": {},
   "outputs": [],
   "source": [
    "label_map = dict(zip(dialects, range(len(dialects))))\n",
    "print(label_map)\n",
    "\n",
    "train_dataset = ClassificationDataset(\n",
    "    dialects_dataset.train[DATA_COLUMN].to_list(),\n",
    "    dialects_dataset.train[LABEL_COLUMN].to_list(),\n",
    "    AraBert_v02_Pytorch,\n",
    "    max_len,\n",
    "    label_map\n",
    "  )\n",
    "\n",
    "test_dataset = ClassificationDataset(\n",
    "    dialects_dataset.test[DATA_COLUMN].to_list(),\n",
    "    dialects_dataset.test[LABEL_COLUMN].to_list(),\n",
    "    AraBert_v02_Pytorch,\n",
    "    max_len,\n",
    "    label_map\n",
    "  )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be9da92d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def model_init():\n",
    "    return AutoModelForSequenceClassification. \\\n",
    "        from_pretrained(AraBert_v02_Pytorch, return_dict=True, num_labels=len(label_map))\n",
    "\n",
    "def set_seed(seed=42):\n",
    "    random.seed(seed)\n",
    "    np.random.seed(seed)\n",
    "    torch.manual_seed(seed)\n",
    "    torch.cuda.manual_seed(seed)\n",
    "    torch.cuda.manual_seed_all(seed)\n",
    "    torch.backends.cudnn.deterministic=True\n",
    "    torch.backends.cudnn.benchmark = False\n",
    "    \n",
    "def compute_metrics(p): #p should be of type EvalPrediction\n",
    "    preds = np.argmax(p.predictions, axis=1)\n",
    "    assert len(preds) == len(p.label_ids)\n",
    "    \n",
    "    macro_f1 = f1_score(p.label_ids, preds, average='macro')\n",
    "    \n",
    "    acc = accuracy_score(p.label_ids, preds)\n",
    "    return {\n",
    "      'macro_f1' : macro_f1,\n",
    "      'accuracy': acc\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8c139d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "training_args = TrainingArguments( \n",
    "    output_dir= './train',    \n",
    "    adam_epsilon = 1e-8,\n",
    "    learning_rate = 2e-5,\n",
    "    fp16 = False, # enable this when using V100 or T4 GPU # Kaggle runs on P100\n",
    "    per_device_train_batch_size = 64, # up to 64 on 16GB with max len of 128\n",
    "    per_device_eval_batch_size = 128,\n",
    "    gradient_accumulation_steps = 2, # use this to scale batch size without needing more memory\n",
    "    num_train_epochs = 1,\n",
    "    warmup_ratio = 0,\n",
    "    do_eval = True,\n",
    "    evaluation_strategy = 'epoch',\n",
    "    save_strategy = 'epoch',\n",
    "    load_best_model_at_end = True, \n",
    "    metric_for_best_model = 'macro_f1',\n",
    "    greater_is_better = True,\n",
    "    seed = 42\n",
    "  )\n",
    "\n",
    "set_seed(training_args.seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83e27e86",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = Trainer(\n",
    "    model = model_init(),\n",
    "    args = training_args,\n",
    "    train_dataset = train_dataset,\n",
    "    eval_dataset = test_dataset,\n",
    "    compute_metrics = compute_metrics,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f6c6046",
   "metadata": {},
   "outputs": [],
   "source": [
    "#start the training\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa12bd67",
   "metadata": {},
   "outputs": [],
   "source": [
    "inv_label_map = dict(enumerate(dialects))\n",
    "print(inv_label_map)\n",
    "\n",
    "trainer.model.config.label2id = label_map\n",
    "trainer.model.config.id2label = inv_label_map\n",
    "trainer.save_model('output_dir')\n",
    "train_dataset.tokenizer.save_pretrained('output_dir')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "212e26bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "!tar -czvf ArabiceDialect_BERT.tar.gz output_dir/*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd049fdf",
   "metadata": {},
   "outputs": [],
   "source": [
    "pipe = pipeline('sentiment-analysis', model=f'output_dir/', \n",
    "                device=0, return_all_scores =True, max_length=max_len, \n",
    "                truncation=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d68503e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "preds = pipe('يسطا رحت فين')\n",
    "sorted([(x['score'], x['label']) for x in preds[0]], reverse=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe7b2dce",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}