day_4/day_4b.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Day 4b\n",
    "In this notebook we will evaluate the ability of a large language model (LLM) to model demographic differences. We will use a dataset on [vaccine hesitancy](https://www.kaggle.com/datasets/christianhritter/vaccine-hesitancy-canada-cosmo-survey?select=COSMO_in_Canada_Waves_1-8_FINAL.csv) to evaluate this.\n",
    "\n",
    "By the end of this notebook, you will have learned to:\n",
    "- Compare the distribution of LLM-generated responses to a survey question with the actual data\n",
    "- Evaluate the effect of demographic information on the generated responses\n",
    " "
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "8a846041309c0b3d"
  },
  {
   "cell_type": "markdown",
   "source": [
    "\n",
    "## Environment Setup\n",
    "Make sure to set your runtime to use a GPU by going to `Runtime` -> `Change runtime type` -> `Hardware accelerator` -> `T4 GPU`"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "5e87e9c46775bdc7"
  },
  {
   "cell_type": "code",
   "source": [
    "import sys\n",
    "if 'google.colab' in sys.modules:  # If in Google Colab environment\n",
    "    # Mount google drive to enable access to data files\n",
    "    from google.colab import drive\n",
    "    drive.mount('/content/drive')\n",
    "    \n",
    "    # Installing requisite packages\n",
    "    !pip install transformers accelerate &> /dev/null\n",
    "    \n",
    "    # Change working directory to day_4\n",
    "    %cd /content/drive/MyDrive/LLM4BeSci_Ljubljana2025/day_4"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "394aa56d7e9b7c31",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "source": [
    "import pandas as pd\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline\n",
    "import torch\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "from tqdm.notebook import tqdm_notebook as tqdm"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "e2a1bee0497b7c26",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Loading the data\n",
    "The code begins by loading the dataset as a `pandas.DataFrame`. The dataset contains the following columns:\n",
    "1. `'age'`\n",
    "2. `'gender'`\n",
    "3. `'education'`\n",
    "4. `'take_vaccine'`: Survey question on willingness to take a COVID-19 vaccine (1 (Strongly Disagree) to 7 (Strongly Agree))\n",
    "5. `'mandatory_vaccine'`\n",
    "\n",
    "It also contains three '`persona`' columns, which contain the demographic information in a prompt format:\n",
    "\n",
    "1. `'persona_a'`: Gender, Age, Education\n",
    "2. `'persona_b'`: Gender, Age\n",
    "3. `'persona_c'`: Age, Education"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "1c7eb38d7b1666f7"
  },
  {
   "cell_type": "code",
   "source": [
    "vaccine = pd.read_csv('vaccine_hesitancy.csv')\n",
    "vaccine"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "3bae385c289f8ff6",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "You can use the `value_counts()` method to get an idea of the distribution of the data. For instance, we can look at the distribution of the `'take_vaccine'` column:"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "b9d5a19b2e72844c"
  },
  {
   "cell_type": "code",
   "source": "vaccine['take_vaccine'].value_counts()",
   "metadata": {
    "collapsed": false
   },
   "id": "a3f4cc9301aa4c31",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "**TASK 1**: Replace `'age`' above with other demographic column names (`'gender'`, `'education'`) to get an idea of the distribution of the data.\n",
    "\n",
    "**TASK 2**: Replace `'take_vaccine'` above with other survey questions to get an idea of the distribution of the data.\n",
    "\n",
    "**TASK 3 (BONUS)**: Try to plot the distributions of the data using a histogram. You can use the `sns.histplot()` function from the `seaborn` library."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "71f9095df2983fe5"
  },
  {
   "cell_type": "markdown",
   "source": [
    "In the next section, we will generate responses to the survey question using the LLM. The survey question is as follows:"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "23c81e298cb4e644"
  },
  {
   "cell_type": "code",
   "source": [
    "survey_q = \"\"\"\n",
    "    Please give your opinion on the following statement:\n",
    "    \n",
    "    'If an effective COVID-19 vaccine becomes available and is recommended for me, I would get it.'\n",
    "    \n",
    "    Choice:\n",
    "    \n",
    "    1 = Strongly Disagree\n",
    "    2\n",
    "    3\n",
    "    4\n",
    "    5\n",
    "    6\n",
    "    7 = Strongly Agree \n",
    "    \n",
    "    Strictly only respond with the number corresponding to your choice and nothing else.\n",
    "\"\"\""
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "5f82d7fd82d7556a",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Comparing the model's default distribution with the data\n",
    "This section generates responses to the survey question using the LLM and compares the distribution of the generated responses with the actual data. We again use Llama 3.2 (don't forget to copy-paste your access token). We begin by loading the model and tokenizer:"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "391a9e1b4d8677b"
  },
  {
   "cell_type": "code",
   "source": [
    "torch.random.manual_seed(42) # Set seed for reproducibility\n",
    "\n",
    "access_token = '<your access token here>'\n",
    "\n",
    "# Load the model\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    \"meta-llama/Llama-3.2-3B-Instruct\",\n",
    "    device_map=\"cuda\", # Use GPU\n",
    "    torch_dtype=torch.float16, # Use half-precision\n",
    "    trust_remote_code=True,\n",
    "    attn_implementation='eager', # For faster inference on T4 GPUs\n",
    "    use_auth_token=access_token\n",
    ")\n",
    "\n",
    "# Load the tokenizer\n",
    "tokenizer = AutoTokenizer.from_pretrained(\n",
    "    \"meta-llama/Llama-3.2-3B-Instruct\",\n",
    "    use_auth_token=access_token,\n",
    ")"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "fbc5cb91d927eddc",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "The code then uses the `pipeline` class to generate responses to the survey question. The `generation_args` dictionary contains the following arguments:"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "80d3133861a77ecc"
  },
  {
   "cell_type": "code",
   "source": [
    "pipe = pipeline(\n",
    "    \"text-generation\",\n",
    "    model=model,\n",
    "    tokenizer=tokenizer,\n",
    ")\n",
    "\n",
    "generation_args = {\n",
    "    \"max_new_tokens\": 10,  # Maximum number of tokens to generate\n",
    "    \"return_full_text\": False, # Return only the generated text\n",
    "    \"do_sample\": True, # Allow sampling from the softmax distribution\n",
    "    \"temperature\": 1.0,  # Temperature parameter\n",
    "    \"pad_token_id\": pipe.tokenizer.eos_token_id\n",
    "}"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "e80059b11d5d40cb",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "It firstly generates 100 responses to the survey question. The code then detects which integer (`1`-`7`) the model responded with and appends this to the `take_vaccine_preds` list. "
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "6ce9e99bdaf9e5a"
  },
  {
   "cell_type": "code",
   "source": [
    "# Generate responses\n",
    "take_vaccine_preds = [] # List to store generated responses\n",
    "n_samples = 100 # Number of samples to generate\n",
    "for i in tqdm(range(n_samples)):\n",
    "    \n",
    "    # Define prompt with JSON structure\n",
    "    prompt = [{\"role\": \"user\", \"content\": survey_q}]  \n",
    "    \n",
    "    # Generate response\n",
    "    response = pipe(prompt, **generation_args)[0]['generated_text']\n",
    "    \n",
    "    # Checks which integer corresponds to the response and appends to list\n",
    "    possibles = ['1', '2', '3', '4', '5', '6', '7']\n",
    "    pred = [int(x) for x in possibles if x in response]\n",
    "    if len(pred) == 1:\n",
    "        pred = pred[0]\n",
    "        take_vaccine_preds.append(pred)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "6315cfa51c6449a3",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "To compare the distribution of the generated responses (`take_vaccine_preds`) with the actual data (`vaccine['take_vaccine']`), the code uses a histogram:"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "9310e018e884f15"
  },
  {
   "cell_type": "code",
   "source": [
    "# Comparing distribution of generated responses with actual data\n",
    "figs, axs = plt.subplots(1, 2, figsize=(12, 6))\n",
    "\n",
    "# Plotting actual data\n",
    "sns.histplot(vaccine['take_vaccine'], stat='percent', bins=7, ax=axs[0])\n",
    "axs[0].set(title='Actual Data', xlabel='Take vaccine', xlim=(1, 7))\n",
    "\n",
    "# Plotting generated data\n",
    "sns.histplot(take_vaccine_preds, stat='percent', bins=7, ax=axs[1])\n",
    "axs[1].set(title='Generated Data', xlabel='Take vaccine', xlim=(1, 7))"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "d7cc8691e3529b6a",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "As you can see, with a temperature of `1.0`, the model's output distribution is considerably narrower than that of the survey participants. In other words, we need to increase the `temperature` parameter to allow for more realistic individual variability in the generated responses. \n",
    "\n",
    "**TASK 1**: Try playing around with larger values of `\"temperature\"` (e.g. `3.0`, `5.0`, `10.0`) to see how it affects the generated responses. **Don't forget to re-run the cell that defines `generation_args` after changing the temperature.**"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "10d34d37e16fde48"
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Demographic Steering: Gender, Age, Education\n",
    "As mentioned, our dataset also has 3 \"personas\", which contain different kinds of demographic information in a prompt format:\".\n",
    "\n",
    "1. `'persona_a'`: Gender, Age, Education\n",
    "2. `'persona_b'`: Gender, Age\n",
    "3. `'persona_c'`: Age, Education\n",
    "\n",
    "We will now evaluate each persona's effect on the generated responses. For each \"participant\", we will append the demographic information to `survey_q` and generate a response. For instance, the prompt for the first \"participant\" would look like this:"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "ad0307dc5fa0a47a"
  },
  {
   "cell_type": "code",
   "source": [
    "prompt = f\"You are {vaccine['persona_a'][0]}.\\n-------------------------------\\n{survey_q}\"\n",
    "print(prompt)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "d140d1fd835a12d",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": "The code next generates responses for each persona and compares the generated responses with the actual data. It starts with `'persona_a'` and a temperature of `1.0` (can you think of why we don't choose `3.0`, despite the fact that it generated a more realistic distribution in the previous task?):",
   "metadata": {
    "collapsed": false
   },
   "id": "474a9a09f8c91248"
  },
  {
   "cell_type": "code",
   "source": [
    "generation_args = {\n",
    "    \"max_new_tokens\": 10,  # Maximum number of tokens to generate\n",
    "    \"return_full_text\": False, # Return only the generated text\n",
    "    \"do_sample\": True, # Allow sampling from the softmax distribution\n",
    "    \"temperature\": 1.0,  # Temperature parameter\n",
    "    \"pad_token_id\": pipe.tokenizer.eos_token_id\n",
    "}\n",
    "\n",
    "demog_col = 'persona_a'  # Replace this for TASK 1\n",
    "\n",
    "# Generate responses\n",
    "take_vaccine_preds = [] # List to store generated responses\n",
    "for demog_prompt in tqdm(vaccine[demog_col]):\n",
    "    prompt = f\"You are {demog_prompt}.\\n-------------------------------\\n{survey_q}\"\n",
    "    \n",
    "    # Define prompt with JSON structure\n",
    "    prompt = [{\"role\": \"user\", \"content\": prompt}]  # Define prompt with JSON structure\n",
    "    \n",
    "    # Generate response\n",
    "    response = pipe(prompt, **generation_args)[0]['generated_text']\n",
    "    \n",
    "    # Checks which integer corresponds to the response\n",
    "    possibles = ['1', '2', '3', '4', '5', '6', '7']\n",
    "    pred = [int(x) for x in possibles if x in response]\n",
    "    if len(pred) == 1:\n",
    "        pred = pred[0]\n",
    "        take_vaccine_preds.append(pred)\n",
    "    else: # If no integer found, or too many found, return None\n",
    "        take_vaccine_preds.append(None)\n",
    "    \n",
    "# Append generated responses to dataframe\n",
    "vaccine[f'{demog_col}_preds'] = take_vaccine_preds\n",
    "vaccine[f'{demog_col}_preds'].value_counts()"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "4108d06f5623db49",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "The generated responses can be compared with the actual data using a regression plot:"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "81282ded685f68bf"
  },
  {
   "cell_type": "code",
   "source": [
    "# Comparing generated and actual data with a regression plot\n",
    "ax = sns.regplot(\n",
    "    x=f'{demog_col}_preds', y='take_vaccine', x_jitter=.1, y_jitter=.1, data=vaccine,\n",
    "    scatter_kws={'alpha': 0.5}\n",
    ")\n",
    "x_y_lim = (.5, 7.5)\n",
    "ax.set(xlim=x_y_lim, ylim=x_y_lim)\n",
    "print(f\"Correlation: {vaccine['take_vaccine'].corr(vaccine[f'{demog_col}_preds'])}\")"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "44813bcb063fe68a",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "**TASK 1**: Unfortunately,`'persona_a'` did not perform well. Try using other demographic columns (`'persona_b'`, `'persona_c'`) by replacing `'persona_a'` above to see how the generated responses change.\n",
    "\n",
    "As you have now discovered, it can be challenging to generate responses to survey questions that reflect individual-level differences, especially when the LLM only has access to group-level (demographic) information. Neveretheless, were we to evaluate performance at the group level, we would likely find that the model is able to capture some of the differences between groups."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "ab785c76b21ddec"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}