notebooks/create_dataset.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2b17588a",
   "metadata": {},
   "source": [
    "# Settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14ca09f7",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "# Desired percentage of NIL mentions\n",
    "perc_nil = 0.1\n",
    "\n",
    "# Number of batch to create\n",
    "n_batch = 10\n",
    "\n",
    "# Desired NIL mentions per batch (in dev and test)\n",
    "desired_nil = 50\n",
    "\n",
    "# Desired size of a single dev and test batch\n",
    "dev_test_desired_batch_size = 10000\n",
    "\n",
    "# Random state\n",
    "random_state = 1234\n",
    "\n",
    "# Output dir\n",
    "outdir = 'incremental_dataset'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0a36413",
   "metadata": {},
   "source": [
    "# Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "859e39d1",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from tqdm import tqdm,trange\n",
    "from glob import glob\n",
    "import numpy as np\n",
    "from sklearn.model_selection import StratifiedKFold\n",
    "import os\n",
    "import json"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "885f0b51",
   "metadata": {},
   "source": [
    "# Download"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4c073b3",
   "metadata": {},
   "source": [
    "Download the original dataset from https://github.com/yasumasaonoe/ET4EL"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5982edb1",
   "metadata": {},
   "source": [
    "# Extract it"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce8898c8",
   "metadata": {},
   "source": [
    "```\n",
    "tar -xzf unseen_mentions.tar.gz\n",
    "```\n",
    "\n",
    "You should be able to see a folder named `unseen_mentions` with the following content:\n",
    "```\n",
    "unseen_mentions/\n",
    "unseen_mentions/dev.json\n",
    "unseen_mentions/test.json\n",
    "unseen_mentions/train\n",
    "unseen_mentions/train/train_2.json\n",
    "unseen_mentions/train/train_1.json\n",
    "unseen_mentions/train/train_4.json\n",
    "unseen_mentions/train/train_3.json\n",
    "unseen_mentions/train/train_5.json\n",
    "unseen_mentions/train/train_0.json\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c82f0d0",
   "metadata": {},
   "source": [
    "# Load Original Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d2b5794",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "train_files = glob('unseen_mentions/train/*.json')\n",
    "train_files.sort()\n",
    "train_files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c482e1c",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "train_df = pd.DataFrame()\n",
    "for file in tqdm(train_files):\n",
    "    _temp = pd.read_json(file, lines=True)\n",
    "    print(file, 'shape: ', _temp.shape)\n",
    "    train_df = pd.concat([train_df, _temp], ignore_index=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17d56984",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "dev_df = pd.read_json('unseen_mentions/dev.json', lines=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df5f5088",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "test_df = pd.read_json('unseen_mentions/test.json', lines=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4da2361f",
   "metadata": {},
   "source": [
    "# Proceed"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35d47495",
   "metadata": {},
   "source": [
    "### Calculate the absolute frequency with which each entity is mentioned"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82110f21",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "mention_frequency = train_df.groupby('wikiId').size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1890f190",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "mention_frequency"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4c31b70",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "freq_df = pd.DataFrame(mention_frequency)\n",
    "freq_df.columns = ['freq']\n",
    "freq_df.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ee70a0d",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "freq_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b1b7d0c",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "# calculate the median frequency with which all the entities are mentioned\n",
    "med_freq = np.median(freq_df['freq'])\n",
    "med_freq"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eed4d9b1",
   "metadata": {},
   "source": [
    "Calculate the probability that each entity is NIL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae1fcb31",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "np.random.seed(random_state)\n",
    "\n",
    "freq_df['p_formula'] = perc_nil ** (freq_df['freq'] / med_freq)\n",
    "s = np.random.uniform(0, 1, freq_df.shape[0])\n",
    "freq_df['p_uniform'] = s\n",
    "freq_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf54768e",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "freq_df['NIL'] = freq_df['p_uniform'] < freq_df['p_formula']\n",
    "freq_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8dbfd88",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "print('N# NIL:', freq_df.eval('NIL').sum(), 'percentage:', freq_df.eval('NIL').sum() / freq_df.shape[0] * 100, '%')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a651d70",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "# Setting NIL freq to 0 so that when we split in batches each batch has the same number of NILs\n",
    "freq_df.loc[freq_df['NIL'] == True, 'freq'] = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65c7ea80",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "train_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ad8732b",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "train_df_merged = train_df.join(freq_df, how='left', on='wikiId')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e905eb08",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "print('Percentage of NIL mentions in train:', train_df_merged.eval('NIL').sum() / train_df_merged.shape[0] * 100)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04f46ae9",
   "metadata": {},
   "source": [
    "# Propagate NIL mentions in test and dev"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62d4e994",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "dev_df_merged = dev_df.join(freq_df, how='left', on='wikiId')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5e9783a",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "dev_df_merged['p_formula'] = dev_df_merged['p_formula'].fillna(-1)\n",
    "dev_df_merged['p_uniform'] = dev_df_merged['p_uniform'].fillna(-1)\n",
    "dev_df_merged['NIL'] = dev_df_merged['NIL'].fillna(False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eafa0322",
   "metadata": {},
   "source": [
    "Compute entity frequencies in dev for a correct sampling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b022384e",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "mention_frequency_dev = dev_df_merged.query('~NIL').groupby('wikiId').size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91336d5f",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "mention_frequency_dev"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "deeb44b4",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "freq_df_dev = pd.DataFrame(mention_frequency_dev)\n",
    "freq_df_dev.columns = ['freq']\n",
    "freq_df_dev.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3647986b",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "freq_df_dev.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6850e900",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "dev_df_merged = dev_df_merged.drop(columns=['freq']).join(freq_df_dev, how='left', on='wikiId')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "153b45c6",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "# set NIL entities freq to 0 for stratifying the NIL class\n",
    "dev_df_merged.loc[dev_df_merged['NIL'], 'freq'] = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2289da4",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "dev_df_merged['freq'] = dev_df_merged['freq'].astype(int)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f24c2cb",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "assert not dev_df_merged['NIL'].isna().any()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb5c2534",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "print('dev NIL N#:', dev_df_merged.eval('NIL').sum(), 'Percentage:', dev_df_merged.eval('NIL').sum() / dev_df_merged.shape[0] * 100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "825422ec",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "test_df_merged = test_df.join(freq_df, how='left', on='wikiId')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "051b6da0",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "test_df_merged['p_formula'] = test_df_merged['p_formula'].fillna(-1)\n",
    "test_df_merged['p_uniform'] = test_df_merged['p_uniform'].fillna(-1)\n",
    "test_df_merged['NIL'] = test_df_merged['NIL'].fillna(False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3122b16",
   "metadata": {},
   "source": [
    "Compute entity frequencies in test for a correct sampling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99c86357",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "mention_frequency_test = test_df_merged.query('~NIL').groupby('wikiId').size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b69f4d4",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "mention_frequency_test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1ee52d1",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "freq_df_test = pd.DataFrame(mention_frequency_test)\n",
    "freq_df_test.columns = ['freq']\n",
    "freq_df_test.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7efdcd11",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "freq_df_test.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7dbe2f7",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "test_df_merged = test_df_merged.drop(columns=['freq']).join(freq_df_test, how='left', on='wikiId')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c2af0f0",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "# set NIL entities freq to 0 for stratifying the NIL class\n",
    "test_df_merged.loc[test_df_merged['NIL'], 'freq'] = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c306ee4",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "test_df_merged['freq'] = test_df_merged['freq'].astype(int)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a04136fa",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "assert not test_df_merged['NIL'].isna().any()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c797eca8",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "print('test NIL N#:', test_df_merged.eval('NIL').sum(), 'Percentage:', test_df_merged.eval('NIL').sum() / test_df_merged.shape[0] * 100)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6bfe1203",
   "metadata": {},
   "source": [
    "# Transplant"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "baaf69dd",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "# I want 500 NIL in dev and test\n",
    "# I want 100k total in both dev and test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4082c893",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "train_unique = train_df[['word', 'wikiId']].drop_duplicates(keep='first')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36db344a",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "dev_unique = dev_df[['word', 'wikiId']].drop_duplicates(keep='first')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a34a5f8",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "test_unique = test_df[['word', 'wikiId']].drop_duplicates(keep='first')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1fbce08",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "# ensure no duplicates\n",
    "assert not train_unique.duplicated().any()\n",
    "assert not dev_unique.duplicated().any()\n",
    "assert not test_unique.duplicated().any()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fde9161c",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "# ensure unseen mentions holds\n",
    "assert not pd.concat([train_unique, dev_unique]).duplicated().any()\n",
    "assert not pd.concat([train_unique, test_unique]).duplicated().any()\n",
    "assert not pd.concat([dev_unique, test_unique]).duplicated().any()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e60c91d6",
   "metadata": {},
   "source": [
    "## NIL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "780a2bf6",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "# prendo n menzioni NIL dal train che non siano nel dev/test e ce le metto"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "518aaf71",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "train_df_nil = train_df_merged.query('NIL')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2db8196",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "desired_nil_total = desired_nil * n_batch\n",
    "nil_to_add_dev = desired_nil_total - dev_df_merged.eval('NIL').sum()\n",
    "nil_to_add_dev"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eba18553",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "nil_to_add_test = desired_nil_total - test_df_merged.eval('NIL').sum()\n",
    "nil_to_add_test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6272f9b2",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "def get_subset_sum_ge(df, sum_, random_state):\n",
    "    \"\"\"\n",
    "    Gets a subset whose sum is greater than sum_\n",
    "    \"\"\"\n",
    "    \n",
    "    #print(sum_)\n",
    "    \n",
    "    if df.shape[0] == 0:\n",
    "        return None\n",
    "    \n",
    "    # get a sample\n",
    "    item = df.sample(n=1, random_state=random_state)\n",
    "    random_state += 1\n",
    "    df = df.drop(item.index) # remove just extracted item\n",
    "    \n",
    "    subset = item.index\n",
    "\n",
    "    current_sum = item['count'].values[0]\n",
    "    while current_sum < sum_:\n",
    "        print(f'\\r{current_sum}/{sum_}',end='')\n",
    "        # get a sample\n",
    "        item = df.sample(n=1, random_state=random_state)\n",
    "        random_state += 1\n",
    "        df = df.drop(item.index) # remove just extracted item\n",
    "\n",
    "        current_sum += item['count'].values[0]\n",
    "        \n",
    "        subset = subset.union(item.index)\n",
    "    print()\n",
    "    return subset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44de2e7e",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "train_df_nil_unique_count = pd.DataFrame(\n",
    "    train_df_nil[['wikiId', 'word']].value_counts(), columns=['count']).reset_index()\n",
    "train_df_nil_unique_count.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11d52a24",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "sample_transplant_to_dev = get_subset_sum_ge(train_df_nil_unique_count, nil_to_add_dev, random_state)\n",
    "print(train_df_nil_unique_count.loc[sample_transplant_to_dev]['count'].sum(), '>=', nil_to_add_dev)\n",
    "assert train_df_nil_unique_count.loc[sample_transplant_to_dev]['count'].sum() >= nil_to_add_dev"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d6e637c",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "# remove sample_transplant to dev before extracting a sample for test\n",
    "train_df_nil_unique_count_no_dev = train_df_nil_unique_count.drop(sample_transplant_to_dev)\n",
    "sample_transplant_to_test = get_subset_sum_ge(train_df_nil_unique_count_no_dev, nil_to_add_test, random_state+10)\n",
    "print(train_df_nil_unique_count_no_dev.loc[sample_transplant_to_test]['count'].sum(), '>=', nil_to_add_test)\n",
    "assert train_df_nil_unique_count_no_dev.loc[sample_transplant_to_test]['count'].sum() >= nil_to_add_test"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9d7cfa5",
   "metadata": {},
   "source": [
    "## Transplant NIL to dev and test"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93989b41",
   "metadata": {},
   "source": [
    "### Dev"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa9ff32a",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "mentions_to_dev = train_df_nil.join(\n",
    "    train_df_nil_unique_count.loc[sample_transplant_to_dev].set_index(['wikiId', 'word']),\n",
    "    on=['wikiId', 'word'], how='inner')\n",
    "mentions_to_dev = mentions_to_dev.drop(columns=['count'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f76e43c",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "dev_df_transplant = pd.concat([dev_df_merged, mentions_to_dev], ignore_index=True)\n",
    "dev_df_transplant = dev_df_transplant.sample(frac=1, random_state=random_state) # randomize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c540975c",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "dev_df_transplant.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7ccbb1e",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "print(dev_df_transplant.eval('NIL').sum(), '>=', desired_nil_total)\n",
    "assert dev_df_transplant.eval('NIL').sum() >= desired_nil_total"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00320a01",
   "metadata": {},
   "source": [
    "### Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa24beed",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "mentions_to_test = train_df_nil.join(\n",
    "    train_df_nil_unique_count.loc[sample_transplant_to_test].set_index(['wikiId', 'word']),\n",
    "    on=['wikiId', 'word'], how='inner')\n",
    "mentions_to_test = mentions_to_test.drop(columns=['count'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "228fc356",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "test_df_transplant = pd.concat([test_df_merged, mentions_to_test], ignore_index=True)\n",
    "test_df_transplant = test_df_transplant.sample(frac=1, random_state=random_state) # randomize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2936f6a",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "test_df_transplant.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37ba5c81",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "print(test_df_transplant.eval('NIL').sum(), '>=', desired_nil_total)\n",
    "assert test_df_transplant.eval('NIL').sum() >= desired_nil_total"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5377a67b",
   "metadata": {},
   "source": [
    "### Remove from train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ed4748b",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "train_df_transplant = train_df_merged.drop(mentions_to_dev.index.union(mentions_to_test.index))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0cb74d11",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "train_unique_transplant = train_df_transplant[['word', 'wikiId']].drop_duplicates(keep='first')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d68bdf94",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "dev_unique_transplant = dev_df_transplant[['word', 'wikiId']].drop_duplicates(keep='first')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "218e3c77",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "test_unique_transplant = test_df_transplant[['word', 'wikiId']].drop_duplicates(keep='first')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c66684d",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "# ensure no duplicates\n",
    "assert not train_unique_transplant.duplicated().any()\n",
    "assert not dev_unique_transplant.duplicated().any()\n",
    "assert not test_unique_transplant.duplicated().any()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a155314c",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "# ensure unseen mentions holds\n",
    "assert not pd.concat([train_unique_transplant, dev_unique_transplant]).duplicated().any()\n",
    "assert not pd.concat([train_unique_transplant, test_unique_transplant]).duplicated().any()\n",
    "assert not pd.concat([dev_unique_transplant, test_unique_transplant]).duplicated().any()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af118754",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "print('Train', train_df_transplant.shape)\n",
    "print('Dev', dev_df_transplant.shape)\n",
    "print('Test', test_df_transplant.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b339808e",
   "metadata": {},
   "source": [
    "# Transplant not-NIL mentions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3baab260",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "# prendo n menzioni NIL dal train che non siano nel dev/test e ce le metto"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8c2e49e",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "train_df_not_nil = train_df_merged.query('~NIL')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86f5542b",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "desired_total = dev_test_desired_batch_size * n_batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dac66966",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "not_nil_to_add_dev = desired_total - dev_df_transplant.shape[0]\n",
    "not_nil_to_add_dev"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58f0bb02",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "not_nil_to_add_test = desired_total - test_df_transplant.shape[0]\n",
    "not_nil_to_add_test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88f69470",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "train_df_not_nil_unique_count = pd.DataFrame(\n",
    "    train_df_not_nil[['wikiId', 'word']].value_counts(), columns=['count']).reset_index()\n",
    "train_df_not_nil_unique_count.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a689ffd4",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "sample_transplant_to_dev_not_nil = get_subset_sum_ge(train_df_not_nil_unique_count, not_nil_to_add_dev, random_state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b73b684",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "print(train_df_not_nil_unique_count.loc[sample_transplant_to_dev_not_nil]['count'].sum(), \">=\", not_nil_to_add_dev)\n",
    "assert train_df_not_nil_unique_count.loc[sample_transplant_to_dev_not_nil]['count'].sum() >= not_nil_to_add_dev"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7de91012",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "# remove sample_transplant to dev before extracting a sample for test\n",
    "train_df_not_nil_unique_count_no_dev = train_df_not_nil_unique_count.drop(sample_transplant_to_dev_not_nil)\n",
    "sample_transplant_to_test_not_nil = get_subset_sum_ge(train_df_not_nil_unique_count_no_dev, not_nil_to_add_test, random_state+10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33df7972",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "print(train_df_not_nil_unique_count_no_dev.loc[sample_transplant_to_test_not_nil]['count'].sum(), \">=\", not_nil_to_add_test)\n",
    "assert train_df_not_nil_unique_count_no_dev.loc[sample_transplant_to_test_not_nil]['count'].sum() >= not_nil_to_add_test"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d049444f",
   "metadata": {},
   "source": [
    "### Dev"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5bd675d",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "mentions_to_dev_not_nil = train_df_not_nil.join(\n",
    "    train_df_not_nil_unique_count.loc[sample_transplant_to_dev_not_nil].set_index(['wikiId', 'word']),\n",
    "    on=['wikiId', 'word'], how='inner')\n",
    "mentions_to_dev_not_nil = mentions_to_dev_not_nil.drop(columns=['count'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f61dbf6d",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "assert not mentions_to_dev_not_nil.eval('NIL').any()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abbcaeb4",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "dev_df_transplant_final = pd.concat([dev_df_transplant, mentions_to_dev_not_nil], ignore_index=True)\n",
    "dev_df_transplant_final = dev_df_transplant_final.sample(frac=1, random_state=random_state) # randomize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50b7cadd",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "dev_df_transplant_final.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c5001bf",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "assert dev_df_transplant_final.eval('NIL').sum() >= desired_nil_total\n",
    "assert dev_df_transplant_final.shape[0] >= desired_total"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4784a194",
   "metadata": {},
   "source": [
    "### Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "250846ec",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "mentions_to_test_not_nil = train_df_not_nil.join(\n",
    "    train_df_not_nil_unique_count.loc[sample_transplant_to_test_not_nil].set_index(['wikiId', 'word']),\n",
    "    on=['wikiId', 'word'], how='inner')\n",
    "mentions_to_test_not_nil = mentions_to_test_not_nil.drop(columns=['count'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d716f8b4",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "test_df_transplant_final = pd.concat([test_df_transplant, mentions_to_test_not_nil], ignore_index=True)\n",
    "test_df_transplant_final = test_df_transplant_final.sample(frac=1, random_state=random_state) # randomize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96fab6d8",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "test_df_transplant_final.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5225ca66",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "assert test_df_transplant_final.eval('NIL').sum() >= desired_nil_total\n",
    "assert test_df_transplant_final.shape[0] >= desired_total"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "171239ef",
   "metadata": {},
   "source": [
    "### Remove from train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31516335",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "train_df_transplant_final = train_df_transplant.drop(mentions_to_dev_not_nil.index.union(mentions_to_test_not_nil.index))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcd23c98",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "train_unique_transplant_final = train_df_transplant_final[['word', 'wikiId']].drop_duplicates(keep='first')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2fa2f2d6",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "dev_unique_transplant_final = dev_df_transplant_final[['word', 'wikiId']].drop_duplicates(keep='first')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2639993",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "test_unique_transplant_final = test_df_transplant_final[['word', 'wikiId']].drop_duplicates(keep='first')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09d121c7",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "# ensure no duplicates\n",
    "assert not train_unique_transplant_final.duplicated().any()\n",
    "assert not dev_unique_transplant_final.duplicated().any()\n",
    "assert not test_unique_transplant_final.duplicated().any()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abd693a0",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "# ensure unseen mentions holds\n",
    "assert not pd.concat([train_unique_transplant_final, dev_unique_transplant_final]).duplicated().any()\n",
    "assert not pd.concat([train_unique_transplant_final, test_unique_transplant_final]).duplicated().any()\n",
    "assert not pd.concat([dev_unique_transplant_final, test_unique_transplant_final]).duplicated().any()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1efe0177",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "print('Train', train_df_transplant_final.shape)\n",
    "print('Dev', dev_df_transplant_final.shape)\n",
    "print('Test', test_df_transplant_final.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "365f10bc",
   "metadata": {},
   "source": [
    "# Remove samples from dev test to reach the desired number\n",
    "If required, here it is possible to remove samples from the dev and test sets to reach a precise number of samples."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00cfc994",
   "metadata": {},
   "source": [
    "# Divide train and test in batches\n",
    "Train is divided in batch similarly as it is in the original dataset to avoid a single very big file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8e4a5ff",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "skf = StratifiedKFold(n_splits=n_batch, shuffle=True, random_state=random_state)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57a569ec",
   "metadata": {},
   "source": [
    "# Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aeb90381",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df_transplant_final = train_df_transplant_final.reset_index(drop=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c83f0de",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_batch_indexes = []                                                                                                                                              \n",
    "# the NIL class has freq = 0 so it is fairly distributed among the batches                                                                                            \n",
    "for _, _index in skf.split(np.zeros(train_df_transplant_final.shape[0]), train_df_transplant_final['freq']):                                                          \n",
    "    train_batch_indexes.append(_index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b64f760e",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i, batch in enumerate(train_batch_indexes):                                                                                                                       \n",
    "    train_df_transplant_final.loc[batch, 'batch'] = i                                                                                                                 \n",
    "train_df_transplant_final['batch'] = train_df_transplant_final['batch'].astype(int)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a984d0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df_transplant_final.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "061592ac",
   "metadata": {},
   "source": [
    "# Reset Dev index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a2b22a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "dev_df_transplant_final = dev_df_transplant_final.reset_index(drop=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "924493e8",
   "metadata": {},
   "source": [
    "## Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39c5d8d0",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "test_df_transplant_final = test_df_transplant_final.reset_index(drop=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "362d6a4d",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "test_batch_indexes = []\n",
    "# the NIL class has freq = 0 so it is fairly distributed among the batches\n",
    "for _, _index in skf.split(np.zeros(test_df_transplant_final.shape[0]), test_df_transplant_final['freq']):\n",
    "    test_batch_indexes.append(_index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c67ac96e",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "for i, batch in enumerate(test_batch_indexes):\n",
    "    test_df_transplant_final.loc[batch, 'batch'] = i\n",
    "test_df_transplant_final['batch'] = test_df_transplant_final['batch'].astype(int)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "111e63dd",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "test_df_transplant_final.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ffc5c8c",
   "metadata": {},
   "source": [
    "# Prepare for BLINK"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df6f34a6",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "train_df_transplant_final = train_df_transplant_final.rename(columns = {\n",
    "    'left_context_text': 'context_left',\n",
    "    'word': 'mention',\n",
    "    'right_context_text': 'context_right',\n",
    "    'ex_id': 'query_id',\n",
    "    'url':'label_id',\n",
    "    'wikiId':'Wikipedia_ID',\n",
    "    'wikiurl':'Wikipedia_URL',\n",
    "    'y_title': 'Wikipedia_title'\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "951dd53a",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "dev_df_transplant_final = dev_df_transplant_final.rename(columns = {\n",
    "    'left_context_text': 'context_left',\n",
    "    'word': 'mention',\n",
    "    'right_context_text': 'context_right',\n",
    "    'ex_id': 'query_id',\n",
    "    'url':'label_id',\n",
    "    'wikiId':'Wikipedia_ID',\n",
    "    'wikiurl':'Wikipedia_URL',\n",
    "    'y_title': 'Wikipedia_title'\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75f282d4",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "test_df_transplant_final = test_df_transplant_final.rename(columns = {\n",
    "    'left_context_text': 'context_left',\n",
    "    'word': 'mention',\n",
    "    'right_context_text': 'context_right',\n",
    "    'ex_id': 'query_id',\n",
    "    'url':'label_id',\n",
    "    'wikiId':'Wikipedia_ID',\n",
    "    'wikiurl':'Wikipedia_URL',\n",
    "    'y_title': 'Wikipedia_title'\n",
    "})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3e4d860",
   "metadata": {},
   "source": [
    "# Save to disk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43920e0b",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "train_basedir = os.path.join(outdir, 'train')\n",
    "os.makedirs(train_basedir, exist_ok=True)\n",
    "print('Saving train...')\n",
    "for batch in range(n_batch):\n",
    "    print('Batch {} of {}'.format(batch + 1, n_batch))\n",
    "    train_batch = train_df_transplant_final[train_df_transplant_final['batch'] == batch]\n",
    "    with open(os.path.join(train_basedir, 'train_{}.jsonl'.format(batch)), 'w') as fd:\n",
    "        for i, row in tqdm(train_batch.iterrows(), total=train_batch.shape[0]):\n",
    "            fd.write(row.to_json())\n",
    "            fd.write('\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d60a186",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "dev_basedir = os.path.join(outdir, 'dev')\n",
    "os.makedirs(dev_basedir, exist_ok=True)\n",
    "print('Saving dev...')\n",
    "with open(os.path.join(dev_basedir, 'dev.jsonl'), 'w') as fd:\n",
    "    for i, row in tqdm(dev_df_transplant_final.iterrows(), total=dev_df_transplant_final.shape[0]):\n",
    "        fd.write(row.to_json())\n",
    "        fd.write('\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e30e51ab",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "test_basedir = os.path.join(outdir, 'test')\n",
    "os.makedirs(test_basedir, exist_ok=True)\n",
    "print('Saving test...')\n",
    "for batch in range(n_batch):\n",
    "    print('Batch {} of {}'.format(batch + 1, n_batch))\n",
    "    test_batch = test_df_transplant_final[test_df_transplant_final['batch'] == batch]\n",
    "    with open(os.path.join(test_basedir, 'test_{}.jsonl'.format(batch)), 'w') as fd:\n",
    "        for i, row in tqdm(test_batch.iterrows(), total=test_batch.shape[0]):\n",
    "            fd.write(row.to_json())\n",
    "            fd.write('\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "335ed25b",
   "metadata": {},
   "source": [
    "# Statistics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9b12554",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "report = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd67cafd",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "for batch in trange(n_batch):\n",
    "    test_batch = test_df_transplant_final[test_df_transplant_final['batch'] == batch]\n",
    "    \n",
    "    batch_report = {}\n",
    "    batch_report['batch'] = batch\n",
    "    \n",
    "    # mentions\n",
    "    batch_report['test_mentions'] = test_batch.shape[0]\n",
    "\n",
    "    # entities\n",
    "    batch_report['test_entities'] = test_batch['Wikipedia_ID'].drop_duplicates().shape[0]\n",
    "\n",
    "    # NIL mentions\n",
    "    batch_report['test_nil_mentions'] = test_batch.eval('NIL').sum()\n",
    "    \n",
    "    # not-NIL mentions\n",
    "    batch_report['test_not_nil_mentions'] = batch_report['test_mentions'] - batch_report['test_nil_mentions']\n",
    "\n",
    "    # NIL entities\n",
    "    batch_report['test_nil_entities'] = test_batch.query('NIL')['Wikipedia_ID'].drop_duplicates().shape[0]\n",
    "    \n",
    "    # not-NIL entities\n",
    "    batch_report['test_not_nil_entities'] = batch_report['test_entities'] - batch_report['test_nil_entities']\n",
    "    \n",
    "    # NIL entities present in a previous batch\n",
    "    test_indexes = set(test_batch.query('NIL')['Wikipedia_ID'].drop_duplicates()\n",
    "           ).intersection(\n",
    "            # the set of all the nil entities found in the previous batches\n",
    "            set(test_df_transplant_final.query('NIL and batch < {}'.format(batch))['Wikipedia_ID'].drop_duplicates())\n",
    "        )\n",
    "    batch_report['test_nil_entities_found_in_previous_batch'] = len(test_indexes)\n",
    "    \n",
    "    # NIL mentions present in a previous batch\n",
    "    batch_report['test_nil_mentions_found_in_previous_batch'] = test_batch['Wikipedia_ID'].isin(test_indexes).sum()\n",
    "    \n",
    "    report.append(batch_report)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34660936",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "all_dataset = {}\n",
    "all_dataset['batch'] = 'ALL'\n",
    "\n",
    "# mentions\n",
    "all_dataset['train_mentions'] = train_df_transplant_final.shape[0]\n",
    "all_dataset['dev_mentions'] = dev_df_transplant_final.shape[0]\n",
    "all_dataset['test_mentions'] = test_df_transplant_final.shape[0]\n",
    "\n",
    "# entities\n",
    "all_dataset['train_entities'] = train_df_transplant_final['Wikipedia_ID'].drop_duplicates().shape[0]\n",
    "all_dataset['dev_entities'] = dev_df_transplant_final['Wikipedia_ID'].drop_duplicates().shape[0]\n",
    "all_dataset['test_entities'] = test_df_transplant_final['Wikipedia_ID'].drop_duplicates().shape[0]\n",
    "\n",
    "# NIL mentions\n",
    "all_dataset['train_nil_mentions'] = train_df_transplant_final.eval('NIL').sum()\n",
    "all_dataset['dev_nil_mentions'] = dev_df_transplant_final.eval('NIL').sum()\n",
    "all_dataset['test_nil_mentions'] = test_df_transplant_final.eval('NIL').sum()\n",
    "\n",
    "# not-NIL mentions\n",
    "all_dataset['train_not_nil_mentions'] = all_dataset['train_mentions'] - all_dataset['train_nil_mentions']\n",
    "all_dataset['dev_not_nil_mentions'] = all_dataset['dev_mentions'] - all_dataset['dev_nil_mentions']\n",
    "all_dataset['test_not_nil_mentions'] = all_dataset['test_mentions'] - all_dataset['test_nil_mentions']\n",
    "\n",
    "# NIL entities\n",
    "all_dataset['train_nil_entities'] = train_df_transplant_final.query('NIL')['Wikipedia_ID'].drop_duplicates().shape[0]\n",
    "all_dataset['dev_nil_entities'] = dev_df_transplant_final.query('NIL')['Wikipedia_ID'].drop_duplicates().shape[0]\n",
    "all_dataset['test_nil_entities'] = test_df_transplant_final.query('NIL')['Wikipedia_ID'].drop_duplicates().shape[0]\n",
    "\n",
    "# not-NIL entities\n",
    "all_dataset['train_not_nil_entities'] = all_dataset['train_entities'] - all_dataset['train_nil_entities']\n",
    "all_dataset['dev_not_nil_entities'] = all_dataset['dev_entities'] - all_dataset['dev_nil_entities']\n",
    "all_dataset['test_not_nil_entities'] = all_dataset['test_entities'] - all_dataset['test_nil_entities']\n",
    "\n",
    "# NIL entities present in a previous batch\n",
    "all_dataset['train_nil_entities_found_in_previous_batch'] = 0\n",
    "all_dataset['dev_nil_entities_found_in_previous_batch'] = 0\n",
    "all_dataset['test_nil_entities_found_in_previous_batch'] = 0\n",
    "\n",
    "# NIL mentions present in a previous batch\n",
    "all_dataset['train_nil_mentions_found_in_previous_batch'] = 0\n",
    "all_dataset['dev_nil_mentions_found_in_previous_batch'] = 0\n",
    "all_dataset['test_nil_mentions_found_in_previous_batch'] = 0\n",
    "\n",
    "report.append(all_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fcadbfea",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "original_dataset = {}\n",
    "original_dataset['batch'] = 'ORIGINAL'\n",
    "\n",
    "# mentions\n",
    "original_dataset['train_mentions'] = train_df.shape[0]\n",
    "original_dataset['dev_mentions'] = dev_df.shape[0]\n",
    "original_dataset['test_mentions'] = test_df.shape[0]\n",
    "\n",
    "# entities\n",
    "original_dataset['train_entities'] = train_df['wikiId'].drop_duplicates().shape[0]\n",
    "original_dataset['dev_entities'] = dev_df['wikiId'].drop_duplicates().shape[0]\n",
    "original_dataset['test_entities'] = test_df['wikiId'].drop_duplicates().shape[0]\n",
    "\n",
    "# NIL mentions\n",
    "original_dataset['train_nil_mentions'] =  0\n",
    "original_dataset['dev_nil_mentions'] = 0\n",
    "original_dataset['test_nil_mentions'] = 0\n",
    "\n",
    "# not-NIL mentions\n",
    "original_dataset['train_not_nil_mentions'] = original_dataset['train_mentions'] - original_dataset['train_nil_mentions']\n",
    "original_dataset['dev_not_nil_mentions'] = original_dataset['dev_mentions'] - original_dataset['dev_nil_mentions']\n",
    "original_dataset['test_not_nil_mentions'] = original_dataset['test_mentions'] - original_dataset['test_nil_mentions']\n",
    "\n",
    "# NIL entities\n",
    "original_dataset['train_nil_entities'] = 0\n",
    "original_dataset['dev_nil_entities'] = 0\n",
    "original_dataset['test_nil_entities'] = 0\n",
    "\n",
    "# not-NIL entities\n",
    "original_dataset['train_not_nil_entities'] = original_dataset['train_entities'] - original_dataset['train_nil_entities']\n",
    "original_dataset['dev_not_nil_entities'] = original_dataset['dev_entities'] - original_dataset['dev_nil_entities']\n",
    "original_dataset['test_not_nil_entities'] = original_dataset['test_entities'] - original_dataset['test_nil_entities']\n",
    "\n",
    "# NIL entities present in a previous batch\n",
    "original_dataset['train_nil_entities_found_in_previous_batch'] = 0\n",
    "original_dataset['dev_nil_entities_found_in_previous_batch'] = 0\n",
    "original_dataset['test_nil_entities_found_in_previous_batch'] = 0\n",
    "\n",
    "# NIL mentions present in a previous batch\n",
    "original_dataset['train_nil_mentions_found_in_previous_batch'] = 0\n",
    "original_dataset['dev_nil_mentions_found_in_previous_batch'] = 0\n",
    "original_dataset['test_nil_mentions_found_in_previous_batch'] = 0\n",
    "\n",
    "report.append(original_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "efa7007b",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "report_df = pd.DataFrame(report)\n",
    "report_df"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e9939ff",
   "metadata": {},
   "source": [
    "## Save statistics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60126e89",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "stats_basedir = os.path.join(outdir, 'statistics')\n",
    "os.makedirs(stats_basedir, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7e7345d",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "report_df.to_csv(os.path.join(stats_basedir, 'statistics.csv'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d7068b8",
   "metadata": {},
   "source": [
    " # Remove NIL entities from the KB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "246231bd",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "NIL_entities = []\n",
    "NIL_entities.extend(\n",
    "    train_df_transplant_final.query('NIL')['Wikipedia_ID'].tolist())\n",
    "NIL_entities.extend(\n",
    "    dev_df_transplant_final.query('NIL')['Wikipedia_ID'].tolist())\n",
    "NIL_entities.extend(\n",
    "    test_df_transplant_final.query('NIL')['Wikipedia_ID'].tolist())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14537966",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "NIL_entities = set(NIL_entities) # remove duplicates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f78916f6",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "sql_list = \",\".join(str(i) for i in NIL_entities)\n",
    "indexer = 10 # ensure to put the correct indexer\n",
    "sql_query = 'DELETE FROM entities WHERE indexer = {} AND wikipedia_id in ({});'.format(indexer, sql_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ad59d90",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "with open(os.path.join(outdir, 'delete_nil_entities.sql'), 'w') as fd:\n",
    "    fd.write(sql_query + '\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fcce484a",
   "metadata": {},
   "source": [
    "Now run the sql query in the database (e.g.):\n",
    "```\n",
    "sudo docker-compose exec -T postgres psql -U postgres < /path/to/incremental_dataset/delete_nil_entities.sql\n",
    "```\n",
    "\n",
    "At this point NIL entities are removed from the db and even if retrieved by the faiss indexer they are discarded."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}