arabert_finetuning.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "This notebook contains finetuning an [AraBERT variant](https://huggingface.co/aubmindlab/bert-large-arabertv02-twitter), which is pre-trained on Arabic Dialects from tweets, for Arabic Sentiment Analysis. \n",
        "\n",
        "It achieved state-of-the-art performance and the first place solution on a Kaggle university-wide NLP competition.\n",
        "\n",
        "![Kaggle competition leaderboard](kaggle_leaderboard.png  \"First place on a university-wide Kaggle competition\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Import necessary Libraries "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "WotCaOwKvg6G",
        "outputId": "b26b5ccf-e19c-44b3-c3ca-6a3571c29439"
      },
      "outputs": [],
      "source": [
        "!pip install transformers[torch]\n",
        "!pip install accelerate -U\n",
        "!pip install farasapy\n",
        "!pip install pyarabic\n",
        "!pip install arabert"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZR2ujvP7rG8d"
      },
      "outputs": [],
      "source": [
        "from transformers import BertTokenizer, BertForSequenceClassification, Trainer, TrainingArguments\n",
        "from arabert.preprocess import ArabertPreprocessor\n",
        "from torch.utils.data import Dataset, DataLoader \n",
        "from torch.utils.data import Dataset\n",
        "import torch\n",
        "import pandas as pd"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### GPU "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yQLEStOd5Fd4"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "os.environ['CUDA_LAUNCH_BLOCKING'] = \"1\"\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Preparing the pre-trained model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "model_name = \"aubmindlab/bert-large-arabertv02-twitter\"\n",
        "tokenizer = BertTokenizer.from_pretrained(model_name)\n",
        "# Labels = {-1, 0, 1}\n",
        "model = BertForSequenceClassification.from_pretrained(model_name, num_labels=3).to(device) \n",
        "arabert_prep = ArabertPreprocessor(model_name=model_name)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "class SimpleDataset(Dataset):\n",
        "    def __init__(self, input_ids, attention_masks, labels=None):\n",
        "        self.input_ids = input_ids\n",
        "        self.attention_masks = attention_masks\n",
        "        self.labels = labels\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.input_ids)\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        data = {\n",
        "            'input_ids': self.input_ids[idx],\n",
        "            'attention_mask': self.attention_masks[idx],\n",
        "        }\n",
        "        if self.labels is not None:\n",
        "            data['labels'] = self.labels[idx]\n",
        "        return data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "def prepare_dataset(df, tokenizer, max_len=64, include_labels=True):\n",
        "    input_ids = []\n",
        "    attention_masks = []\n",
        "    labels = []\n",
        "\n",
        "    for _, row in df.iterrows():\n",
        "        encoded_data = tokenizer.encode_plus(\n",
        "            row['review_description'],\n",
        "            add_special_tokens=True,\n",
        "            max_length=max_len,\n",
        "            padding='max_length',\n",
        "            truncation=True,\n",
        "            return_attention_mask=True,\n",
        "            return_tensors='pt',\n",
        "        )\n",
        "        input_ids.append(encoded_data['input_ids'])\n",
        "        attention_masks.append(encoded_data['attention_mask'])\n",
        "\n",
        "        # Add label if 'rating' column is present and include_labels is True\n",
        "        if include_labels and 'rating' in df.columns:\n",
        "            labels.append(row['rating'])\n",
        "\n",
        "    input_ids = torch.cat(input_ids, dim=0)\n",
        "    attention_masks = torch.cat(attention_masks, dim=0)\n",
        "\n",
        "    if labels:\n",
        "        labels = torch.tensor(labels, dtype=torch.long)\n",
        "        return SimpleDataset(input_ids, attention_masks, labels)\n",
        "    else:\n",
        "        return SimpleDataset(input_ids, attention_masks, None)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Loading training dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "h7QO78GQx00J"
      },
      "outputs": [],
      "source": [
        "train_df = pd.read_excel(\"datasets/train.xlsx\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PkhjwW4m4ivF"
      },
      "outputs": [],
      "source": [
        "# It is required for the labels to start from 0\n",
        "train_df['rating'] = train_df['rating']+1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2rcq1CYgvvxF"
      },
      "outputs": [],
      "source": [
        "train_df['review_description'] = train_df['review_description'].apply(arabert_prep.preprocess)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-tRjOZ76yHpM"
      },
      "outputs": [],
      "source": [
        "train_dataset = prepare_dataset(train_df, tokenizer, include_labels=True)\n",
        "train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Finetuning the model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AOfetjL0yJPe"
      },
      "outputs": [],
      "source": [
        "training_args = TrainingArguments(\n",
        "    output_dir='./results',\n",
        "    num_train_epochs=3,\n",
        "    per_device_train_batch_size=16,\n",
        "    warmup_steps=500,\n",
        "    weight_decay=0.01,\n",
        "    logging_dir='./logs',\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ssamy2s-yN4w"
      },
      "outputs": [],
      "source": [
        "trainer = Trainer(\n",
        "    model=model,\n",
        "    args=training_args,\n",
        "    train_dataset=train_loader.dataset\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 503
        },
        "id": "wR_6tIom2DZz",
        "outputId": "fe0bc10d-279f-486d-f602-a0939e5903e9"
      },
      "outputs": [],
      "source": [
        "trainer.train()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Loading test dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oa-WQb072FgK"
      },
      "outputs": [],
      "source": [
        "test_df = pd.read_csv(\"datasets/test.csv\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "test_df['review_description'] = test_df['review_description'].apply(lambda x: arabert_prep.preprocess(x))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "test_dataset = prepare_dataset(test_df, tokenizer, include_labels=False)\n",
        "test_loader = DataLoader(test_dataset, batch_size=16)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Rating prediction"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "PqduFC7arZx6",
        "outputId": "7d3e3144-d5da-4bd1-efaa-78edf68da4e1"
      },
      "outputs": [],
      "source": [
        "def predict_sentiments(model, data_loader):\n",
        "    model.eval()\n",
        "    predictions = []\n",
        "\n",
        "    with torch.no_grad():\n",
        "        for batch in data_loader:\n",
        "            input_ids = batch['input_ids'].to(device)\n",
        "            attention_mask = batch['attention_mask'].to(device)\n",
        "\n",
        "            outputs = model(input_ids, attention_mask=attention_mask)\n",
        "            _, preds = torch.max(outputs.logits, dim=1)\n",
        "\n",
        "            predictions.extend(preds.tolist())\n",
        "\n",
        "    return predictions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "test_predictions = predict_sentiments(model, test_loader)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "test_df['rating'] = test_predictions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "test_df.head()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "test_df.to_csv('predicted_arabert_twitter.csv', index=False)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.11.6"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}