{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Sentiment analysis with BERT\n",
    "\n",
    "Using transformers with the distilled bert-base model on the IMDB dataset, to perform continuous score sentiment analysis."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Retrieving IMDB training and testing dataset from datasets directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pandas as pd\n",
    "\n",
    "imdb_train_dataset = 'aclimdb/train'\n",
    "imdb_test_dataset = 'aclimdb/test'\n",
    "\n",
    "train_reviews = []\n",
    "train_scores = []\n",
    "test_reviews = []\n",
    "test_scores = []\n",
    "\n",
    "for dataset, reviews, scores in [(imdb_train_dataset, train_reviews, train_scores), (imdb_test_dataset, test_reviews, test_scores)]:\n",
    "    for sentiment in ['pos','neg']:\n",
    "        sentiment_dir = os.path.join(dataset,sentiment)\n",
    "\n",
    "        for filename in os.listdir(sentiment_dir):\n",
    "            if filename.endswith('.txt'):\n",
    "                with open(os.path.join(sentiment_dir,filename),'r',encoding='utf-8') as file:\n",
    "                    review = file.read()\n",
    "                    sentiment_score = int(filename[:-4].split('_')[1])\n",
    "\n",
    "                    scores.append(sentiment_score)\n",
    "                    reviews.append(review)\n",
    "\n",
    "df_train = pd.DataFrame({'text': train_reviews, 'sentiment': train_scores})\n",
    "df_test = pd.DataFrame({'text': test_reviews, 'sentiment': test_scores}).sample(5000)\n",
    "\n",
    "df_train.info()\n",
    "print('')\n",
    "df_test.info()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Normalizing the training and testing dataset to a range of -1 to 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def normalize(n):\n",
    "    normal_n = (n - 5) / 5\n",
    "    return normal_n\n",
    "\n",
    "df_train['s'] = normalize(df_train['sentiment'])\n",
    "df_test['s'] = normalize(df_test['sentiment'])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Evaluating the summary statistics of the training and testing dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_train['s'].describe()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_test['s'].describe()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Determining the length of the training and testing dataset, to set a proper batch size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f'Length training set: {len(df_train)}\\nLength testing set: {len(df_test)}')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Preparing the data for BERT, this includes tokenization, encoding and creating dataloaders for both training and testing datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.data import DataLoader, TensorDataset\n",
    "from transformers import (\n",
    "    DistilBertTokenizerFast,\n",
    "    DistilBertForSequenceClassification\n",
    ")\n",
    "\n",
    "# Set device cuda\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "print('Using:', device)\n",
    "\n",
    "# Define model\n",
    "model_name = 'distilbert-base-uncased'\n",
    "tokenizer = DistilBertTokenizerFast.from_pretrained(model_name)\n",
    "\n",
    "# Tokenize and encode the text data\n",
    "train_encodings = tokenizer(\n",
    "    df_train['text'].tolist(), \n",
    "    truncation=True, \n",
    "    padding=True, \n",
    "    return_tensors='pt'\n",
    ")\n",
    "\n",
    "test_encodings = tokenizer(\n",
    "    df_test['text'].tolist(), \n",
    "    truncation=True, \n",
    "    padding=True, \n",
    "    return_tensors='pt'\n",
    ")\n",
    "\n",
    "# Create tensor dataset\n",
    "train_dataset = TensorDataset(\n",
    "    train_encodings['input_ids'], \n",
    "    train_encodings['attention_mask'], \n",
    "    torch.tensor(df_train['s'].values, dtype=torch.float32)\n",
    ")\n",
    "\n",
    "test_dataset = TensorDataset(\n",
    "    test_encodings['input_ids'], \n",
    "    test_encodings['attention_mask'], \n",
    "    torch.tensor(df_test['s'].values, dtype=torch.float32)\n",
    ")\n",
    "\n",
    "# Create data loaders\n",
    "if device == 'cuda':\n",
    "    en_pin_memory = True\n",
    "else:\n",
    "    en_pin_memory = False\n",
    "\n",
    "train_dataloader = DataLoader(\n",
    "    train_dataset, \n",
    "    batch_size=10, \n",
    "    shuffle=True,\n",
    "    pin_memory=en_pin_memory\n",
    ")\n",
    "\n",
    "test_dataloader = DataLoader(\n",
    "    test_dataset, \n",
    "    batch_size=10, \n",
    "    shuffle=False,\n",
    "    pin_memory=en_pin_memory\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Defining the model: distilbert."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = DistilBertForSequenceClassification.from_pretrained(\n",
    "    model_name, \n",
    "    num_labels=1\n",
    ").to(device)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Defining the optimizer and loss function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.optim import AdamW\n",
    "from torch.nn import L1Loss\n",
    "\n",
    "optimizer = AdamW(model.parameters(), lr=1e-5)\n",
    "loss_fn = L1Loss()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Training loop. Here BERT will be trained with the training dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.tensorboard import SummaryWriter\n",
    "\n",
    "log_dir = 'bert-aclimdb/logs'\n",
    "writer = SummaryWriter(log_dir)\n",
    "global_step = 0\n",
    "\n",
    "num_epochs = 30\n",
    "\n",
    "early_stop_patience = 2\n",
    "best_validation_loss = float('inf')\n",
    "no_improvement_counter = 0\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    model.train()\n",
    "    total_loss = 0\n",
    "    num_batches = 0\n",
    "\n",
    "    for batch in train_dataloader:\n",
    "        global_step += 1\n",
    "        num_batches += 1\n",
    "        input_ids, attention_mask, target_scores = [x.to(device) for x in batch]\n",
    "\n",
    "        # Forward pass\n",
    "        output = model(input_ids=input_ids, attention_mask=attention_mask)\n",
    "        predicted_scores = output.logits.view(-1)\n",
    "\n",
    "        # Calculating the loss\n",
    "        loss = loss_fn(predicted_scores, target_scores)\n",
    "\n",
    "        # The total loss per epoch\n",
    "        total_loss += loss.item()\n",
    "\n",
    "        # Determining the average loss in the epoch\n",
    "        average_loss = total_loss / num_batches\n",
    "\n",
    "        # Tensorboard logging\n",
    "        writer.add_scalar('batch-loss-train', average_loss, global_step)\n",
    "\n",
    "        # Backward pass and optimization\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "    # Determining the average loss for the epoch\n",
    "    average_loss = total_loss / len(train_dataloader)\n",
    "\n",
    "    # Logging\n",
    "    writer.add_scalar('epoch-loss-train', average_loss, epoch + 1)\n",
    "\n",
    "    # Validation\n",
    "    model.eval()\n",
    "    total_loss = 0\n",
    "\n",
    "    for batch in test_dataloader:\n",
    "        with torch.no_grad():\n",
    "            input_ids, attention_mask, target_scores = [x.to(device) for x in batch]\n",
    "\n",
    "            # Obtaining the scores\n",
    "            output = model(input_ids=input_ids, attention_mask=attention_mask)\n",
    "            predicted_scores = output.logits.view(-1)\n",
    "\n",
    "            # Calculating the loss\n",
    "            loss = loss_fn(predicted_scores, target_scores)\n",
    "\n",
    "            # The total loss per epoch\n",
    "            total_loss += loss.item()\n",
    "\n",
    "    # Determining the average loss for the epoch\n",
    "    average_loss = total_loss / len(test_dataloader)\n",
    "\n",
    "    # Logging  \n",
    "    writer.add_scalar('epoch-loss-validation', average_loss, epoch + 1)\n",
    "    print(f\"Epoch {epoch + 1}/{num_epochs}, Validation loss: {average_loss:.4f}\\n\")\n",
    "\n",
    "    # Saving the model\n",
    "    torch.save(model.state_dict(), f'bert-aclimdb/{epoch + 1}.pth')\n",
    "\n",
    "    # Early stopping check\n",
    "    if average_loss < best_validation_loss:\n",
    "        best_validation_loss = average_loss\n",
    "        no_improvement_counter = 0\n",
    "    else:\n",
    "        no_improvement_counter += 1\n",
    "\n",
    "    if no_improvement_counter >= early_stop_patience:\n",
    "        print('Early stopping triggered')\n",
    "        break\n",
    "\n",
    "writer.close()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Loading a version of the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = DistilBertForSequenceClassification.from_pretrained(model_name, num_labels=1)\n",
    "model.load_state_dict(torch.load('...', map_location=device))\n",
    "model.to(device)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Evaluating the model, with as output the MAE, MSE and R-value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.stats import pearsonr\n",
    "from sklearn.metrics import mean_squared_error, mean_absolute_error\n",
    "\n",
    "model.eval()\n",
    "list_predicted_scores = []\n",
    "\n",
    "for batch in test_dataloader:\n",
    "    with torch.no_grad():\n",
    "        input_ids, attention_mask, scores = [x.to(device) for x in batch]\n",
    "\n",
    "        # Obtaining the scores\n",
    "        output = model(input_ids=input_ids, attention_mask=attention_mask)\n",
    "        predicted_scores = output.logits.view(-1)\n",
    "\n",
    "        # Writing the scores to the list\n",
    "        list_predicted_scores.extend(predicted_scores.tolist())\n",
    "    \n",
    "# Inserting the scores in df_test\n",
    "df_test['sp'] = list_predicted_scores\n",
    "\n",
    "# Computing the R, MSE and MAE vlaues\n",
    "correlation, _ = pearsonr(df_test['s'], df_test['sp'])\n",
    "print(f\"Pearson Correlation Coefficient (R) s: {correlation:.4f}\")\n",
    "print(f\"Mean Absolute Error (MAE) s: {mean_absolute_error(df_test['s'], df_test['sp']):.4f}\")\n",
    "print(f\"Root mean Squared Error (RMSE) s: {(mean_squared_error(df_test['s'], df_test['sp'])**(1/2)):.4f}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Evaluating the summary statistics of the testing dataset and the predicted values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_test[['s', 'sp']].describe()"
   ]
  }
 ],
 "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.14.6"
  },
  "vscode": {
   "interpreter": {
    "hash": "8c75c0fdd1a718867cdcb84b32adcfdbeaad00b3a4e00a59385211aeed084d4c"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}