{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Assignment 3: Fine-Tuning a Language Model for Text Classification\n",
    "\n",
    "### Exercise 1: Introduction to Fine-Tuning a BERT Model for Text Classification\n",
    "In this exercise, you will gain practical experience in fine-tuning a pre-trained language model for a text classification task. You will use the `bert-base-cased` model and the `yelp_review_full` dataset, which contains Yelp business reviews classified into five-star ratings. You are also free to choose a different pre-trained model or dataset if you prefer, but in that case, you will need to adapt the provided code accordingly.\n",
    "\n",
    "Your goal is to fine-tune the model to predict the rating based on the text of the review. You are required to complete the following tasks:\n",
    "- Load and prepare a text dataset for fine-tuning.\n",
    "- Fine-tune a pre-trained language model for text classification.\n",
    "- Evaluate the model's performance.\n",
    "- Reflect on the model's strengths and weaknesses.\n",
    "\n",
    "Example code snippets are provided below to help you complete the tasks. You can use them as a starting point for your solution. Alternatively, you are welcome to use other libraries and resources if you prefer.\n",
    "\n",
    "### Exercise 2: Fine-Tuning a BERT Model for Text Classification with Your Own Dataset\n",
    "In this exercise, you will use your own classified State of the Union dataset to fine-tune a BERT model for text classification. You can use the `bert-base-cased` model or select a different pre-trained model if you prefer. The goal is to fine-tune the model to predict the sentiment you assigned to paragraphs (or other subsets of the speeches in the first assignment) based on the text of the speech.\n",
    "\n",
    "Although we do not provide specific code for this exercise, you should be able to reuse much of the code from Exercise 1 with some modifications. Here are some suggestions to guide your approach:\n",
    "- **Dataset Creation**: Create a dataset with text samples and corresponding sentiment labels. This can be a binary classification task (e.g., positive vs. negative sentiment) or a multi-class classification task (e.g., positive vs. neutral vs. negative sentiment). You will need to choose appropriate cutoffs for the sentiment scores to create the labels.\n",
    "- **Fine-Tuning**: Fine-tune a pre-trained language model for text classification. You can adapt the code from the previous exercise to work with your dataset and classification task.\n",
    "- **Evaluation**: Evaluate the model's performance using the same evaluation metrics from Exercise 1 or other suitable metrics. How well does this approach work for your dataset compared to the simple approach you used in the first assignment? If you did not use a simple approach in the first assignment, you should do one now to compare to the results here. **Make sure you are using the same data for the different approaches.** Use e.g., a dictionary based approach, a boolean approach, or a simple machine learning model like the Vader sentiment analysis model. \n",
    "- **Analysis**: Discuss the strengths and weaknesses of your model. How well does it perform on the sentiment classification task? What factors influence its performance? Do you see any benefits or limitations of using a pre-trained language model for this task over a simpler approach?\n",
    "\n",
    "\n",
    "**Note 1**: *This exercise, especially Exercise 1, can be computationally demanding. During the lecture \"Practical Day 2,\" we will discuss how to potentially run this on a server to make it feasible to include more training data.*\n",
    "\n",
    "**Note 2**: *The code provided below is based on the Hugging Face documentation. You can find more details here: [Hugging Face Transformers Training Documentation](https://huggingface.co/docs/transformers/en/training).*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'label': [4, 1, 3],\n",
       " 'text': [\"dr. goldberg offers everything i look for in a general practitioner.  he's nice and easy to talk to without being patronizing; he's always on time in seeing his patients; he's affiliated with a top-notch hospital (nyu) which my parents have explained to me is very important in case something happens and you need surgery; and you can get referrals to see specialists without having to see him first.  really, what more do you need?  i'm sitting here trying to think of any complaints i have about him, but i'm really drawing a blank.\",\n",
       "  \"Unfortunately, the frustration of being Dr. Goldberg's patient is a repeat of the experience I've had with so many other doctors in NYC -- good doctor, terrible staff.  It seems that his staff simply never answers the phone.  It usually takes 2 hours of repeated calling to get an answer.  Who has time for that or wants to deal with it?  I have run into this problem with many other doctors and I just don't get it.  You have office workers, you have patients with medical needs, why isn't anyone answering the phone?  It's incomprehensible and not work the aggravation.  It's with regret that I feel that I have to give Dr. Goldberg 2 stars.\",\n",
       "  \"Been going to Dr. Goldberg for over 10 years. I think I was one of his 1st patients when he started at MHMG. He's been great over the years and is really all about the big picture. It is because of him, not my now former gyn Dr. Markoff, that I found out I have fibroids. He explores all options with you and is very patient and understanding. He doesn't judge and asks all the right questions. Very thorough and wants to be kept in the loop on every aspect of your medical health and your life.\"]}"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Import the load_dataset function from the datasets library to load and work with text datasets (you may need to install the library first)\n",
    "from datasets import load_dataset\n",
    "\n",
    "# Load the 'yelp_review_full' dataset, which contains Yelp business reviews with five-star ratings\n",
    "dataset = load_dataset(\"yelp_review_full\")\n",
    "\n",
    "# Access and display the 100th sample from the 'train' split of the dataset\n",
    "dataset[\"train\"][0:3]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from datasets import Dataset\n",
    "\n",
    "#dataset = pd.read_csv('data.csv')\n",
    "\n",
    "#dataset = pd.DataFrame([['test', 0],\n",
    "#                        ['test2', 1]], columns=['text', 'label'])\n",
    "\n",
    "# Convert the pandas DataFrame to a Hugging Face Dataset\n",
    "#dataset_hf = Dataset.from_pandas(dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#dataset_hf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "51d8bad445c549f4a3d804afb571dffb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/650000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ee53343415f3432eb6edb5d4890f1aa3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/50000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Import the AutoTokenizer class from the transformers library for loading a tokenizer (you need to install the tokenizer library first)\n",
    "from transformers import AutoTokenizer\n",
    "\n",
    "# Load a pre-trained tokenizer from the 'google-bert/bert-base-cased' model\n",
    "# This tokenizer will be used to tokenize the text data for input into the model\n",
    "# We will use this model for fine-tuning on the Yelp review dataset, if you have a different model in mind, you can replace it here\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"google-bert/bert-base-cased\")\n",
    "\n",
    "# Define a function to tokenize input examples from the dataset\n",
    "def tokenize_function(examples):\n",
    "    # Tokenize the 'text' field of the examples with padding and truncation\n",
    "    # 'padding=\"max_length\"' pads the input sequences to the model's maximum length\n",
    "    # 'truncation=True' ensures that input sequences longer than the maximum length are truncated\n",
    "    return tokenizer(examples[\"text\"], padding=\"max_length\", truncation=True)\n",
    "\n",
    "# Apply the tokenize_function to the entire dataset using the map() method\n",
    "# 'batched=True' processes examples in batches for faster computation\n",
    "tokenized_datasets = dataset.map(tokenize_function)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train with PyTorch "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import PyTorch (needs to be installed)\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Remove the text column because the model does not accept raw text as an input\n",
    "tokenized_datasets = tokenized_datasets.remove_columns([\"text\"])\n",
    "\n",
    "# Rename the label column to labels because the model expects the argument to be named labels\n",
    "tokenized_datasets = tokenized_datasets.rename_column(\"label\", \"labels\")\n",
    "\n",
    "# Set the format of the dataset to return PyTorch tensors instead of lists\n",
    "tokenized_datasets.set_format(\"torch\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenized_datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a smaller subset of the dataset to speed up the fine-tuning \n",
    "# You can increase the size of the subset or use the full dataset for better results but it will take longer\n",
    "# Here, we are using 1000 examples each for training and evaluation\n",
    "small_train_dataset = tokenized_datasets[\"train\"].shuffle(seed=42).select(range(1000))\n",
    "small_eval_dataset = tokenized_datasets[\"test\"].shuffle(seed=42).select(range(1000))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import the DataLoader class from PyTorch's torch.utils.data module\n",
    "# DataLoader allows efficient batching, shuffling, and loading of data for model training and evaluation\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "# Create a DataLoader for the training dataset\n",
    "# small_train_dataset is assumed to be a subset of your tokenized training data\n",
    "# shuffle=True shuffles the data at each epoch to improve model generalization\n",
    "# batch_size=8 specifies that each batch will contain 8 samples\n",
    "train_dataloader = DataLoader(small_train_dataset, shuffle=True, batch_size=8)\n",
    "\n",
    "# Create a DataLoader for the evaluation dataset\n",
    "# small_eval_dataset is assumed to be a subset of your tokenized evaluation data\n",
    "# shuffle is not used for evaluation data to maintain consistency in data order\n",
    "# batch_size=8 specifies that each batch will contain 8 samples\n",
    "eval_dataloader = DataLoader(small_eval_dataset, batch_size=8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import the AutoModelForSequenceClassification class from the transformers library\n",
    "# This class allows loading a pre-trained transformer model specifically for classification tasks\n",
    "from transformers import AutoModelForSequenceClassification\n",
    "\n",
    "# Load a pre-trained BERT model ('google-bert/bert-base-cased') for sequence classification\n",
    "# The num_labels parameter specifies the number of output labels/classes for classification\n",
    "# Here, num_labels=5 indicates that the model is being fine-tuned for a task with five classes (e.g., a five-class classification problem)\n",
    "model = AutoModelForSequenceClassification.from_pretrained(\"google-bert/bert-base-cased\", num_labels=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import the AdamW optimizer from PyTorch's torch.optim module\n",
    "# AdamW is a variant of the Adam optimizer with weight decay, commonly used for fine-tuning transformers\n",
    "from torch.optim import AdamW\n",
    "\n",
    "# Create an optimizer for fine-tuning the model\n",
    "# optimizer updates the model parameters during training to minimize the loss function\n",
    "# model.parameters() specifies which parameters (weights and biases) to optimize\n",
    "# lr=5e-5 sets the learning rate for the optimizer, controlling the step size at each update\n",
    "optimizer = AdamW(model.parameters(), lr=5e-5)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import the get_scheduler function from the transformers library\n",
    "# This function allows the creation of different learning rate schedulers\n",
    "from transformers import get_scheduler\n",
    "\n",
    "# Define the number of epochs (full passes through the training data)\n",
    "num_epochs = 3\n",
    "\n",
    "# Calculate the total number of training steps (batches)\n",
    "# This is the product of the number of epochs and the number of batches in the training DataLoader\n",
    "num_training_steps = num_epochs * len(train_dataloader)\n",
    "\n",
    "# Create a learning rate scheduler\n",
    "# This scheduler linearly decreases the learning rate from its initial value (set in the optimizer) to zero over the course of training\n",
    "# name=\"linear\" specifies a linear learning rate decay schedule\n",
    "# optimizer specifies the optimizer to update (AdamW in this case)\n",
    "# num_warmup_steps=0 means there is no warmup period where the learning rate increases from a small value to the initial value\n",
    "# num_training_steps specifies the total number of steps for the decay schedule\n",
    "lr_scheduler = get_scheduler(\n",
    "    name=\"linear\", optimizer=optimizer, num_warmup_steps=0, num_training_steps=num_training_steps\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check if a GPU (CUDA) is available and set the device accordingly\n",
    "# torch.device(\"cuda\") selects a GPU if available for faster computation\n",
    "# If a GPU is not available, it defaults to using the CPU for computation\n",
    "device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n",
    "\n",
    "# Move the model's parameters to the specified device (GPU or CPU)\n",
    "# This ensures that all model computations happen on the chosen device\n",
    "model.to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Note: on my machine, the following code block took about 30 minutes to run\n",
    "\n",
    "# Import tqdm for creating a progress bar to visually track training progress\n",
    "# tqdm.auto automatically chooses the appropriate version (e.g., Jupyter-friendly) for display\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "# Create a progress bar for the total number of training steps\n",
    "# This will visually track the training progress over all epochs and batches\n",
    "progress_bar = tqdm(range(num_training_steps))\n",
    "\n",
    "# Set the model to training mode\n",
    "# This activates behaviors like dropout and batch normalization in training mode\n",
    "model.train()\n",
    "\n",
    "# Loop over the number of epochs (full passes through the dataset)\n",
    "for epoch in range(num_epochs):\n",
    "    # Loop over each batch in the training DataLoader\n",
    "    for batch in train_dataloader:\n",
    "        # Move each batch of data to the specified device (CPU or GPU)\n",
    "        batch = {k: v.to(device) for k, v in batch.items()}\n",
    "\n",
    "        # Forward pass: Pass the input data through the model to obtain outputs\n",
    "        outputs = model(**batch)\n",
    "\n",
    "        # Compute the loss (automatically computed based on the specified loss function for the model)\n",
    "        loss = outputs.loss\n",
    "\n",
    "        # Backpropagation: Compute gradients for model parameters with respect to the loss\n",
    "        loss.backward()\n",
    "\n",
    "        # Update model parameters using the optimizer based on computed gradients\n",
    "        optimizer.step()\n",
    "\n",
    "        # Update the learning rate using the learning rate scheduler\n",
    "        lr_scheduler.step()\n",
    "\n",
    "        # Reset the gradients of the model parameters to zero\n",
    "        # This prevents the accumulation of gradients from multiple backward passes\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        # Update the progress bar to reflect the completion of another training step\n",
    "        progress_bar.update(1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import the evaluate library to facilitate the computation of metrics\n",
    "import evaluate\n",
    "\n",
    "# Load the accuracy metric for evaluation\n",
    "# This metric will be used to evaluate the model's performance on the evaluation dataset\n",
    "metric = evaluate.load(\"accuracy\")\n",
    "\n",
    "# Set the model to evaluation mode\n",
    "# This deactivates training-specific behaviors such as dropout for more consistent predictions\n",
    "model.eval()\n",
    "\n",
    "# Loop over each batch in the evaluation DataLoader\n",
    "for batch in eval_dataloader:\n",
    "    # Move the batch of data to the specified device (CPU or GPU) for computation\n",
    "    batch = {k: v.to(device) for k, v in batch.items()}\n",
    "\n",
    "    # Disable gradient calculation during evaluation to save memory and computation\n",
    "    # This ensures no gradients are computed, as they are not needed during evaluation\n",
    "    with torch.no_grad():\n",
    "        # Perform a forward pass through the model to obtain outputs (predictions)\n",
    "        outputs = model(**batch)\n",
    "\n",
    "    # Extract logits (raw output scores) from the model's output\n",
    "    logits = outputs.logits\n",
    "\n",
    "    # Convert the logits to predicted class labels\n",
    "    # torch.argmax selects the index of the maximum value along the specified dimension (dim=-1)\n",
    "    predictions = torch.argmax(logits, dim=-1)\n",
    "\n",
    "    # Add the batch of predictions and reference (true) labels to the metric for computation\n",
    "    metric.add_batch(predictions=predictions, references=batch[\"labels\"])\n",
    "\n",
    "# Compute and return the final accuracy metric based on all evaluated batches\n",
    "metric.compute()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Make a prediction using a new text example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to make predictions for input text\n",
    "def predict_class(input_text, model, tokenizer, device):\n",
    "    # Tokenize the input text\n",
    "    inputs = tokenizer(\n",
    "        input_text,\n",
    "        return_tensors=\"pt\",        # Return PyTorch tensors\n",
    "        padding=\"max_length\",       # Pad to the maximum length used during training\n",
    "        truncation=True,            # Truncate if the input exceeds the maximum length\n",
    "        max_length=128              # Ensure this matches the max length used during training\n",
    "    )\n",
    "    \n",
    "    # Move input tensors to the same device as the model\n",
    "    inputs = {k: v.to(device) for k, v in inputs.items()}\n",
    "    \n",
    "    # Put the model in evaluation mode\n",
    "    model.eval()\n",
    "    \n",
    "    # Make predictions without computing gradients\n",
    "    with torch.no_grad():\n",
    "        outputs = model(**inputs)\n",
    "    \n",
    "    # Extract logits (raw scores) from the model's output\n",
    "    logits = outputs.logits\n",
    "    \n",
    "    # Convert logits to predicted class (taking the index of the max logit)\n",
    "    predicted_class = torch.argmax(logits, dim=-1).item()\n",
    "    \n",
    "    return predicted_class\n",
    "\n",
    "# Example usage\n",
    "input_text = \"The food at this restaurant was absolutely amazing!\"\n",
    "predicted_class = predict_class(input_text, model, tokenizer, device)\n",
    "\n",
    "# Display the predicted class, 0 to 4, based on the input text 4 being the highest rating\n",
    "print(f\"The predicted class for the input text is: {predicted_class}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch",
   "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.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}