{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "outputs": [],
   "source": [
    "import os\n",
    "import pyarrow.parquet as pq\n",
    "import pandas as pd\n",
    "import tqdm\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "\n",
    "folder_path = r\"D:\\github\\oa_test_gpu\\wiki_queries\""
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "outputs": [],
   "source": [
    "# parquet_files = [os.path.join(folder_path, f) for f in os.listdir(folder_path) if f.endswith(\".parquet\")]\n",
    "# parquet_files"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "outputs": [],
   "source": [
    "# dfs = [pq.read_table(f).to_pandas() for f in parquet_files]\n",
    "# concatenated_df = pd.concat(dfs)\n",
    "# concatenated_df.to_parquet(os.path.join(folder_path, \"wiki_queries_950k.parquet\"))"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "outputs": [
    {
     "data": {
      "text/plain": "              id                                                url  \\\n1000000  5106811  https://en.wikipedia.org/wiki/Shahtaj%20Qizilbash   \n1000001  5106829   https://en.wikipedia.org/wiki/Aurat%20Foundation   \n1000002  5106839          https://en.wikipedia.org/wiki/BloomingOUT   \n1000003  5106855  https://en.wikipedia.org/wiki/The%20Moment%20o...   \n1000004  5106875      https://en.wikipedia.org/wiki/Daily%20Emerald   \n...          ...                                                ...   \n89995     156641          https://en.wikipedia.org/wiki/Stalag%2017   \n89996     156644  https://en.wikipedia.org/wiki/Die%20Another%20Day   \n89997     156653             https://en.wikipedia.org/wiki/District   \n89998     156656  https://en.wikipedia.org/wiki/Oath%20of%20Alle...   \n89999     156658           https://en.wikipedia.org/wiki/Wiktionary   \n\n                                                      text  \\\n1000000  Shahtaj Qizilbash  was a women's right advocat...   \n1000001  Aurat Foundation, founded in 1986, is a women'...   \n1000002  bloomingOUT is an LGBT+ radio show broadcast o...   \n1000003  The Moment of the Magician  is a fantasy novel...   \n1000004  The Daily Emerald is the independent, student-...   \n...                                                    ...   \n89995    Stalag 17 is a 1953 American drama war film wh...   \n89996    Die Another Day is a 2002 spy film and the twe...   \n89997    A district is a type of administrative divisio...   \n89998    The Irish Oath of Allegiance  was a controvers...   \n89999    Wiktionary is a multilingual, web-based projec...   \n\n                                title  word_num  \\\n1000000             Shahtaj Qizilbash       139   \n1000001              Aurat Foundation        86   \n1000002                   BloomingOUT       153   \n1000003    The Moment of the Magician        49   \n1000004                 Daily Emerald        52   \n...                               ...       ...   \n89995                       Stalag 17       157   \n89996                 Die Another Day       211   \n89997                        District        43   \n89998    Oath of Allegiance (Ireland)        90   \n89999                      Wiktionary       167   \n\n                                                   queries  \n1000000  [who was shahtaj qizilbash, who is shahtaj qiz...  \n1000001  [where is aurat?, aurat foundation, what is au...  \n1000002  [where is blooming out, blooming out indiana, ...  \n1000003  [moment of the magician book, who wrote moment...  \n1000004  [where is the emerald newspaper in oregon, wha...  \n...                                                    ...  \n89995    [when was stalag 17 movie filmed, who was the ...  \n89996    [who acted in die another day movie, what movi...  \n89997    [what is the difference between district and c...  \n89998    [who swore the irish oath of allegiance?, when...  \n89999    [where is wikikidnary, definition of wikipedia...  \n\n[950000 rows x 6 columns]",
      "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>url</th>\n      <th>text</th>\n      <th>title</th>\n      <th>word_num</th>\n      <th>queries</th>\n    </tr>\n  </thead>\n  <tbody>\n    <tr>\n      <th>1000000</th>\n      <td>5106811</td>\n      <td>https://en.wikipedia.org/wiki/Shahtaj%20Qizilbash</td>\n      <td>Shahtaj Qizilbash  was a women's right advocat...</td>\n      <td>Shahtaj Qizilbash</td>\n      <td>139</td>\n      <td>[who was shahtaj qizilbash, who is shahtaj qiz...</td>\n    </tr>\n    <tr>\n      <th>1000001</th>\n      <td>5106829</td>\n      <td>https://en.wikipedia.org/wiki/Aurat%20Foundation</td>\n      <td>Aurat Foundation, founded in 1986, is a women'...</td>\n      <td>Aurat Foundation</td>\n      <td>86</td>\n      <td>[where is aurat?, aurat foundation, what is au...</td>\n    </tr>\n    <tr>\n      <th>1000002</th>\n      <td>5106839</td>\n      <td>https://en.wikipedia.org/wiki/BloomingOUT</td>\n      <td>bloomingOUT is an LGBT+ radio show broadcast o...</td>\n      <td>BloomingOUT</td>\n      <td>153</td>\n      <td>[where is blooming out, blooming out indiana, ...</td>\n    </tr>\n    <tr>\n      <th>1000003</th>\n      <td>5106855</td>\n      <td>https://en.wikipedia.org/wiki/The%20Moment%20o...</td>\n      <td>The Moment of the Magician  is a fantasy novel...</td>\n      <td>The Moment of the Magician</td>\n      <td>49</td>\n      <td>[moment of the magician book, who wrote moment...</td>\n    </tr>\n    <tr>\n      <th>1000004</th>\n      <td>5106875</td>\n      <td>https://en.wikipedia.org/wiki/Daily%20Emerald</td>\n      <td>The Daily Emerald is the independent, student-...</td>\n      <td>Daily Emerald</td>\n      <td>52</td>\n      <td>[where is the emerald newspaper in oregon, wha...</td>\n    </tr>\n    <tr>\n      <th>...</th>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n    </tr>\n    <tr>\n      <th>89995</th>\n      <td>156641</td>\n      <td>https://en.wikipedia.org/wiki/Stalag%2017</td>\n      <td>Stalag 17 is a 1953 American drama war film wh...</td>\n      <td>Stalag 17</td>\n      <td>157</td>\n      <td>[when was stalag 17 movie filmed, who was the ...</td>\n    </tr>\n    <tr>\n      <th>89996</th>\n      <td>156644</td>\n      <td>https://en.wikipedia.org/wiki/Die%20Another%20Day</td>\n      <td>Die Another Day is a 2002 spy film and the twe...</td>\n      <td>Die Another Day</td>\n      <td>211</td>\n      <td>[who acted in die another day movie, what movi...</td>\n    </tr>\n    <tr>\n      <th>89997</th>\n      <td>156653</td>\n      <td>https://en.wikipedia.org/wiki/District</td>\n      <td>A district is a type of administrative divisio...</td>\n      <td>District</td>\n      <td>43</td>\n      <td>[what is the difference between district and c...</td>\n    </tr>\n    <tr>\n      <th>89998</th>\n      <td>156656</td>\n      <td>https://en.wikipedia.org/wiki/Oath%20of%20Alle...</td>\n      <td>The Irish Oath of Allegiance  was a controvers...</td>\n      <td>Oath of Allegiance (Ireland)</td>\n      <td>90</td>\n      <td>[who swore the irish oath of allegiance?, when...</td>\n    </tr>\n    <tr>\n      <th>89999</th>\n      <td>156658</td>\n      <td>https://en.wikipedia.org/wiki/Wiktionary</td>\n      <td>Wiktionary is a multilingual, web-based projec...</td>\n      <td>Wiktionary</td>\n      <td>167</td>\n      <td>[where is wikikidnary, definition of wikipedia...</td>\n    </tr>\n  </tbody>\n</table>\n<p>950000 rows × 6 columns</p>\n</div>"
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "concatenated_df = pd.read_parquet(os.path.join(folder_path, \"wiki_queries_950k.parquet\"))\n",
    "concatenated_df"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 950000/950000 [01:16<00:00, 12381.89it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": "              id                                                url  \\\n0        5106811  https://en.wikipedia.org/wiki/Shahtaj%20Qizilbash   \n1        5106811  https://en.wikipedia.org/wiki/Shahtaj%20Qizilbash   \n2        5106811  https://en.wikipedia.org/wiki/Shahtaj%20Qizilbash   \n3        5106829   https://en.wikipedia.org/wiki/Aurat%20Foundation   \n4        5106829   https://en.wikipedia.org/wiki/Aurat%20Foundation   \n...          ...                                                ...   \n2849995   156656  https://en.wikipedia.org/wiki/Oath%20of%20Alle...   \n2849996   156656  https://en.wikipedia.org/wiki/Oath%20of%20Alle...   \n2849997   156658           https://en.wikipedia.org/wiki/Wiktionary   \n2849998   156658           https://en.wikipedia.org/wiki/Wiktionary   \n2849999   156658           https://en.wikipedia.org/wiki/Wiktionary   \n\n                                                      text  \\\n0        Shahtaj Qizilbash  was a women's right advocat...   \n1        Shahtaj Qizilbash  was a women's right advocat...   \n2        Shahtaj Qizilbash  was a women's right advocat...   \n3        Aurat Foundation, founded in 1986, is a women'...   \n4        Aurat Foundation, founded in 1986, is a women'...   \n...                                                    ...   \n2849995  The Irish Oath of Allegiance  was a controvers...   \n2849996  The Irish Oath of Allegiance  was a controvers...   \n2849997  Wiktionary is a multilingual, web-based projec...   \n2849998  Wiktionary is a multilingual, web-based projec...   \n2849999  Wiktionary is a multilingual, web-based projec...   \n\n                                title  word_num  \\\n0                   Shahtaj Qizilbash       139   \n1                   Shahtaj Qizilbash       139   \n2                   Shahtaj Qizilbash       139   \n3                    Aurat Foundation        86   \n4                    Aurat Foundation        86   \n...                               ...       ...   \n2849995  Oath of Allegiance (Ireland)        90   \n2849996  Oath of Allegiance (Ireland)        90   \n2849997                    Wiktionary       167   \n2849998                    Wiktionary       167   \n2849999                    Wiktionary       167   \n\n                                                    query  \n0                               who was shahtaj qizilbash  \n1                               who is shahtaj qizilbash?  \n2                                who is shahtaj qizilbash  \n3                                         where is aurat?  \n4                                        aurat foundation  \n...                                                   ...  \n2849995  when did the irish oath of allegiance take place  \n2849996             what was the irish oath of allegiance  \n2849997                              where is wikikidnary  \n2849998                           definition of wikipedia  \n2849999                 what is the meaning of wiktionary  \n\n[2850000 rows x 6 columns]",
      "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>url</th>\n      <th>text</th>\n      <th>title</th>\n      <th>word_num</th>\n      <th>query</th>\n    </tr>\n  </thead>\n  <tbody>\n    <tr>\n      <th>0</th>\n      <td>5106811</td>\n      <td>https://en.wikipedia.org/wiki/Shahtaj%20Qizilbash</td>\n      <td>Shahtaj Qizilbash  was a women's right advocat...</td>\n      <td>Shahtaj Qizilbash</td>\n      <td>139</td>\n      <td>who was shahtaj qizilbash</td>\n    </tr>\n    <tr>\n      <th>1</th>\n      <td>5106811</td>\n      <td>https://en.wikipedia.org/wiki/Shahtaj%20Qizilbash</td>\n      <td>Shahtaj Qizilbash  was a women's right advocat...</td>\n      <td>Shahtaj Qizilbash</td>\n      <td>139</td>\n      <td>who is shahtaj qizilbash?</td>\n    </tr>\n    <tr>\n      <th>2</th>\n      <td>5106811</td>\n      <td>https://en.wikipedia.org/wiki/Shahtaj%20Qizilbash</td>\n      <td>Shahtaj Qizilbash  was a women's right advocat...</td>\n      <td>Shahtaj Qizilbash</td>\n      <td>139</td>\n      <td>who is shahtaj qizilbash</td>\n    </tr>\n    <tr>\n      <th>3</th>\n      <td>5106829</td>\n      <td>https://en.wikipedia.org/wiki/Aurat%20Foundation</td>\n      <td>Aurat Foundation, founded in 1986, is a women'...</td>\n      <td>Aurat Foundation</td>\n      <td>86</td>\n      <td>where is aurat?</td>\n    </tr>\n    <tr>\n      <th>4</th>\n      <td>5106829</td>\n      <td>https://en.wikipedia.org/wiki/Aurat%20Foundation</td>\n      <td>Aurat Foundation, founded in 1986, is a women'...</td>\n      <td>Aurat Foundation</td>\n      <td>86</td>\n      <td>aurat foundation</td>\n    </tr>\n    <tr>\n      <th>...</th>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n    </tr>\n    <tr>\n      <th>2849995</th>\n      <td>156656</td>\n      <td>https://en.wikipedia.org/wiki/Oath%20of%20Alle...</td>\n      <td>The Irish Oath of Allegiance  was a controvers...</td>\n      <td>Oath of Allegiance (Ireland)</td>\n      <td>90</td>\n      <td>when did the irish oath of allegiance take place</td>\n    </tr>\n    <tr>\n      <th>2849996</th>\n      <td>156656</td>\n      <td>https://en.wikipedia.org/wiki/Oath%20of%20Alle...</td>\n      <td>The Irish Oath of Allegiance  was a controvers...</td>\n      <td>Oath of Allegiance (Ireland)</td>\n      <td>90</td>\n      <td>what was the irish oath of allegiance</td>\n    </tr>\n    <tr>\n      <th>2849997</th>\n      <td>156658</td>\n      <td>https://en.wikipedia.org/wiki/Wiktionary</td>\n      <td>Wiktionary is a multilingual, web-based projec...</td>\n      <td>Wiktionary</td>\n      <td>167</td>\n      <td>where is wikikidnary</td>\n    </tr>\n    <tr>\n      <th>2849998</th>\n      <td>156658</td>\n      <td>https://en.wikipedia.org/wiki/Wiktionary</td>\n      <td>Wiktionary is a multilingual, web-based projec...</td>\n      <td>Wiktionary</td>\n      <td>167</td>\n      <td>definition of wikipedia</td>\n    </tr>\n    <tr>\n      <th>2849999</th>\n      <td>156658</td>\n      <td>https://en.wikipedia.org/wiki/Wiktionary</td>\n      <td>Wiktionary is a multilingual, web-based projec...</td>\n      <td>Wiktionary</td>\n      <td>167</td>\n      <td>what is the meaning of wiktionary</td>\n    </tr>\n  </tbody>\n</table>\n<p>2850000 rows × 6 columns</p>\n</div>"
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# flatten the queries\n",
    "query_list, id_list, url_list, text_list, title_list, word_num_list = [], [], [], [], [], []\n",
    "for i, row in tqdm.tqdm(concatenated_df.iterrows(), total=len(concatenated_df)):\n",
    "    for query in row[\"queries\"]:\n",
    "        id_list.append(row[\"id\"])\n",
    "        url_list.append(row[\"url\"])\n",
    "        text_list.append(row[\"text\"])\n",
    "        title_list.append(row[\"title\"])\n",
    "        word_num_list.append(row[\"word_num\"])\n",
    "        query_list.append(query)\n",
    "query_df = pd.DataFrame(\n",
    "    {\n",
    "        \"id\": id_list,\n",
    "        \"url\": url_list,\n",
    "        \"text\": text_list,\n",
    "        \"title\": title_list,\n",
    "        \"word_num\": word_num_list,\n",
    "        \"query\": query_list,\n",
    "    }\n",
    ")\n",
    "query_df"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\MW\\AppData\\Local\\Temp\\ipykernel_33600\\2992388678.py:12: SettingWithCopyWarning: \n",
      "A value is trying to be set on a copy of a slice from a DataFrame.\n",
      "Try using .loc[row_indexer,col_indexer] = value instead\n",
      "\n",
      "See the caveats in the documentation: https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#returning-a-view-versus-a-copy\n",
      "  query_df['parent'] = query_df['query'].apply(get_parent_tag)\n"
     ]
    },
    {
     "data": {
      "text/plain": "              id                                                url  \\\n0        5106811  https://en.wikipedia.org/wiki/Shahtaj%20Qizilbash   \n1        5106811  https://en.wikipedia.org/wiki/Shahtaj%20Qizilbash   \n2        5106811  https://en.wikipedia.org/wiki/Shahtaj%20Qizilbash   \n3        5106829   https://en.wikipedia.org/wiki/Aurat%20Foundation   \n4        5106829   https://en.wikipedia.org/wiki/Aurat%20Foundation   \n...          ...                                                ...   \n2849995   156656  https://en.wikipedia.org/wiki/Oath%20of%20Alle...   \n2849996   156656  https://en.wikipedia.org/wiki/Oath%20of%20Alle...   \n2849997   156658           https://en.wikipedia.org/wiki/Wiktionary   \n2849998   156658           https://en.wikipedia.org/wiki/Wiktionary   \n2849999   156658           https://en.wikipedia.org/wiki/Wiktionary   \n\n                                                      text  \\\n0        Shahtaj Qizilbash  was a women's right advocat...   \n1        Shahtaj Qizilbash  was a women's right advocat...   \n2        Shahtaj Qizilbash  was a women's right advocat...   \n3        Aurat Foundation, founded in 1986, is a women'...   \n4        Aurat Foundation, founded in 1986, is a women'...   \n...                                                    ...   \n2849995  The Irish Oath of Allegiance  was a controvers...   \n2849996  The Irish Oath of Allegiance  was a controvers...   \n2849997  Wiktionary is a multilingual, web-based projec...   \n2849998  Wiktionary is a multilingual, web-based projec...   \n2849999  Wiktionary is a multilingual, web-based projec...   \n\n                                title  word_num  \\\n0                   Shahtaj Qizilbash       139   \n1                   Shahtaj Qizilbash       139   \n2                   Shahtaj Qizilbash       139   \n3                    Aurat Foundation        86   \n4                    Aurat Foundation        86   \n...                               ...       ...   \n2849995  Oath of Allegiance (Ireland)        90   \n2849996  Oath of Allegiance (Ireland)        90   \n2849997                    Wiktionary       167   \n2849998                    Wiktionary       167   \n2849999                    Wiktionary       167   \n\n                                                    query   parent  \n0                               who was shahtaj qizilbash      who  \n1                               who is shahtaj qizilbash?      who  \n2                                who is shahtaj qizilbash      who  \n3                                         where is aurat?    where  \n4                                        aurat foundation  Unknown  \n...                                                   ...      ...  \n2849995  when did the irish oath of allegiance take place     when  \n2849996             what was the irish oath of allegiance     what  \n2849997                              where is wikikidnary    where  \n2849998                           definition of wikipedia   define  \n2849999                 what is the meaning of wiktionary     what  \n\n[2648487 rows x 7 columns]",
      "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>url</th>\n      <th>text</th>\n      <th>title</th>\n      <th>word_num</th>\n      <th>query</th>\n      <th>parent</th>\n    </tr>\n  </thead>\n  <tbody>\n    <tr>\n      <th>0</th>\n      <td>5106811</td>\n      <td>https://en.wikipedia.org/wiki/Shahtaj%20Qizilbash</td>\n      <td>Shahtaj Qizilbash  was a women's right advocat...</td>\n      <td>Shahtaj Qizilbash</td>\n      <td>139</td>\n      <td>who was shahtaj qizilbash</td>\n      <td>who</td>\n    </tr>\n    <tr>\n      <th>1</th>\n      <td>5106811</td>\n      <td>https://en.wikipedia.org/wiki/Shahtaj%20Qizilbash</td>\n      <td>Shahtaj Qizilbash  was a women's right advocat...</td>\n      <td>Shahtaj Qizilbash</td>\n      <td>139</td>\n      <td>who is shahtaj qizilbash?</td>\n      <td>who</td>\n    </tr>\n    <tr>\n      <th>2</th>\n      <td>5106811</td>\n      <td>https://en.wikipedia.org/wiki/Shahtaj%20Qizilbash</td>\n      <td>Shahtaj Qizilbash  was a women's right advocat...</td>\n      <td>Shahtaj Qizilbash</td>\n      <td>139</td>\n      <td>who is shahtaj qizilbash</td>\n      <td>who</td>\n    </tr>\n    <tr>\n      <th>3</th>\n      <td>5106829</td>\n      <td>https://en.wikipedia.org/wiki/Aurat%20Foundation</td>\n      <td>Aurat Foundation, founded in 1986, is a women'...</td>\n      <td>Aurat Foundation</td>\n      <td>86</td>\n      <td>where is aurat?</td>\n      <td>where</td>\n    </tr>\n    <tr>\n      <th>4</th>\n      <td>5106829</td>\n      <td>https://en.wikipedia.org/wiki/Aurat%20Foundation</td>\n      <td>Aurat Foundation, founded in 1986, is a women'...</td>\n      <td>Aurat Foundation</td>\n      <td>86</td>\n      <td>aurat foundation</td>\n      <td>Unknown</td>\n    </tr>\n    <tr>\n      <th>...</th>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n    </tr>\n    <tr>\n      <th>2849995</th>\n      <td>156656</td>\n      <td>https://en.wikipedia.org/wiki/Oath%20of%20Alle...</td>\n      <td>The Irish Oath of Allegiance  was a controvers...</td>\n      <td>Oath of Allegiance (Ireland)</td>\n      <td>90</td>\n      <td>when did the irish oath of allegiance take place</td>\n      <td>when</td>\n    </tr>\n    <tr>\n      <th>2849996</th>\n      <td>156656</td>\n      <td>https://en.wikipedia.org/wiki/Oath%20of%20Alle...</td>\n      <td>The Irish Oath of Allegiance  was a controvers...</td>\n      <td>Oath of Allegiance (Ireland)</td>\n      <td>90</td>\n      <td>what was the irish oath of allegiance</td>\n      <td>what</td>\n    </tr>\n    <tr>\n      <th>2849997</th>\n      <td>156658</td>\n      <td>https://en.wikipedia.org/wiki/Wiktionary</td>\n      <td>Wiktionary is a multilingual, web-based projec...</td>\n      <td>Wiktionary</td>\n      <td>167</td>\n      <td>where is wikikidnary</td>\n      <td>where</td>\n    </tr>\n    <tr>\n      <th>2849998</th>\n      <td>156658</td>\n      <td>https://en.wikipedia.org/wiki/Wiktionary</td>\n      <td>Wiktionary is a multilingual, web-based projec...</td>\n      <td>Wiktionary</td>\n      <td>167</td>\n      <td>definition of wikipedia</td>\n      <td>define</td>\n    </tr>\n    <tr>\n      <th>2849999</th>\n      <td>156658</td>\n      <td>https://en.wikipedia.org/wiki/Wiktionary</td>\n      <td>Wiktionary is a multilingual, web-based projec...</td>\n      <td>Wiktionary</td>\n      <td>167</td>\n      <td>what is the meaning of wiktionary</td>\n      <td>what</td>\n    </tr>\n  </tbody>\n</table>\n<p>2648487 rows × 7 columns</p>\n</div>"
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query_df = query_df.drop_duplicates()\n",
    "\n",
    "\n",
    "def get_parent_tag(query):\n",
    "    question_tag = [\"who\", \"when\", \"what\", \"which\", \"where\", \"how\", \"why\"]\n",
    "    for tag in question_tag:\n",
    "        if tag in query:\n",
    "            return tag\n",
    "    if \"is\" == query[:2] or \"are\" == query[:3] or \"was\" == query[:3] or \"were\" == query[:4]:\n",
    "        return \"is/are/was/were\"\n",
    "    if \"define\" in query or \"definition\" in query:\n",
    "        return \"define\"\n",
    "    return \"Unknown\"\n",
    "\n",
    "\n",
    "query_df[\"parent\"] = query_df[\"query\"].apply(get_parent_tag)\n",
    "query_df"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "outputs": [
    {
     "data": {
      "text/plain": "([<matplotlib.patches.Wedge at 0x2038f10add0>,\n  <matplotlib.patches.Wedge at 0x2038f10ace0>,\n  <matplotlib.patches.Wedge at 0x2038f10bac0>,\n  <matplotlib.patches.Wedge at 0x2038f148190>,\n  <matplotlib.patches.Wedge at 0x2038f148820>,\n  <matplotlib.patches.Wedge at 0x2038f148eb0>,\n  <matplotlib.patches.Wedge at 0x2038f149540>,\n  <matplotlib.patches.Wedge at 0x2038f149bd0>,\n  <matplotlib.patches.Wedge at 0x2038f14a260>,\n  <matplotlib.patches.Wedge at 0x2038f14a8f0>],\n [Text(0.6639542683200764, 0.8770203701052515, 'what'),\n  Text(-1.0210867085051922, 0.4091233722412259, 'who'),\n  Text(-0.30258321384542464, -1.057564843732513, 'where'),\n  Text(0.7146903846452939, -0.8361923547220234, 'when'),\n  Text(0.939944335886077, -0.5714058500187776, 'Unknown'),\n  Text(1.0180705476485639, -0.41657215462696645, 'which'),\n  Text(1.0629847816910225, -0.2829546852294363, 'define'),\n  Text(1.0867731997921213, -0.170070609493805, 'how'),\n  Text(1.0972074519613497, -0.07833139447553998, 'why'),\n  Text(1.0998588401189247, -0.017621912843214328, 'is/are/was/were')],\n [Text(0.36215687362913257, 0.47837474733013713, '777952'),\n  Text(-0.5569563864573776, 0.22315820304066866, '771311'),\n  Text(-0.16504538937023158, -0.5768535511268251, '639276'),\n  Text(0.389831118897433, -0.4561049207574672, '191862'),\n  Text(0.5126969104833147, -0.31167591819206053, '75688'),\n  Text(0.5553112078083076, -0.22722117525107258, '57365'),\n  Text(0.5798098809223758, -0.15433891921605614, '50744'),\n  Text(0.5927853817047933, -0.09276578699662091, '37711'),\n  Text(0.598476791978918, -0.04272621516847635, '33072'),\n  Text(0.5999230037012316, -0.009611952459935088, '13506')])"
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": "<Figure size 1000x1000 with 1 Axes>",
      "image/png": "[encoded data, 88724 characters]\n"
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# create pie chart for query_df\n",
    "# text should be like 28.8% (n=1234)\n",
    "plt.figure(figsize=(10, 10))\n",
    "values = query_df[\"parent\"].value_counts()\n",
    "plt.pie(\n",
    "    values, labels=query_df[\"parent\"].value_counts().index, autopct=lambda x: \"{:.0f}\".format(x * values.sum() / 100)\n",
    ")"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "outputs": [
    {
     "data": {
      "text/plain": "              id                                                url  \\\n4        5106829   https://en.wikipedia.org/wiki/Aurat%20Foundation   \n7        5106839          https://en.wikipedia.org/wiki/BloomingOUT   \n9        5106855  https://en.wikipedia.org/wiki/The%20Moment%20o...   \n124      5107172        https://en.wikipedia.org/wiki/Aubin%20Codex   \n169      5107366  https://en.wikipedia.org/wiki/Speed%20Demon%20...   \n...          ...                                                ...   \n2849905   156587        https://en.wikipedia.org/wiki/Death%20Eater   \n2849928   156598     https://en.wikipedia.org/wiki/Great%20Northern   \n2849967   156628  https://en.wikipedia.org/wiki/The%20Bad%20and%...   \n2849968   156628  https://en.wikipedia.org/wiki/The%20Bad%20and%...   \n2849976   156635         https://en.wikipedia.org/wiki/Glossophobia   \n\n                                                      text  \\\n4        Aurat Foundation, founded in 1986, is a women'...   \n7        bloomingOUT is an LGBT+ radio show broadcast o...   \n9        The Moment of the Magician  is a fantasy novel...   \n124      The Aubin Codex is an 81-leaf Aztec codex writ...   \n169      Speed Demon is a 2003 Canadian-American horror...   \n...                                                    ...   \n2849905  The Death Eaters are characters featured in th...   \n2849928                       Great Northern may refer to:   \n2849967  The Bad and the Beautiful is a 1952 American m...   \n2849968  The Bad and the Beautiful is a 1952 American m...   \n2849976  Glossophobia or speech anxiety is the fear of ...   \n\n                              title  word_num  \\\n4                  Aurat Foundation        86   \n7                       BloomingOUT       153   \n9        The Moment of the Magician        49   \n124                     Aubin Codex        48   \n169         Speed Demon (2003 film)        23   \n...                             ...       ...   \n2849905                 Death Eater       204   \n2849928              Great Northern         5   \n2849967   The Bad and the Beautiful       270   \n2849968   The Bad and the Beautiful       270   \n2849976                Glossophobia       175   \n\n                                                     query   parent  \n4                                         aurat foundation  Unknown  \n7                                     blooming out indiana  Unknown  \n9                              moment of the magician book  Unknown  \n124                                        the aubin codex  Unknown  \n169                              id speed demon movie cast  Unknown  \n...                                                    ...      ...  \n2849905                    harry potter death eaters names  Unknown  \n2849928                       great northern slang meaning  Unknown  \n2849967  the song for the bad and the beautiful theme o...  Unknown  \n2849968                      cast of the bad and beautiful  Unknown  \n2849976                                glossophobia causes  Unknown  \n\n[75688 rows x 7 columns]",
      "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>url</th>\n      <th>text</th>\n      <th>title</th>\n      <th>word_num</th>\n      <th>query</th>\n      <th>parent</th>\n    </tr>\n  </thead>\n  <tbody>\n    <tr>\n      <th>4</th>\n      <td>5106829</td>\n      <td>https://en.wikipedia.org/wiki/Aurat%20Foundation</td>\n      <td>Aurat Foundation, founded in 1986, is a women'...</td>\n      <td>Aurat Foundation</td>\n      <td>86</td>\n      <td>aurat foundation</td>\n      <td>Unknown</td>\n    </tr>\n    <tr>\n      <th>7</th>\n      <td>5106839</td>\n      <td>https://en.wikipedia.org/wiki/BloomingOUT</td>\n      <td>bloomingOUT is an LGBT+ radio show broadcast o...</td>\n      <td>BloomingOUT</td>\n      <td>153</td>\n      <td>blooming out indiana</td>\n      <td>Unknown</td>\n    </tr>\n    <tr>\n      <th>9</th>\n      <td>5106855</td>\n      <td>https://en.wikipedia.org/wiki/The%20Moment%20o...</td>\n      <td>The Moment of the Magician  is a fantasy novel...</td>\n      <td>The Moment of the Magician</td>\n      <td>49</td>\n      <td>moment of the magician book</td>\n      <td>Unknown</td>\n    </tr>\n    <tr>\n      <th>124</th>\n      <td>5107172</td>\n      <td>https://en.wikipedia.org/wiki/Aubin%20Codex</td>\n      <td>The Aubin Codex is an 81-leaf Aztec codex writ...</td>\n      <td>Aubin Codex</td>\n      <td>48</td>\n      <td>the aubin codex</td>\n      <td>Unknown</td>\n    </tr>\n    <tr>\n      <th>169</th>\n      <td>5107366</td>\n      <td>https://en.wikipedia.org/wiki/Speed%20Demon%20...</td>\n      <td>Speed Demon is a 2003 Canadian-American horror...</td>\n      <td>Speed Demon (2003 film)</td>\n      <td>23</td>\n      <td>id speed demon movie cast</td>\n      <td>Unknown</td>\n    </tr>\n    <tr>\n      <th>...</th>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n    </tr>\n    <tr>\n      <th>2849905</th>\n      <td>156587</td>\n      <td>https://en.wikipedia.org/wiki/Death%20Eater</td>\n      <td>The Death Eaters are characters featured in th...</td>\n      <td>Death Eater</td>\n      <td>204</td>\n      <td>harry potter death eaters names</td>\n      <td>Unknown</td>\n    </tr>\n    <tr>\n      <th>2849928</th>\n      <td>156598</td>\n      <td>https://en.wikipedia.org/wiki/Great%20Northern</td>\n      <td>Great Northern may refer to:</td>\n      <td>Great Northern</td>\n      <td>5</td>\n      <td>great northern slang meaning</td>\n      <td>Unknown</td>\n    </tr>\n    <tr>\n      <th>2849967</th>\n      <td>156628</td>\n      <td>https://en.wikipedia.org/wiki/The%20Bad%20and%...</td>\n      <td>The Bad and the Beautiful is a 1952 American m...</td>\n      <td>The Bad and the Beautiful</td>\n      <td>270</td>\n      <td>the song for the bad and the beautiful theme o...</td>\n      <td>Unknown</td>\n    </tr>\n    <tr>\n      <th>2849968</th>\n      <td>156628</td>\n      <td>https://en.wikipedia.org/wiki/The%20Bad%20and%...</td>\n      <td>The Bad and the Beautiful is a 1952 American m...</td>\n      <td>The Bad and the Beautiful</td>\n      <td>270</td>\n      <td>cast of the bad and beautiful</td>\n      <td>Unknown</td>\n    </tr>\n    <tr>\n      <th>2849976</th>\n      <td>156635</td>\n      <td>https://en.wikipedia.org/wiki/Glossophobia</td>\n      <td>Glossophobia or speech anxiety is the fear of ...</td>\n      <td>Glossophobia</td>\n      <td>175</td>\n      <td>glossophobia causes</td>\n      <td>Unknown</td>\n    </tr>\n  </tbody>\n</table>\n<p>75688 rows × 7 columns</p>\n</div>"
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query_df[(query_df[\"parent\"] == \"Unknown\")]\n",
    "# query_df[(query_df['parent'] == 'Unknown') & (query_df['query'].str.contains('define'))]"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "outputs": [
    {
     "data": {
      "text/plain": "               id                                                url  \\\n2430719   2300253  https://en.wikipedia.org/wiki/Passenger%20to%2...   \n1009479    744577  https://en.wikipedia.org/wiki/Indoor%20Footbal...   \n1848132  50950148       https://en.wikipedia.org/wiki/Katie%20Umback   \n176027     262411  https://en.wikipedia.org/wiki/Anton%20Alberts%...   \n289371     354408      https://en.wikipedia.org/wiki/Arshile%20Gorky   \n...           ...                                                ...   \n2582496   2582061       https://en.wikipedia.org/wiki/A.%20M.%20Naik   \n2338613   2194726   https://en.wikipedia.org/wiki/Jesse%20E.%20James   \n1904500  51160949  https://en.wikipedia.org/wiki/Amanda%20Benoit-...   \n2328121   2177821  https://en.wikipedia.org/wiki/Will%20to%20Powe...   \n1700683   1768814   https://en.wikipedia.org/wiki/Antoine%20Duquesne   \n\n                                                      text  \\\n2430719  Passenger to Frankfurt: An Extravaganza is a s...   \n1009479  The Indoor Football League  began in 1999 as a...   \n1848132  Katie-Marie Umback  is an Australian para-eque...   \n176027   Anton  Alberts  was a Dutch architect best kno...   \n289371   Arshile Gorky  was an Armenian-American painte...   \n...                                                    ...   \n2582496  Anil Manibhai Naik  is an Indian industrialist...   \n2338613  Jesse Edwards \"Tim\" James  was the only surviv...   \n1904500  Amanda Benoit-Wark  is a Canadian ice hockey p...   \n2328121  Will to Power is an American dance-pop group t...   \n1700683  Antoine Duquesne  was a Belgian politician and...   \n\n                                      title  word_num  \\\n2430719              Passenger to Frankfurt       272   \n1009479  Indoor Football League (1999–2000)       209   \n1848132                        Katie Umback        18   \n176027            Anton Alberts (architect)       185   \n289371                        Arshile Gorky        87   \n...                                     ...       ...   \n2582496                          A. M. Naik        82   \n2338613                      Jesse E. James        48   \n1904500                  Amanda Benoit-Wark        41   \n2328121                Will to Power (band)       117   \n1700683                    Antoine Duquesne        42   \n\n                                               query           parent  \n2430719             who wrote passenger to frankfurt              who  \n1009479                             who owns the ifl              who  \n1848132                   who is the para equestrian              who  \n176027            who is architect alexander alberts              who  \n289371                                  who is gorky              who  \n...                                              ...              ...  \n2582496               is anil manibhai naik an india  is/are/was/were  \n2338613            was jesse james born in tennessee  is/are/was/were  \n1904500           is amanda benoit ark hockey player  is/are/was/were  \n2328121  is will to power in the united states music  is/are/was/were  \n1700683                 is anthony duquesne a french  is/are/was/were  \n\n[10000 rows x 7 columns]",
      "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>url</th>\n      <th>text</th>\n      <th>title</th>\n      <th>word_num</th>\n      <th>query</th>\n      <th>parent</th>\n    </tr>\n  </thead>\n  <tbody>\n    <tr>\n      <th>2430719</th>\n      <td>2300253</td>\n      <td>https://en.wikipedia.org/wiki/Passenger%20to%2...</td>\n      <td>Passenger to Frankfurt: An Extravaganza is a s...</td>\n      <td>Passenger to Frankfurt</td>\n      <td>272</td>\n      <td>who wrote passenger to frankfurt</td>\n      <td>who</td>\n    </tr>\n    <tr>\n      <th>1009479</th>\n      <td>744577</td>\n      <td>https://en.wikipedia.org/wiki/Indoor%20Footbal...</td>\n      <td>The Indoor Football League  began in 1999 as a...</td>\n      <td>Indoor Football League (1999–2000)</td>\n      <td>209</td>\n      <td>who owns the ifl</td>\n      <td>who</td>\n    </tr>\n    <tr>\n      <th>1848132</th>\n      <td>50950148</td>\n      <td>https://en.wikipedia.org/wiki/Katie%20Umback</td>\n      <td>Katie-Marie Umback  is an Australian para-eque...</td>\n      <td>Katie Umback</td>\n      <td>18</td>\n      <td>who is the para equestrian</td>\n      <td>who</td>\n    </tr>\n    <tr>\n      <th>176027</th>\n      <td>262411</td>\n      <td>https://en.wikipedia.org/wiki/Anton%20Alberts%...</td>\n      <td>Anton  Alberts  was a Dutch architect best kno...</td>\n      <td>Anton Alberts (architect)</td>\n      <td>185</td>\n      <td>who is architect alexander alberts</td>\n      <td>who</td>\n    </tr>\n    <tr>\n      <th>289371</th>\n      <td>354408</td>\n      <td>https://en.wikipedia.org/wiki/Arshile%20Gorky</td>\n      <td>Arshile Gorky  was an Armenian-American painte...</td>\n      <td>Arshile Gorky</td>\n      <td>87</td>\n      <td>who is gorky</td>\n      <td>who</td>\n    </tr>\n    <tr>\n      <th>...</th>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n    </tr>\n    <tr>\n      <th>2582496</th>\n      <td>2582061</td>\n      <td>https://en.wikipedia.org/wiki/A.%20M.%20Naik</td>\n      <td>Anil Manibhai Naik  is an Indian industrialist...</td>\n      <td>A. M. Naik</td>\n      <td>82</td>\n      <td>is anil manibhai naik an india</td>\n      <td>is/are/was/were</td>\n    </tr>\n    <tr>\n      <th>2338613</th>\n      <td>2194726</td>\n      <td>https://en.wikipedia.org/wiki/Jesse%20E.%20James</td>\n      <td>Jesse Edwards \"Tim\" James  was the only surviv...</td>\n      <td>Jesse E. James</td>\n      <td>48</td>\n      <td>was jesse james born in tennessee</td>\n      <td>is/are/was/were</td>\n    </tr>\n    <tr>\n      <th>1904500</th>\n      <td>51160949</td>\n      <td>https://en.wikipedia.org/wiki/Amanda%20Benoit-...</td>\n      <td>Amanda Benoit-Wark  is a Canadian ice hockey p...</td>\n      <td>Amanda Benoit-Wark</td>\n      <td>41</td>\n      <td>is amanda benoit ark hockey player</td>\n      <td>is/are/was/were</td>\n    </tr>\n    <tr>\n      <th>2328121</th>\n      <td>2177821</td>\n      <td>https://en.wikipedia.org/wiki/Will%20to%20Powe...</td>\n      <td>Will to Power is an American dance-pop group t...</td>\n      <td>Will to Power (band)</td>\n      <td>117</td>\n      <td>is will to power in the united states music</td>\n      <td>is/are/was/were</td>\n    </tr>\n    <tr>\n      <th>1700683</th>\n      <td>1768814</td>\n      <td>https://en.wikipedia.org/wiki/Antoine%20Duquesne</td>\n      <td>Antoine Duquesne  was a Belgian politician and...</td>\n      <td>Antoine Duquesne</td>\n      <td>42</td>\n      <td>is anthony duquesne a french</td>\n      <td>is/are/was/were</td>\n    </tr>\n  </tbody>\n</table>\n<p>10000 rows × 7 columns</p>\n</div>"
     },
     "execution_count": 56,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "parent_tags = query_df[\"parent\"].unique()\n",
    "sampled_df_list = []\n",
    "for parent_tag in parent_tags:\n",
    "    df = query_df[query_df[\"parent\"] == parent_tag]\n",
    "    # randomly get 100 rows with seed\n",
    "    df = df.sample(n=1000, random_state=1)\n",
    "    sampled_df_list.append(df)\n",
    "sampled_df = pd.concat(sampled_df_list)\n",
    "sampled_df"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "outputs": [],
   "source": [
    "sampled_df.to_csv(\"wiki_qa_bart_10000row_input.csv\", index=False)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "outputs": [],
   "source": [
    "# plt.figure(figsize=(10,10))\n",
    "# values = sampled_df['parent'].value_counts()\n",
    "# plt.pie(values, labels=sampled_df['parent'].value_counts().index,\n",
    "#         autopct= lambda x: '{:.0f}'.format(x*values.sum()/100)\n",
    "#         )"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(device)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "outputs": [],
   "source": [
    "from transformers import pipeline\n",
    "import tiktoken\n",
    "\n",
    "\n",
    "def num_tokens_from_string(string: str) -> int:\n",
    "    encoding = tiktoken.encoding_for_model(\"gpt-3.5-turbo\")\n",
    "    num_tokens = len(encoding.encode(string))\n",
    "    return num_tokens\n",
    "\n",
    "\n",
    "row = sampled_df.iloc[0]\n",
    "summarizer = pipeline(\"summarization\", model=\"facebook/bart-large-cnn\")"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Your max_length is set to 353, but you input_length is only 346. You might consider decreasing max_length manually, e.g. summarizer('...', max_length=173)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "who wrote passenger to frankfurt\n",
      "text (token: 353)\n",
      "Passenger to Frankfurt: An Extravaganza is a spy novel by Agatha Christie first published in the United Kingdom by the Collins Crime Club in September 1970 and in the United States by Dodd, Mead and Company later in the same year. The UK edition retailed at twenty-five shillings. In preparation for decimalisation on 15 February 1971, it was concurrently priced on the dustjacket at £1.25. The US edition retailed at $5.95.\n",
      "It was published to mark Christie's eightieth birthday and, by counting up both UK and US short-story collections to reach the desired total, was also advertised as her eightieth book. It is the last of her spy novels. At the beginning of the book there is a quote by Jan Smuts, \"Leadership, besides being a great creative force, can be diabolical ...\"\n",
      "Sir Stafford Nye, a middle-aged diplomat, steps into the world of spies, double agents, and secret groups to effect a change in international power centres. He meets a woman who has selected him to aid her at a crucial point, when a weather delay changes where her and his aeroplane flight lands before proceeding to England. There is much commentary on the changes in the world, especially college age youth in Europe, the United States, and South America, in the late 1960s.\n",
      "The novel received mixed reviews at publication and in 1990. In 2017, it was assessed favourably in an essay about speculative spy thriller novels by women. It is one of only four Christie novels not to have received an adaptation of any kind, the others being Death Comes as the End, Destination Unknown and Postern of Fate.\n",
      "output\n",
      "[{'summary_text': \"The Passenger to Frankfurt: An Extravaganza is a spy novel by Agatha Christie. It was published to mark Christie's eightieth birthday and was advertised as her eightieth book. It is one of only four of her novels not to have received an adaptation.\"}]\n"
     ]
    }
   ],
   "source": [
    "# one sample test\n",
    "print(row[\"query\"])\n",
    "num_tokens = num_tokens_from_string(row[\"text\"])\n",
    "print(f\"text (token: {num_tokens})\")\n",
    "print(row[\"text\"])\n",
    "print(\"output\")\n",
    "print(summarizer(row[\"text\"], max_length=num_tokens, min_length=40, do_sample=False))"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "outputs": [
    {
     "data": {
      "text/plain": "\"The Passenger to Frankfurt: An Extravaganza is a spy novel by Agatha Christie. It was published to mark Christie's eightieth birthday and was advertised as her eightieth book. It is one of only four of her novels not to have received an adaptation.\""
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response = summarizer(row[\"text\"], max_length=int(num_tokens * 0.7), min_length=40, do_sample=False)[0][\"summary_text\"]\n",
    "response"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "outputs": [
    {
     "data": {
      "text/plain": "            id                                                url  \\\n0      2300253  https://en.wikipedia.org/wiki/Passenger%20to%2...   \n1       744577  https://en.wikipedia.org/wiki/Indoor%20Footbal...   \n2     50950148       https://en.wikipedia.org/wiki/Katie%20Umback   \n3       262411  https://en.wikipedia.org/wiki/Anton%20Alberts%...   \n4       354408      https://en.wikipedia.org/wiki/Arshile%20Gorky   \n...        ...                                                ...   \n9995   2582061       https://en.wikipedia.org/wiki/A.%20M.%20Naik   \n9996   2194726   https://en.wikipedia.org/wiki/Jesse%20E.%20James   \n9997  51160949  https://en.wikipedia.org/wiki/Amanda%20Benoit-...   \n9998   2177821  https://en.wikipedia.org/wiki/Will%20to%20Powe...   \n9999   1768814   https://en.wikipedia.org/wiki/Antoine%20Duquesne   \n\n                                                   text  \\\n0     Passenger to Frankfurt: An Extravaganza is a s...   \n1     The Indoor Football League  began in 1999 as a...   \n2     Katie-Marie Umback  is an Australian para-eque...   \n3     Anton  Alberts  was a Dutch architect best kno...   \n4     Arshile Gorky  was an Armenian-American painte...   \n...                                                 ...   \n9995  Anil Manibhai Naik  is an Indian industrialist...   \n9996  Jesse Edwards \"Tim\" James  was the only surviv...   \n9997  Amanda Benoit-Wark  is a Canadian ice hockey p...   \n9998  Will to Power is an American dance-pop group t...   \n9999  Antoine Duquesne  was a Belgian politician and...   \n\n                                   title  word_num  \\\n0                 Passenger to Frankfurt       272   \n1     Indoor Football League (1999–2000)       209   \n2                           Katie Umback        18   \n3              Anton Alberts (architect)       185   \n4                          Arshile Gorky        87   \n...                                  ...       ...   \n9995                          A. M. Naik        82   \n9996                      Jesse E. James        48   \n9997                  Amanda Benoit-Wark        41   \n9998                Will to Power (band)       117   \n9999                    Antoine Duquesne        42   \n\n                                            query           parent  \\\n0                who wrote passenger to frankfurt              who   \n1                                who owns the ifl              who   \n2                      who is the para equestrian              who   \n3              who is architect alexander alberts              who   \n4                                    who is gorky              who   \n...                                           ...              ...   \n9995               is anil manibhai naik an india  is/are/was/were   \n9996            was jesse james born in tennessee  is/are/was/were   \n9997           is amanda benoit ark hockey player  is/are/was/were   \n9998  is will to power in the united states music  is/are/was/were   \n9999                 is anthony duquesne a french  is/are/was/were   \n\n                                               response  \n0     The Passenger to Frankfurt: An Extravaganza is...  \n1     The Indoor Football League was founded in 1999...  \n2     Katie-Marie Umback is an Australian para-eques...  \n3     Anton Alberts was a Dutch architect best known...  \n4     Arshile Gorky was an Armenian-American painter...  \n...                                                 ...  \n9995  Anil Manibhai Naik is an Indian industrialist,...  \n9996  Jesse Edwards \"Tim\" James was the only survivi...  \n9997  Amanda Benoit-Wark played for the Canada women...  \n9998  Will to Power is an American dance-pop group t...  \n9999  Antoine Duquesne was a Belgian politician.Memb...  \n\n[10000 rows x 8 columns]",
      "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>url</th>\n      <th>text</th>\n      <th>title</th>\n      <th>word_num</th>\n      <th>query</th>\n      <th>parent</th>\n      <th>response</th>\n    </tr>\n  </thead>\n  <tbody>\n    <tr>\n      <th>0</th>\n      <td>2300253</td>\n      <td>https://en.wikipedia.org/wiki/Passenger%20to%2...</td>\n      <td>Passenger to Frankfurt: An Extravaganza is a s...</td>\n      <td>Passenger to Frankfurt</td>\n      <td>272</td>\n      <td>who wrote passenger to frankfurt</td>\n      <td>who</td>\n      <td>The Passenger to Frankfurt: An Extravaganza is...</td>\n    </tr>\n    <tr>\n      <th>1</th>\n      <td>744577</td>\n      <td>https://en.wikipedia.org/wiki/Indoor%20Footbal...</td>\n      <td>The Indoor Football League  began in 1999 as a...</td>\n      <td>Indoor Football League (1999–2000)</td>\n      <td>209</td>\n      <td>who owns the ifl</td>\n      <td>who</td>\n      <td>The Indoor Football League was founded in 1999...</td>\n    </tr>\n    <tr>\n      <th>2</th>\n      <td>50950148</td>\n      <td>https://en.wikipedia.org/wiki/Katie%20Umback</td>\n      <td>Katie-Marie Umback  is an Australian para-eque...</td>\n      <td>Katie Umback</td>\n      <td>18</td>\n      <td>who is the para equestrian</td>\n      <td>who</td>\n      <td>Katie-Marie Umback is an Australian para-eques...</td>\n    </tr>\n    <tr>\n      <th>3</th>\n      <td>262411</td>\n      <td>https://en.wikipedia.org/wiki/Anton%20Alberts%...</td>\n      <td>Anton  Alberts  was a Dutch architect best kno...</td>\n      <td>Anton Alberts (architect)</td>\n      <td>185</td>\n      <td>who is architect alexander alberts</td>\n      <td>who</td>\n      <td>Anton Alberts was a Dutch architect best known...</td>\n    </tr>\n    <tr>\n      <th>4</th>\n      <td>354408</td>\n      <td>https://en.wikipedia.org/wiki/Arshile%20Gorky</td>\n      <td>Arshile Gorky  was an Armenian-American painte...</td>\n      <td>Arshile Gorky</td>\n      <td>87</td>\n      <td>who is gorky</td>\n      <td>who</td>\n      <td>Arshile Gorky was an Armenian-American painter...</td>\n    </tr>\n    <tr>\n      <th>...</th>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n    </tr>\n    <tr>\n      <th>9995</th>\n      <td>2582061</td>\n      <td>https://en.wikipedia.org/wiki/A.%20M.%20Naik</td>\n      <td>Anil Manibhai Naik  is an Indian industrialist...</td>\n      <td>A. M. Naik</td>\n      <td>82</td>\n      <td>is anil manibhai naik an india</td>\n      <td>is/are/was/were</td>\n      <td>Anil Manibhai Naik is an Indian industrialist,...</td>\n    </tr>\n    <tr>\n      <th>9996</th>\n      <td>2194726</td>\n      <td>https://en.wikipedia.org/wiki/Jesse%20E.%20James</td>\n      <td>Jesse Edwards \"Tim\" James  was the only surviv...</td>\n      <td>Jesse E. James</td>\n      <td>48</td>\n      <td>was jesse james born in tennessee</td>\n      <td>is/are/was/were</td>\n      <td>Jesse Edwards \"Tim\" James was the only survivi...</td>\n    </tr>\n    <tr>\n      <th>9997</th>\n      <td>51160949</td>\n      <td>https://en.wikipedia.org/wiki/Amanda%20Benoit-...</td>\n      <td>Amanda Benoit-Wark  is a Canadian ice hockey p...</td>\n      <td>Amanda Benoit-Wark</td>\n      <td>41</td>\n      <td>is amanda benoit ark hockey player</td>\n      <td>is/are/was/were</td>\n      <td>Amanda Benoit-Wark played for the Canada women...</td>\n    </tr>\n    <tr>\n      <th>9998</th>\n      <td>2177821</td>\n      <td>https://en.wikipedia.org/wiki/Will%20to%20Powe...</td>\n      <td>Will to Power is an American dance-pop group t...</td>\n      <td>Will to Power (band)</td>\n      <td>117</td>\n      <td>is will to power in the united states music</td>\n      <td>is/are/was/were</td>\n      <td>Will to Power is an American dance-pop group t...</td>\n    </tr>\n    <tr>\n      <th>9999</th>\n      <td>1768814</td>\n      <td>https://en.wikipedia.org/wiki/Antoine%20Duquesne</td>\n      <td>Antoine Duquesne  was a Belgian politician and...</td>\n      <td>Antoine Duquesne</td>\n      <td>42</td>\n      <td>is anthony duquesne a french</td>\n      <td>is/are/was/were</td>\n      <td>Antoine Duquesne was a Belgian politician.Memb...</td>\n    </tr>\n  </tbody>\n</table>\n<p>10000 rows × 8 columns</p>\n</div>"
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# multi sample output\n",
    "sampled_output_df = pd.read_csv(\"wiki_qa_bart_10000row.csv\")\n",
    "sampled_output_df\n",
    "# for index, row in sampled_output_df[:50].iterrows():\n",
    "#     print(index)\n",
    "#     print(row['query'])\n",
    "#     print(row['response'])"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 who\n",
      "who wrote passenger to frankfurt\n",
      "The Passenger to Frankfurt: An Extravaganza is a spy novel by Agatha Christie. It was published to mark Christie's eightieth birthday and was advertised as her eightieth book. It is one of only four of her novels not to have received an adaptation.\n",
      "1 who\n",
      "who owns the ifl\n",
      "The Indoor Football League was founded in 1999. It was an offshoot of the troubled PIFL. The league was successful enough for a major expansion in 2000. After the season, the league was purchased by the Arena Football League's Orlando Predators.\n",
      "2 who\n",
      "who is the para equestrian\n",
      "Katie-Marie Umback is an Australian para-equestrian. She represented Australia at the 2016 Rio Paralympics. Umback won a silver medal in the women's individual\n",
      "3 who\n",
      "who is architect alexander alberts\n",
      "Anton Alberts was a Dutch architect best known for the ING Bank in Amsterdam and the Gas Corporation headquarters in Groningen. He was involved with Situationist International before being expelled in 1960. In 1963 he founded the firm Alberts and Van Huut with Max van Huut.\n",
      "4 who\n",
      "who is gorky\n",
      "Arshile Gorky was an Armenian-American painter who had a seminal influence on Abstract Expressionism. Along with Mark Rothko, Jackson Pollock and Willem de Kooning, Gorky has been hailed as one of the most powerful American painters of the 20th century.\n",
      "1000 where\n",
      "where is harcourt house\n",
      "Harcourt House Artist Run Centre is an artist-run, charitable organization that promotes contemporary visual art. The Centre includes two public gallery spaces; and as the single largest community of visual artists in Edmonton, it offers 42 low-rent studios.\n",
      "1001 where\n",
      "where is choye commune\n",
      "Choye is a commune in the Haute-Saône department in the region of Bourgogne-Franche-Comté in eastern France.It is\n",
      "1002 where\n",
      "where was kea bouman born\n",
      "Cornelia \"Kea\" Tiedemann-Bouman was a female tennis player from the Netherlands. She won the singles title at the 1927 French Championships. She is the only Dutch woman to have won a Grand Slam singles title.\n",
      "1003 where\n",
      "where is perseus arm\n",
      "Perseus Arm is one of two major spiral arms of the Milky Way galaxy. Previously thought to be 13,000 light-years away, it is now thought to lie 6,400 light years from the Solar System.\n",
      "1004 where\n",
      "where is st-livrade france\n",
      "Sainte-Livrade is a commune in the Haute-Garonne department in southwestern France.It is located in the south-west of the province of Ha\n",
      "2000 Unknown\n",
      "ningal meaning\n",
      "Ningal was a Mesopotamian goddess of Sumerian origin regarded as the wife of the moon god, Nanna/Sin. She was particularly closely associated with his main cult centers, Ur and Harran, but they were also worshiped together in other cities of Mesopotamia.\n",
      "2001 Unknown\n",
      "vers sur meuge france\n",
      "Vers-sur-Méouge is a commune in the Drôme department in southeastern France.It is located in the south-east of the department, near the border\n",
      "2002 Unknown\n",
      "carl barron comedian\n",
      "Carl Barron is an Australian theatre and television comedian. His style is based on observational humour. He is best known for his work on the ABC's The Biggest Loser. He has also\n",
      "2003 Unknown\n",
      "paradoris caerulea\n",
      "Paradoris caerulea is a species of sea slug.It is a dorid nudibranch, a shell-less marine opisthobranch gastropod m\n",
      "2004 Unknown\n",
      "cast of london town\n",
      "The film stars Daniel Huttlestone, Dougray Scott, Natascha McElhone, Nell Williams and Jonathan Rhys Meyers. It had its world premiere at the Los Angeles Film Festival on June 3, 2016. The film was released in a limited release and through video on demand on October 7, 2016, by\n",
      "3000 what\n",
      "what was longmont potion castle\n",
      "Longmont Potion Castle is the pseudonym of an anonymous prank caller active in Colorado and the Los Angeles areas since 1986. Details about his personal life are scarce, and his real name is unknown to the public. Over the years\n",
      "3001 what\n",
      "what year did the jungle show premiere\n",
      "\"The Jungle\" is episode 77 of the American television anthology series The Twilight Zone. It first aired on December 1, 1961. It is the 77th episode of the series and the\n",
      "3002 what\n",
      "what is bae systems largest export agreement\n",
      "The first sales occurred in September 1985 and the most recent contract for 72 Eurofighter Typhoon multirole fighters was signed in August 2006. It is Britain's largest ever export agreement, and employs at least 5,000 people in Saudi Arabia.\n",
      "3003 what\n",
      "what type of university is quenzhou normal\n",
      "Quanzhou Normal University is a public university in Quanzhou, Fujian province, People's Republic of China.It is located in the city centre of the city of Quanz\n",
      "3004 what\n",
      "what is israel astronomical association\n",
      "The Israeli Astronomical Association  is an Israeli nonprofit organization. Its purpose is to deepen and distribute the awareness for the field of astronomy among the Israeli public. The organization is based in Tel Aviv\n",
      "4000 define\n",
      "defence hockey definition\n",
      "Defence or defense is a player whose primary responsibility is to prevent the opposing team from scoring. They are often referred to as defencemen, D, D-men or blueliners. In regular play, two defencemen complement three forwards and a goaltender.\n",
      "4001 define\n",
      "wein's law definition\n",
      "Wien approximation is an equation used to describe the short-wavelength  spectrum of thermal radiation. Wien displacement law describes the relationship between the temperature of an object and the peak wavelength or\n",
      "4002 define\n",
      "lovingkindness definition\n",
      "Chesed is an English translation of Chesed, a term found in the Hebrew Bible. Loving-kindness is a term used in Buddhism to refer to love and kindness. Chesed\n",
      "4003 define\n",
      "definition of fernandina\n",
      "Fernandina is a Spanish name for a woman.It may refer to any of a number of women.It is also known as \"Fernando\" or \"Fernanda\"\n",
      "4004 define\n",
      "definition of pti\n",
      " PTI may refer to: PTI, PTi, PTI or PTI-I.It may also refer to the following people: \"PTI\" or \"PTi\n",
      "5000 when\n",
      "when was south africa founded\n",
      "The following lists events that happened during 1891 in South Africa.It was the first year of the Second Boer War, and the first of the Third Boer era in South African history\n",
      "5001 when\n",
      "when is the cleveland hustles season\n",
      "Cleveland Hustles was an hour-long American reality TV series created by LeBron James and Maverick Carter. In it, aspiring entrepreneurs competed to open one of four physical stores in the Gordon Square Arts District under the mentorship of a Cleveland\n",
      "5002 when\n",
      "when was 50 cent is the future released\n",
      "50 Cent Is the Future is the second mixtape by American rapper 50 Cent and first one by his rap group G-Unit. It was released on June 1, 2002 via Street Dance/Thurd World Muzic. The lone guest appearance is provided by UTP, which marks the first collaboration between the group and future member Young Buck.\n",
      "5003 when\n",
      "when is 343\n",
      "Year 343 was a common year starting on Saturday of the Julian calendar. At the time, it was known as the Year of the Consulship of Memmius and Romulus. The denomination 343 for this year has been used since the early medieval\n",
      "5004 when\n",
      "when was lonely places by august derleth released\n",
      "Lonesome Places is a collection of fantasy and horror short stories by American author August Derleth. It was released in 1962 by Arkham House in an edition of 2,201 copies.\n",
      "6000 how\n",
      "victoria's parliament has how many members\n",
      "The Parliament of Victoria is the bicameral legislature of the Australian state of Victoria. It consists of the Queen, the Legislative Assembly and the Legislative Council. The current Parliament was elected on 24 November 2018, sworn in on 19 December 2018.\n",
      "6001 how\n",
      "how many scottish mps in burgh constituency\n",
      "New constituency boundaries were first used in the 1918 general election. In Scotland the legislation defined 32 burgh constituencies, 38 county constituencies and one university constituency. For the 1950 general election, new boundaries were introduced under the House of Commons  Act 1949.\n",
      "6002 how\n",
      "how many years ago did christopher knowles start painting\n",
      "Christopher Knowles is an American poet and painter. In 1976, his poetry was used by Robert Wilson for the avant-garde minimalist Philip Glass opera, Einstein on the Beach. Knowles has received a diagnosis of possible brain damage. He is often referred to as autistic.\n",
      "6003 how\n",
      "how many universities in turkey\n",
      "There are 207 universities and academies in Turkey. There are 129 state universities, 74 private foundation universities, four two-year granting institutions. This is a list of universities in Turkey, sorted\n",
      "6004 how\n",
      "how many federated states\n",
      "The following table lists Mexico's 32 federated entities, ranked by total continental surface. The table is ordered by the size of the country's continental surface, with the highest ranked nation at the\n",
      "7000 which\n",
      "which region is tettnang\n",
      "Tettnang is a town in the Bodensee district in southern Baden-Württemberg in a region of Germany known as Swabia. It lies 7 kilometres from Lake Constance. The region\n",
      "7001 which\n",
      "which family is the antbirds\n",
      "The Grallariidae are a family of small passerine birds of Central and South America. They lay 1 to 6 eggs in a nest in a tree, both sexes incubating. The long, powerful legs and an essentially vestigial tail aid this lifestyle.\n",
      "7002 which\n",
      "which ncaa division is eastern michigan\n",
      "The 1981 Eastern Michigan Hurons football team represented Eastern Michigan University in the 1981 NCAA Division I-A football season. The Hurons finished in last place in the Mid-American Conference. The team's statistical leaders included J.F. Green with 1,391 passing yards, Ricky Calhoun with 971 rushing yards, and Jeff Dackin with 440 receiving yards.\n",
      "7003 which\n",
      "which nobleman was george gordon\n",
      "George Gordon, 1st Marquess of Huntly was a Scottish nobleman. He took a leading role in the political and military life of Scotland in the late 16th century, and around\n",
      "7004 which\n",
      "which champions are defending champions?\n",
      "Karen Khachanov was the defending champion but chose not to participate.Malek Jaziri won the title after defeating Dudi Sela 1-6, 6-1,\n",
      "8000 why\n",
      "why did grant surrender fort henry\n",
      "The Battle of Fort Henry was fought on February 6, 1862, in Donelson, Stewart County, Tennessee. It was the first important victory for the Union and Brig. Gen. Ulysses S. Grant in the Western Theater.\n",
      "8001 why\n",
      "why was the gal and pal invented\n",
      "The Generic Array Logic was an innovation of the PAL. It was able to take the place of many PAL device types. A similar device called a PEEL was introduced by the International CMOS Technology  corporation.\n",
      "8002 why\n",
      "why did the last us space shuttle leave earth after the biosatellite mission\n",
      "Biosatellite 3 was the last unmanned U.S. Biosatellite program for biological research. The intent had been to fly a 6 kg male Southern pig-tailed macaque in Earth-orbit for 30 days. However, after only 8.8 days in orbit, the mission was terminated because of the subject's deteriorating health.\n",
      "8003 why\n",
      "why is austria a good tourist destination\n",
      "Austria has one guest bed for every six inhabitants. In 2007, Austria ranked 9th worldwide in international tourism receipts, with 18.9 billion US$. In international tourist arrivals, Austriaranked 12th with 30.8 million tourists.\n",
      "8004 why\n",
      "why was compass a failure\n",
      "Operation Compass was the first large British military operation of the Western Desert Campaign during the Second World War. British Empire forces attacked Italian forces of the 10th Army  in western Egypt and Cyrenaica, the eastern province of Libya, from December 1940 to February 1941.\n",
      "9000 is/are/was/were\n",
      "washington state route 97\n",
      "Wagon roads along the west side of the Columbia River were plotted in the 1880s by American settlers. The wagon road was designated as State Road 10 and was completely paved by the late 1930s. US 97 was created in 1926 and followed the highway, which was relocated in the 1950s due to the construction of the Rocky Reach Dam.\n",
      "9001 is/are/was/were\n",
      "is red mountain pass colorado\n",
      "Red Mountain Pass is a mountain pass in the San Juan Mountains of western Colorado. Red Mountain Pass elevation  is 1,071 feet. Red mountain pass is in the western San Juan\n",
      "9002 is/are/was/were\n",
      "was clayton waagner a bank robber\n",
      "Clayton Lee Waagner is a convicted bank robber and anti-abortion terrorist. He was the FBI's 467th fugitive to be placed on the Ten Most Wanted list on September 21, 2001. He is currently in prison.\n",
      "9003 is/are/was/were\n",
      "is myrrh a flower?\n",
      "Myrrh is a gum-resin extracted from a number of small, thorny tree species of the genus Commiphora. Myrrh resin has been used throughout history as a perfume, incense and medicine.\n",
      "9004 is/are/was/were\n",
      "are there first nations in saskatchewan canada\n",
      "First Nations ethnicities in the province include the Cree, Assiniboine, Saulteaux, Lakota, Dene and Dakota. Historically, the Atsina and Blackfoot could also be found at various times.\n"
     ]
    }
   ],
   "source": [
    "parents = sampled_output_df[\"parent\"].unique()\n",
    "for parent in parents:\n",
    "    for index, row in sampled_output_df[sampled_output_df[\"parent\"] == parent].head(5).iterrows():\n",
    "        print(index, parent)\n",
    "        print(row[\"query\"])\n",
    "        print(row[\"response\"])"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "outputs": [
    {
     "data": {
      "text/plain": "                                                 query  \\\n0                     who wrote passenger to frankfurt   \n1                                     who owns the ifl   \n2                           who is the para equestrian   \n3                   who is architect alexander alberts   \n4                                         who is gorky   \n..                                                 ...   \n995                    who is the director of cornered   \n996                            who was mikra kassabian   \n997  who was the first leader of the third hellenic...   \n998                                 who was isacovici?   \n999                                   who is ethridge?   \n\n                                              response  \n0    The Passenger to Frankfurt: An Extravaganza is...  \n1    The Indoor Football League was founded in 1999...  \n2    Katie-Marie Umback is an Australian para-eques...  \n3    Anton Alberts was a Dutch architect best known...  \n4    Arshile Gorky was an Armenian-American painter...  \n..                                                 ...  \n995  Cornered is a 1945 film noir starring Dick Pow...  \n996  Mihran Krikor Kassabian was an Armenian-Americ...  \n997  Michail Stasinopoulos was a Greek jurist and p...  \n998  Salomon Isacovici was a Jewish Holocaust survi...  \n999  Ethridge is a British surname.It may refer to ...  \n\n[1000 rows x 2 columns]",
      "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>query</th>\n      <th>response</th>\n    </tr>\n  </thead>\n  <tbody>\n    <tr>\n      <th>0</th>\n      <td>who wrote passenger to frankfurt</td>\n      <td>The Passenger to Frankfurt: An Extravaganza is...</td>\n    </tr>\n    <tr>\n      <th>1</th>\n      <td>who owns the ifl</td>\n      <td>The Indoor Football League was founded in 1999...</td>\n    </tr>\n    <tr>\n      <th>2</th>\n      <td>who is the para equestrian</td>\n      <td>Katie-Marie Umback is an Australian para-eques...</td>\n    </tr>\n    <tr>\n      <th>3</th>\n      <td>who is architect alexander alberts</td>\n      <td>Anton Alberts was a Dutch architect best known...</td>\n    </tr>\n    <tr>\n      <th>4</th>\n      <td>who is gorky</td>\n      <td>Arshile Gorky was an Armenian-American painter...</td>\n    </tr>\n    <tr>\n      <th>...</th>\n      <td>...</td>\n      <td>...</td>\n    </tr>\n    <tr>\n      <th>995</th>\n      <td>who is the director of cornered</td>\n      <td>Cornered is a 1945 film noir starring Dick Pow...</td>\n    </tr>\n    <tr>\n      <th>996</th>\n      <td>who was mikra kassabian</td>\n      <td>Mihran Krikor Kassabian was an Armenian-Americ...</td>\n    </tr>\n    <tr>\n      <th>997</th>\n      <td>who was the first leader of the third hellenic...</td>\n      <td>Michail Stasinopoulos was a Greek jurist and p...</td>\n    </tr>\n    <tr>\n      <th>998</th>\n      <td>who was isacovici?</td>\n      <td>Salomon Isacovici was a Jewish Holocaust survi...</td>\n    </tr>\n    <tr>\n      <th>999</th>\n      <td>who is ethridge?</td>\n      <td>Ethridge is a British surname.It may refer to ...</td>\n    </tr>\n  </tbody>\n</table>\n<p>1000 rows × 2 columns</p>\n</div>"
     },
     "execution_count": 64,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sampled_output_df[sampled_output_df[\"parent\"] == \"who\"][[\"query\", \"response\"]]"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "outputs": [
    {
     "data": {
      "text/plain": "id                                                    1704412\nurl         https://en.wikipedia.org/wiki/1891%20in%20Sout...\ntext        The following lists events that happened durin...\ntitle                                    1891 in South Africa\nword_num                                                   11\nquery                           when was south africa founded\nparent                                                   when\nresponse    The following lists events that happened durin...\nName: 500, dtype: object"
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sampled_output_df.iloc[500]"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "searchgpt_oa_bart",
   "language": "python",
   "display_name": "searchgpt_oa_bart"
  },
  "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": 0
}
