{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import pickle\n",
    "import os; os.chdir('..')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_df = pd.read_csv('wiki_intro_processed.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_prompt(title, starter_text):\n",
    "    return f'''200 word wikipedia style introduction on '{title}'\n",
    "    {starter_text}'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = []\n",
    "\n",
    "for i in range(150):\n",
    "    with open(f'data/result-{i}.pkl', 'rb') as file:\n",
    "        temp = pickle.load(file)\n",
    "        results += list(temp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "150000"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "processed_results = []\n",
    "for dct in results:\n",
    "    for key in dct:\n",
    "        processed_results.append({\n",
    "            'title': key, \n",
    "            'generated_text': dct[key]['choices'][0]['text'],\n",
    "            'prompt_tokens' : dct[key]['usage']['prompt_tokens'],\n",
    "            'completion_tokens' : dct[key]['usage']['completion_tokens'],\n",
    "        })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "processed_results_df = pd.DataFrame(processed_results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create final df\n",
    "final_df = pd.merge(data_df, processed_results_df, on=['title'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create new columns\n",
    "final_df['prompt'] = final_df.apply(lambda row: create_prompt(row['title'], row['starter_text']), axis=1)\n",
    "\n",
    "final_df['generated_text_complete'] = final_df['starter_text'] + final_df['generated_text']\n",
    "\n",
    "final_df['generated_text_len'] = final_df['generated_text_complete'].apply(lambda x: len(x.split(' ')))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Rename columns\n",
    "final_df = final_df.rename(columns={\n",
    "    'intro': 'wiki_intro', 'intro_len': 'wiki_intro_len',\n",
    "    'generated_text_complete': 'generated_intro', 'generated_text_len' : 'generated_intro_len'})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reorder columns\n",
    "final_df = final_df[[ \n",
    "    'id', 'url', 'title', 'wiki_intro', 'generated_intro', 'title_len',\n",
    "    'wiki_intro_len', 'generated_intro_len', 'prompt', 'generated_text',\n",
    "    'prompt_tokens', 'generated_text_tokens']]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Write csv file\n",
    "final_df.to_csv('GPT-wiki-intro.csv', index=False)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "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.10.12"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "3f100d68d9cf80676b1a4c3ace5430b03ae266a1d88e3f101eb196b64b263632"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}