notebooks/transformer_splits_5_folds.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Packages\n",
    "\n",
    "- [`iterstat`](https://github.com/trent-b/iterative-stratification) - multilabel stratification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "import gc\n",
    "import pickle\n",
    "import numpy as np\n",
    "import pandas as ps\n",
    "from tqdm import tqdm\n",
    "from pathlib import Path\n",
    "from itertools import chain\n",
    "from collections import Counter\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sbn\n",
    "from urllib.parse import urlparse\n",
    "from sklearn.preprocessing import OneHotEncoder\n",
    "\n",
    "from sklearn.model_selection import train_test_split, KFold, GroupKFold\n",
    "from iterstrat.ml_stratifiers import MultilabelStratifiedKFold"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_dir = Path('..') / 'data'\n",
    "embeddings_dir = data_dir / 'embeddings'\n",
    "results_dir = data_dir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[01;34m../data\u001b[00m\n",
      "├── \u001b[01;34malbert_base\u001b[00m\n",
      "│   ├── 30k-clean.model\n",
      "│   ├── 30k-clean.vocab\n",
      "│   ├── albert_config.json\n",
      "│   ├── model.ckpt-best.data-00000-of-00001\n",
      "│   ├── model.ckpt-best.index\n",
      "│   ├── model.ckpt-best.meta\n",
      "│   └── spiece.model\n",
      "├── \u001b[01;31malbert_base_v2.tar.gz\u001b[00m\n",
      "├── \u001b[01;31mcased_L-12_H-768_A-12.zip\u001b[00m\n",
      "├── \u001b[01;34membeddings\u001b[00m\n",
      "│   ├── bert_base.pkl\n",
      "│   ├── crawl-300d-2M.pkl\n",
      "│   ├── \u001b[01;31mcrawl-300d-2M.pkl.zip\u001b[00m\n",
      "│   ├── glove.840B.300d.pkl\n",
      "│   ├── \u001b[01;31mglove.840B.300d.pkl.zip\u001b[00m\n",
      "│   ├── glove_crawl_emb.pkl\n",
      "│   ├── glove_crawl_wikinews_emb.pkl\n",
      "│   ├── wiki-news-300d-1M.pkl\n",
      "│   ├── wiki-news-300d-1M.vec\n",
      "│   └── \u001b[01;31mwikinews300d1mvec.zip\u001b[00m\n",
      "├── \u001b[01;34mfolds\u001b[00m\n",
      "│   ├── tgkf_train_00.json\n",
      "│   ├── tgkf_train_00.pkl\n",
      "│   ├── tgkf_train_01.json\n",
      "│   ├── tgkf_train_01.pkl\n",
      "│   ├── tgkf_train_02.json\n",
      "│   ├── tgkf_train_02.pkl\n",
      "│   ├── tgkf_train_03.json\n",
      "│   ├── tgkf_train_03.pkl\n",
      "│   ├── tgkf_train_04.json\n",
      "│   ├── tgkf_train_04.pkl\n",
      "│   ├── tgkf_train_0.json\n",
      "│   ├── tgkf_train_0.pkl\n",
      "│   ├── tgkf_train_1.json\n",
      "│   ├── tgkf_train_1.pkl\n",
      "│   ├── tgkf_train_2.json\n",
      "│   ├── tgkf_train_2.pkl\n",
      "│   ├── tgkf_valid_00.pkl\n",
      "│   ├── tgkf_valid_01.pkl\n",
      "│   ├── tgkf_valid_02.pkl\n",
      "│   ├── tgkf_valid_03.pkl\n",
      "│   ├── tgkf_valid_04.pkl\n",
      "│   ├── tgkf_valid_0.pkl\n",
      "│   ├── tgkf_valid_1.pkl\n",
      "│   ├── tgkf_valid_2.pkl\n",
      "│   ├── tkf_train_0.pkl\n",
      "│   ├── tkf_train_1.pkl\n",
      "│   ├── tkf_train_2.pkl\n",
      "│   ├── tkf_train_3.pkl\n",
      "│   ├── tkf_train_4.pkl\n",
      "│   ├── tkf_valid_0.pkl\n",
      "│   ├── tkf_valid_1.pkl\n",
      "│   ├── tkf_valid_2.pkl\n",
      "│   ├── tkf_valid_3.pkl\n",
      "│   ├── tkf_valid_4.pkl\n",
      "│   ├── train_0.pkl\n",
      "│   ├── train_1.pkl\n",
      "│   ├── train_2.pkl\n",
      "│   ├── train_3.pkl\n",
      "│   ├── train_4.pkl\n",
      "│   ├── transf_train_0.pkl\n",
      "│   ├── transf_train_1.pkl\n",
      "│   ├── transf_train_2.pkl\n",
      "│   ├── transf_train_3.pkl\n",
      "│   ├── transf_train_4.pkl\n",
      "│   ├── transf_valid_0.pkl\n",
      "│   ├── transf_valid_1.pkl\n",
      "│   ├── transf_valid_2.pkl\n",
      "│   ├── transf_valid_3.pkl\n",
      "│   ├── transf_valid_4.pkl\n",
      "│   ├── valid_0.pkl\n",
      "│   ├── valid_1.pkl\n",
      "│   ├── valid_2.pkl\n",
      "│   ├── valid_3.pkl\n",
      "│   └── valid_4.pkl\n",
      "├── \u001b[01;34mroberta_base\u001b[00m\n",
      "│   ├── merges.txt\n",
      "│   └── vocab.json\n",
      "├── sample_submission.csv\n",
      "├── test.csv\n",
      "├── test.pkl\n",
      "├── train.csv\n",
      "├── \u001b[01;31mtrain.csv.zip\u001b[00m\n",
      "├── train.pkl\n",
      "├── trans_test.pkl\n",
      "├── trans_train.pkl\n",
      "├── trans_valid.pkl\n",
      "├── \u001b[01;34muncased_L-12_H-768_A-12\u001b[00m\n",
      "│   ├── bert_config.json\n",
      "│   ├── bert_model.ckpt.data-00000-of-00001\n",
      "│   ├── bert_model.ckpt.index\n",
      "│   ├── bert_model.ckpt.meta\n",
      "│   └── vocab.txt\n",
      "├── \u001b[01;31muncased_L-12_H-768_A-12.zip\u001b[00m\n",
      "├── \u001b[01;34muncased_L-24_H-1024_A-16\u001b[00m\n",
      "│   ├── bert_config.json\n",
      "│   ├── bert_model.ckpt.data-00000-of-00001\n",
      "│   ├── bert_model.ckpt.index\n",
      "│   ├── bert_model.ckpt.meta\n",
      "│   └── vocab.txt\n",
      "├── \u001b[01;31muncased_L-24_H-1024_A-16.zip\u001b[00m\n",
      "├── valid.pkl\n",
      "├── vocab.json\n",
      "├── vocab.pkl\n",
      "├── vocab.txt\n",
      "├── \u001b[01;34mwwm_uncased_L-24_H-1024_A-16\u001b[00m\n",
      "│   ├── bert_config.json\n",
      "│   ├── bert_model.ckpt.data-00000-of-00001\n",
      "│   ├── bert_model.ckpt.index\n",
      "│   ├── bert_model.ckpt.meta\n",
      "│   └── vocab.txt\n",
      "├── \u001b[01;31mwwm_uncased_L-24_H-1024_A-16.zip\u001b[00m\n",
      "├── \u001b[01;34mxlnet\u001b[00m\n",
      "│   └── spiece.model\n",
      "└── \u001b[01;34mxlnet_cased_L-12_H-768_A-12\u001b[00m\n",
      "    ├── spiece.model\n",
      "    ├── xlnet_config.json\n",
      "    ├── xlnet_model.ckpt.data-00000-of-00001\n",
      "    ├── xlnet_model.ckpt.index\n",
      "    └── xlnet_model.ckpt.meta\n",
      "\n",
      "9 directories, 112 files\n"
     ]
    }
   ],
   "source": [
    "!tree ../data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Embedings\n",
    "\n",
    "Embedings downloaded from this kernel - [**Quest Q&A - LSTM Inference Only**](https://www.kaggle.com/chanhu/quest-q-a-lstm-inference-baseline)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(6079, 41)\n"
     ]
    }
   ],
   "source": [
    "train = ps.read_csv(data_dir / 'train.csv')\n",
    "print(train.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Index(['qa_id', 'question_title', 'question_body', 'question_user_name',\n",
       "       'question_user_page', 'answer', 'answer_user_name', 'answer_user_page',\n",
       "       'url', 'category', 'host', 'question_asker_intent_understanding',\n",
       "       'question_body_critical', 'question_conversational',\n",
       "       'question_expect_short_answer', 'question_fact_seeking',\n",
       "       'question_has_commonly_accepted_answer',\n",
       "       'question_interestingness_others', 'question_interestingness_self',\n",
       "       'question_multi_intent', 'question_not_really_a_question',\n",
       "       'question_opinion_seeking', 'question_type_choice',\n",
       "       'question_type_compare', 'question_type_consequence',\n",
       "       'question_type_definition', 'question_type_entity',\n",
       "       'question_type_instructions', 'question_type_procedure',\n",
       "       'question_type_reason_explanation', 'question_type_spelling',\n",
       "       'question_well_written', 'answer_helpful',\n",
       "       'answer_level_of_information', 'answer_plausible', 'answer_relevance',\n",
       "       'answer_satisfaction', 'answer_type_instructions',\n",
       "       'answer_type_procedure', 'answer_type_reason_explanation',\n",
       "       'answer_well_written'],\n",
       "      dtype='object')"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train.columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(476, 11)\n"
     ]
    }
   ],
   "source": [
    "test = ps.read_csv(data_dir / 'test.csv')\n",
    "print(test.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Index(['qa_id', 'question_title', 'question_body', 'question_user_name',\n",
       "       'question_user_page', 'answer', 'answer_user_name', 'answer_user_page',\n",
       "       'url', 'category', 'host'],\n",
       "      dtype='object')"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test.columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "targets = [\n",
    "    'question_asker_intent_understanding',\n",
    "    'question_body_critical',\n",
    "    'question_conversational',\n",
    "    'question_expect_short_answer',\n",
    "    'question_fact_seeking',\n",
    "    'question_has_commonly_accepted_answer',\n",
    "    'question_interestingness_others',\n",
    "    'question_interestingness_self',\n",
    "    'question_multi_intent',\n",
    "    'question_not_really_a_question',\n",
    "    'question_opinion_seeking',\n",
    "    'question_type_choice',\n",
    "    'question_type_compare',\n",
    "    'question_type_consequence',\n",
    "    'question_type_definition',\n",
    "    'question_type_entity',\n",
    "    'question_type_instructions',\n",
    "    'question_type_procedure',\n",
    "    'question_type_reason_explanation',\n",
    "    'question_type_spelling',\n",
    "    'question_well_written',\n",
    "    'answer_helpful',\n",
    "    'answer_level_of_information',\n",
    "    'answer_plausible',\n",
    "    'answer_relevance',\n",
    "    'answer_satisfaction',\n",
    "    'answer_type_instructions',\n",
    "    'answer_type_procedure',\n",
    "    'answer_type_reason_explanation',\n",
    "    'answer_well_written'    \n",
    "]\n",
    "\n",
    "text_columns = [\n",
    "    'question_title', \n",
    "    'question_body', \n",
    "    'answer'\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "65"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "unique_hosts = list(set(train['host'].unique().tolist() + test['host'].unique().tolist()))\n",
    "unique_hosts = sorted(unique_hosts)\n",
    "\n",
    "idx2host = [\"<unk>\"] + unique_hosts\n",
    "host2idx = {host: idx for idx, host in enumerate(idx2host)}\n",
    "\n",
    "len(host2idx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "6"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "unique_categories = list(set(train['category'].unique().tolist() + test['category'].unique().tolist()))\n",
    "unique_categories = sorted(unique_categories)\n",
    "\n",
    "idx2category = [\"<unk>\"] + unique_categories\n",
    "category2idx = {cat: idx for idx, cat in enumerate(idx2category)}\n",
    "\n",
    "len(category2idx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "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>qa_id</th>\n",
       "      <th>question_title</th>\n",
       "      <th>question_body</th>\n",
       "      <th>question_user_name</th>\n",
       "      <th>question_user_page</th>\n",
       "      <th>answer</th>\n",
       "      <th>answer_user_name</th>\n",
       "      <th>answer_user_page</th>\n",
       "      <th>url</th>\n",
       "      <th>category</th>\n",
       "      <th>host</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>0</td>\n",
       "      <td>39</td>\n",
       "      <td>Will leaving corpses lying around upset my pri...</td>\n",
       "      <td>I see questions/information online about how t...</td>\n",
       "      <td>Dylan</td>\n",
       "      <td>https://gaming.stackexchange.com/users/64471</td>\n",
       "      <td>There is no consequence for leaving corpses an...</td>\n",
       "      <td>Nelson868</td>\n",
       "      <td>https://gaming.stackexchange.com/users/97324</td>\n",
       "      <td>http://gaming.stackexchange.com/questions/1979...</td>\n",
       "      <td>CULTURE</td>\n",
       "      <td>gaming.stackexchange.com</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>46</td>\n",
       "      <td>Url link to feature image in the portfolio</td>\n",
       "      <td>I am new to Wordpress. i have issue with Featu...</td>\n",
       "      <td>Anu</td>\n",
       "      <td>https://wordpress.stackexchange.com/users/72927</td>\n",
       "      <td>I think it is possible with custom fields.\\n\\n...</td>\n",
       "      <td>Irina</td>\n",
       "      <td>https://wordpress.stackexchange.com/users/27233</td>\n",
       "      <td>http://wordpress.stackexchange.com/questions/1...</td>\n",
       "      <td>TECHNOLOGY</td>\n",
       "      <td>wordpress.stackexchange.com</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>70</td>\n",
       "      <td>Is accuracy, recoil or bullet spread affected ...</td>\n",
       "      <td>To experiment I started a bot game, toggled in...</td>\n",
       "      <td>Konsta</td>\n",
       "      <td>https://gaming.stackexchange.com/users/37545</td>\n",
       "      <td>You do not have armour in the screenshots. Thi...</td>\n",
       "      <td>Damon Smithies</td>\n",
       "      <td>https://gaming.stackexchange.com/users/70641</td>\n",
       "      <td>http://gaming.stackexchange.com/questions/2154...</td>\n",
       "      <td>CULTURE</td>\n",
       "      <td>gaming.stackexchange.com</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3</td>\n",
       "      <td>132</td>\n",
       "      <td>Suddenly got an I/O error from my external HDD</td>\n",
       "      <td>I have used my Raspberry Pi as a torrent-serve...</td>\n",
       "      <td>robbannn</td>\n",
       "      <td>https://raspberrypi.stackexchange.com/users/17341</td>\n",
       "      <td>Your Western Digital hard drive is disappearin...</td>\n",
       "      <td>HeatfanJohn</td>\n",
       "      <td>https://raspberrypi.stackexchange.com/users/1311</td>\n",
       "      <td>http://raspberrypi.stackexchange.com/questions...</td>\n",
       "      <td>TECHNOLOGY</td>\n",
       "      <td>raspberrypi.stackexchange.com</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4</td>\n",
       "      <td>200</td>\n",
       "      <td>Passenger Name - Flight Booking Passenger only...</td>\n",
       "      <td>I have bought Delhi-London return flights for ...</td>\n",
       "      <td>Amit</td>\n",
       "      <td>https://travel.stackexchange.com/users/29089</td>\n",
       "      <td>I called two persons who work for Saudia (tick...</td>\n",
       "      <td>Nean Der Thal</td>\n",
       "      <td>https://travel.stackexchange.com/users/10051</td>\n",
       "      <td>http://travel.stackexchange.com/questions/4704...</td>\n",
       "      <td>CULTURE</td>\n",
       "      <td>travel.stackexchange.com</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   qa_id                                     question_title  \\\n",
       "0     39  Will leaving corpses lying around upset my pri...   \n",
       "1     46         Url link to feature image in the portfolio   \n",
       "2     70  Is accuracy, recoil or bullet spread affected ...   \n",
       "3    132     Suddenly got an I/O error from my external HDD   \n",
       "4    200  Passenger Name - Flight Booking Passenger only...   \n",
       "\n",
       "                                       question_body question_user_name  \\\n",
       "0  I see questions/information online about how t...              Dylan   \n",
       "1  I am new to Wordpress. i have issue with Featu...                Anu   \n",
       "2  To experiment I started a bot game, toggled in...             Konsta   \n",
       "3  I have used my Raspberry Pi as a torrent-serve...           robbannn   \n",
       "4  I have bought Delhi-London return flights for ...               Amit   \n",
       "\n",
       "                                  question_user_page  \\\n",
       "0       https://gaming.stackexchange.com/users/64471   \n",
       "1    https://wordpress.stackexchange.com/users/72927   \n",
       "2       https://gaming.stackexchange.com/users/37545   \n",
       "3  https://raspberrypi.stackexchange.com/users/17341   \n",
       "4       https://travel.stackexchange.com/users/29089   \n",
       "\n",
       "                                              answer answer_user_name  \\\n",
       "0  There is no consequence for leaving corpses an...        Nelson868   \n",
       "1  I think it is possible with custom fields.\\n\\n...            Irina   \n",
       "2  You do not have armour in the screenshots. Thi...   Damon Smithies   \n",
       "3  Your Western Digital hard drive is disappearin...      HeatfanJohn   \n",
       "4  I called two persons who work for Saudia (tick...    Nean Der Thal   \n",
       "\n",
       "                                   answer_user_page  \\\n",
       "0      https://gaming.stackexchange.com/users/97324   \n",
       "1   https://wordpress.stackexchange.com/users/27233   \n",
       "2      https://gaming.stackexchange.com/users/70641   \n",
       "3  https://raspberrypi.stackexchange.com/users/1311   \n",
       "4      https://travel.stackexchange.com/users/10051   \n",
       "\n",
       "                                                 url    category  \\\n",
       "0  http://gaming.stackexchange.com/questions/1979...     CULTURE   \n",
       "1  http://wordpress.stackexchange.com/questions/1...  TECHNOLOGY   \n",
       "2  http://gaming.stackexchange.com/questions/2154...     CULTURE   \n",
       "3  http://raspberrypi.stackexchange.com/questions...  TECHNOLOGY   \n",
       "4  http://travel.stackexchange.com/questions/4704...     CULTURE   \n",
       "\n",
       "                            host  \n",
       "0       gaming.stackexchange.com  \n",
       "1    wordpress.stackexchange.com  \n",
       "2       gaming.stackexchange.com  \n",
       "3  raspberrypi.stackexchange.com  \n",
       "4       travel.stackexchange.com  "
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train_df, valid_df = train_test_split(train, test_size=0.2, random_state=2019)\n",
    "\n",
    "# print(train_df.shape)\n",
    "# print(valid_df.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# results_dir = data_dir\n",
    "\n",
    "\n",
    "# with open(results_dir / 'trans_train.pkl', 'wb') as f:\n",
    "#     pickle.dump(train_df, f)\n",
    "    \n",
    "    \n",
    "# with open(results_dir / 'trans_valid.pkl', 'wb') as f:\n",
    "#     pickle.dump(valid_df, f)\n",
    "    \n",
    "    \n",
    "# with open(results_dir / 'trans_test.pkl', 'wb') as f:\n",
    "#     pickle.dump(test, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # folds = MultilabelStratifiedKFold(n_splits=5, shuffle=True, random_state=2019)\n",
    "# folds = KFold(n_splits=5, shuffle=True, random_state=2019)\n",
    "\n",
    "# def dump_data(data, fname):\n",
    "#     with open(results_dir / 'folds' / fname, 'wb') as f:\n",
    "#         pickle.dump(data, f)\n",
    "\n",
    "# for idx, (train_idx, valid_idx) in enumerate(folds.split(train, train[targets].values)):\n",
    "#     _train = train.loc[train_idx]\n",
    "#     _valid = train.loc[valid_idx]\n",
    "    \n",
    "#     dump_data(_train, f'tkf_train_{idx}.pkl')\n",
    "#     dump_data(_valid, f'tkf_valid_{idx}.pkl')\n",
    "    \n",
    "#     print(f'Generated split for fold - {idx}', flush=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated split for fold - 0\n",
      "Generated split for fold - 1\n",
      "Generated split for fold - 2\n",
      "Generated split for fold - 3\n",
      "Generated split for fold - 4\n",
      "Generated split for fold - 5\n",
      "Generated split for fold - 6\n"
     ]
    }
   ],
   "source": [
    "folds = GroupKFold(n_splits=7)\n",
    "\n",
    "def dump_data(data, fname):\n",
    "    with open(results_dir / 'folds' / fname, 'wb') as f:\n",
    "        pickle.dump(data, f)\n",
    "\n",
    "for idx, (train_idx, valid_idx) in enumerate(folds.split(X=train['question_body'], groups=train['question_body'])):\n",
    "    _train = train.loc[train_idx]\n",
    "    _valid = train.loc[valid_idx]\n",
    "    \n",
    "    dump_data(_train, f'tgkf_train_7_0{idx}.pkl')\n",
    "    dump_data(_valid, f'tgkf_valid_7_0{idx}.pkl')\n",
    "    \n",
    "    print(f'Generated split for fold - {idx}', flush=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}