generate_dataset.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/abhishekanand/miniforge3/envs/focus-llm/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "\n",
    "from openai_service import get_chat_completion\n",
    "from prompts import *\n",
    "import pandas as pd\n",
    "import json\n",
    "import os\n",
    "from tqdm.auto import tqdm\n",
    "import random\n",
    "import string\n",
    "tqdm.pandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_row(row, answer_key, processed_data):\n",
    "    id = row['id']\n",
    "    if id not in processed_data:\n",
    "        prompt = row['prompt']\n",
    "        text_response, prompt_tokens, completion_tokens, total_tokens = get_chat_completion(prompt=prompt, model=\"gpt-4-0125-preview\")\n",
    "        processed_data[id] = {\n",
    "            'response': text_response,\n",
    "            'prompt_tokens': prompt_tokens,\n",
    "            'completion_tokens': completion_tokens,\n",
    "            'total_tokens': total_tokens,\n",
    "            'correct_answer': row[answer_key]\n",
    "        }\n",
    "        return True\n",
    "    return False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ARC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>id</th>\n",
       "      <th>question</th>\n",
       "      <th>choices</th>\n",
       "      <th>answerKey</th>\n",
       "      <th>options</th>\n",
       "      <th>prompt</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Mercury_SC_415702</td>\n",
       "      <td>George wants to warm his hands quickly by rubb...</td>\n",
       "      <td>{'text': array(['dry palms', 'wet palms', 'pal...</td>\n",
       "      <td>A</td>\n",
       "      <td>['dry palms', 'wet palms', 'palms covered with...</td>\n",
       "      <td>Question: - George wants to warm his hands qui...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>MCAS_2009_5_6516</td>\n",
       "      <td>Which of the following statements best explain...</td>\n",
       "      <td>{'text': array(['The refrigerator door is smoo...</td>\n",
       "      <td>B</td>\n",
       "      <td>['The refrigerator door is smooth.', 'The refr...</td>\n",
       "      <td>Question: - Which of the following statements ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Mercury_7233695</td>\n",
       "      <td>A fold observed in layers of sedimentary rock ...</td>\n",
       "      <td>{'text': array(['cooling of flowing magma.', '...</td>\n",
       "      <td>B</td>\n",
       "      <td>['cooling of flowing magma.', 'converging of c...</td>\n",
       "      <td>Question: - A fold observed in layers of sedim...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Mercury_7041615</td>\n",
       "      <td>Which of these do scientists offer as the most...</td>\n",
       "      <td>{'text': array(['worldwide disease', 'global m...</td>\n",
       "      <td>D</td>\n",
       "      <td>['worldwide disease', 'global mountain buildin...</td>\n",
       "      <td>Question: - Which of these do scientists offer...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Mercury_7041860</td>\n",
       "      <td>A boat is acted on by a river current flowing ...</td>\n",
       "      <td>{'text': array(['west', 'east', 'north', 'sout...</td>\n",
       "      <td>B</td>\n",
       "      <td>['west', 'east', 'north', 'south']</td>\n",
       "      <td>Question: - A boat is acted on by a river curr...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                  id                                           question  \\\n",
       "0  Mercury_SC_415702  George wants to warm his hands quickly by rubb...   \n",
       "1   MCAS_2009_5_6516  Which of the following statements best explain...   \n",
       "2    Mercury_7233695  A fold observed in layers of sedimentary rock ...   \n",
       "3    Mercury_7041615  Which of these do scientists offer as the most...   \n",
       "4    Mercury_7041860  A boat is acted on by a river current flowing ...   \n",
       "\n",
       "                                             choices answerKey  \\\n",
       "0  {'text': array(['dry palms', 'wet palms', 'pal...         A   \n",
       "1  {'text': array(['The refrigerator door is smoo...         B   \n",
       "2  {'text': array(['cooling of flowing magma.', '...         B   \n",
       "3  {'text': array(['worldwide disease', 'global m...         D   \n",
       "4  {'text': array(['west', 'east', 'north', 'sout...         B   \n",
       "\n",
       "                                             options  \\\n",
       "0  ['dry palms', 'wet palms', 'palms covered with...   \n",
       "1  ['The refrigerator door is smooth.', 'The refr...   \n",
       "2  ['cooling of flowing magma.', 'converging of c...   \n",
       "3  ['worldwide disease', 'global mountain buildin...   \n",
       "4                 ['west', 'east', 'north', 'south']   \n",
       "\n",
       "                                              prompt  \n",
       "0  Question: - George wants to warm his hands qui...  \n",
       "1  Question: - Which of the following statements ...  \n",
       "2  Question: - A fold observed in layers of sedim...  \n",
       "3  Question: - Which of these do scientists offer...  \n",
       "4  Question: - A boat is acted on by a river curr...  "
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "arc_df = pd.read_csv('datasets/arc_challenge_train_with_prompt.csv')\n",
    "arc_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1119, 1119)"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(arc_df), len(arc_df['id'].unique())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 129,
   "metadata": {},
   "outputs": [],
   "source": [
    "PROCESSED_DATA_PATH = \"./datasets/arc_challenge_train_gpt4_response.json\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 130,
   "metadata": {},
   "outputs": [],
   "source": [
    "processed_data = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 131,
   "metadata": {},
   "outputs": [],
   "source": [
    "if os.path.exists(PROCESSED_DATA_PATH):\n",
    "    with open(PROCESSED_DATA_PATH, 'r') as file:\n",
    "        processed_data = json.load(file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 132,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "500"
      ]
     },
     "execution_count": 132,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(processed_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 134,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/500 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [2:48:47<00:00, 20.26s/it]  \n"
     ]
    }
   ],
   "source": [
    "processed_status = arc_df.iloc[500:1000].progress_apply(lambda row: process_row(row, answer_key='answerKey', processed_data=processed_data), axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 135,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1000"
      ]
     },
     "execution_count": 135,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(processed_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 136,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(PROCESSED_DATA_PATH, 'w') as file:\n",
    "\tjson.dump(processed_data, file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 137,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "500"
      ]
     },
     "execution_count": 137,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(processed_status[processed_status])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Winogrande"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>question</th>\n",
       "      <th>option1</th>\n",
       "      <th>option2</th>\n",
       "      <th>answer</th>\n",
       "      <th>prompt</th>\n",
       "      <th>id</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>John moved the couch from the garage to the ba...</td>\n",
       "      <td>garage</td>\n",
       "      <td>backyard</td>\n",
       "      <td>1</td>\n",
       "      <td>Question: - John moved the couch from the gara...</td>\n",
       "      <td>GO3wkLp9TC</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>The doctor diagnosed Justin with bipolar and R...</td>\n",
       "      <td>Justin</td>\n",
       "      <td>Robert</td>\n",
       "      <td>2</td>\n",
       "      <td>Question: - The doctor diagnosed Justin with b...</td>\n",
       "      <td>IPQZOkoLLz</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Dennis drew up a business proposal to present ...</td>\n",
       "      <td>Dennis</td>\n",
       "      <td>Logan</td>\n",
       "      <td>1</td>\n",
       "      <td>Question: - Dennis drew up a business proposal...</td>\n",
       "      <td>W85KDUn8gA</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Felicia unexpectedly made fried eggs for break...</td>\n",
       "      <td>Felicia</td>\n",
       "      <td>Katrina</td>\n",
       "      <td>2</td>\n",
       "      <td>Question: - Felicia unexpectedly made fried eg...</td>\n",
       "      <td>UIUCGI68OY</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>My shampoo did not lather easily on my Afro ha...</td>\n",
       "      <td>shampoo</td>\n",
       "      <td>hair</td>\n",
       "      <td>2</td>\n",
       "      <td>Question: - My shampoo did not lather easily o...</td>\n",
       "      <td>WFAUCgvz91</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                            question  option1   option2  \\\n",
       "0  John moved the couch from the garage to the ba...   garage  backyard   \n",
       "1  The doctor diagnosed Justin with bipolar and R...   Justin    Robert   \n",
       "2  Dennis drew up a business proposal to present ...   Dennis     Logan   \n",
       "3  Felicia unexpectedly made fried eggs for break...  Felicia   Katrina   \n",
       "4  My shampoo did not lather easily on my Afro ha...  shampoo      hair   \n",
       "\n",
       "   answer                                             prompt          id  \n",
       "0       1  Question: - John moved the couch from the gara...  GO3wkLp9TC  \n",
       "1       2  Question: - The doctor diagnosed Justin with b...  IPQZOkoLLz  \n",
       "2       1  Question: - Dennis drew up a business proposal...  W85KDUn8gA  \n",
       "3       2  Question: - Felicia unexpectedly made fried eg...  UIUCGI68OY  \n",
       "4       2  Question: - My shampoo did not lather easily o...  WFAUCgvz91  "
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wino_df = pd.read_csv('datasets/winogrande_train_with_prompt.csv')\n",
    "wino_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>question</th>\n",
       "      <th>option1</th>\n",
       "      <th>option2</th>\n",
       "      <th>answer</th>\n",
       "      <th>prompt</th>\n",
       "      <th>id</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Laura gave a class on drawing to a group inclu...</td>\n",
       "      <td>Laura</td>\n",
       "      <td>Kayla</td>\n",
       "      <td>2</td>\n",
       "      <td>Question: - Laura gave a class on drawing to a...</td>\n",
       "      <td>NcYXNaT5zN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>John cleaned the table with his bare hand and ...</td>\n",
       "      <td>table</td>\n",
       "      <td>water</td>\n",
       "      <td>2</td>\n",
       "      <td>Question: - John cleaned the table with his ba...</td>\n",
       "      <td>HkUfjZSeAx</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>So _ was embarrassed because Mary wears a wig ...</td>\n",
       "      <td>Mary</td>\n",
       "      <td>Kayla</td>\n",
       "      <td>1</td>\n",
       "      <td>Question: - So _ was embarrassed because Mary ...</td>\n",
       "      <td>z79MTrfX9f</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>I like working on the old operating systems ov...</td>\n",
       "      <td>new operating systems</td>\n",
       "      <td>old operating systems</td>\n",
       "      <td>2</td>\n",
       "      <td>Question: - I like working on the old operatin...</td>\n",
       "      <td>33wGXmn6z5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>When visiting Canadian National Parks, I purch...</td>\n",
       "      <td>Annual Park Pass</td>\n",
       "      <td>Daily Park Pass</td>\n",
       "      <td>1</td>\n",
       "      <td>Question: - When visiting Canadian National Pa...</td>\n",
       "      <td>OLcn9DE9IC</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                            question                option1  \\\n",
       "0  Laura gave a class on drawing to a group inclu...                  Laura   \n",
       "1  John cleaned the table with his bare hand and ...                  table   \n",
       "2  So _ was embarrassed because Mary wears a wig ...                   Mary   \n",
       "3  I like working on the old operating systems ov...  new operating systems   \n",
       "4  When visiting Canadian National Parks, I purch...       Annual Park Pass   \n",
       "\n",
       "                 option2  answer  \\\n",
       "0                  Kayla       2   \n",
       "1                  water       2   \n",
       "2                  Kayla       1   \n",
       "3  old operating systems       2   \n",
       "4        Daily Park Pass       1   \n",
       "\n",
       "                                              prompt          id  \n",
       "0  Question: - Laura gave a class on drawing to a...  NcYXNaT5zN  \n",
       "1  Question: - John cleaned the table with his ba...  HkUfjZSeAx  \n",
       "2  Question: - So _ was embarrassed because Mary ...  z79MTrfX9f  \n",
       "3  Question: - I like working on the old operatin...  33wGXmn6z5  \n",
       "4  Question: - When visiting Canadian National Pa...  OLcn9DE9IC  "
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wino_df_sampled = wino_df.sample(n=990).reset_index(drop=True)\n",
    "wino_df_sampled.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# wino_df['id'] = [''.join(random.choices(string.ascii_letters + string.digits, k=10)) for _ in range(len(wino_df))]\n",
    "# wino_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1000"
      ]
     },
     "execution_count": 57,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "WINO_PROCESSED_DATA_PATH = \"./datasets/winogrande_train_gpt4_response.json\"\n",
    "\n",
    "processed_data = {}\n",
    "if os.path.exists(WINO_PROCESSED_DATA_PATH):\n",
    "    with open(WINO_PROCESSED_DATA_PATH, 'r') as file:\n",
    "        processed_data = json.load(file)\n",
    "len(processed_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:17<00:00, 17.96s/it]\n"
     ]
    }
   ],
   "source": [
    "processed_status = wino_df[wino_df['id'] == 'G1iDf6VN4H'].progress_apply(lambda row: process_row(row, answer_key='answer', processed_data=processed_data), axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(WINO_PROCESSED_DATA_PATH, 'w') as file:\n",
    "\tjson.dump(processed_data, file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1"
      ]
     },
     "execution_count": 55,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(processed_status[processed_status])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "focus-llm",
   "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}