Model_v36bert.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c9b4e7e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys, os\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from sklearn.preprocessing import LabelEncoder\n",
    "from sklearn.model_selection import train_test_split\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import torch.nn as nn\n",
    "import math\n",
    "\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n",
    "\n",
    "def seed_everything(seed):\n",
    "    os.environ['PYTHONHASHSEED'] = str(seed)\n",
    "    np.random.seed(seed)\n",
    "    torch.manual_seed(seed)\n",
    "    torch.cuda.manual_seed(seed)\n",
    "    torch.backends.cudnn.deterministic = True\n",
    "    torch.backends.cudnn.benchmark = False\n",
    "    \n",
    "seed_everything(0)\n",
    "\n",
    "MAX_LEN = 64"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b11776d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def adjust_lr(optimizer, epoch):\n",
    "    if epoch >= 2:\n",
    "        lr = 2e-5\n",
    "    else:\n",
    "        lr = 1e-6\n",
    "        \n",
    "    lr *= 4\n",
    "\n",
    "    optimizer.param_groups[0]['lr'] = lr\n",
    "    optimizer.param_groups[1]['lr'] = 100*lr\n",
    "    \n",
    "    return lr\n",
    "\n",
    "\n",
    "def get_optimizer(net):\n",
    "    params = [x[1] for x in filter(lambda kv: \"backbone\" in kv[0], net.named_parameters())]\n",
    "    arc_weight = [x[1] for x in filter(lambda kv: \"backbone\" not in kv[0], net.named_parameters())]\n",
    "\n",
    "    optimizer = torch.optim.Adam([{\"params\": params}, {\"params\": arc_weight}], lr=3e-4, betas=(0.9, 0.999),\n",
    "                                 eps=1e-08)\n",
    "    return optimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "357e6f61",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ArcMarginProduct(nn.Module):\n",
    "    def __init__(self, in_features, out_features, scale=30.0, margin=0.50, easy_margin=False, ls_eps=0.0):\n",
    "        super(ArcMarginProduct, self).__init__()\n",
    "        self.in_features = in_features\n",
    "        self.out_features = out_features\n",
    "        self.scale = scale\n",
    "        \n",
    "        self.ls_eps = ls_eps\n",
    "        self.arc_weight = nn.Parameter(torch.FloatTensor(out_features, in_features))\n",
    "        nn.init.xavier_uniform_(self.arc_weight)\n",
    "\n",
    "        self.easy_margin = easy_margin\n",
    "        self.set_margin(margin)\n",
    "        \n",
    "    def set_margin(self, margin):\n",
    "        self.margin = margin\n",
    "        self.cos_m = math.cos(margin)\n",
    "        self.sin_m = math.sin(margin)\n",
    "        \n",
    "        self.th = nn.Parameter(torch.FloatTensor([math.cos(math.pi - margin)]), requires_grad=False)\n",
    "        self.mm = nn.Parameter(torch.FloatTensor([math.sin(math.pi - margin) * margin]), requires_grad=False)\n",
    "\n",
    "    def forward(self, input, label):\n",
    "        cosine = F.linear(F.normalize(input), F.normalize(self.arc_weight))\n",
    "        sine = torch.sqrt(1.0 - torch.pow(cosine, 2))\n",
    "        phi = cosine * self.cos_m - sine * self.sin_m\n",
    "\n",
    "        if self.easy_margin:\n",
    "            phi = torch.where(cosine > 0, phi, cosine)\n",
    "        else:\n",
    "            phi = torch.where(cosine > self.th, phi, cosine - self.mm)\n",
    "\n",
    "        one_hot = label #torch.zeros(cosine.size(), device='cuda')\n",
    "        #one_hot.scatter_(1, label.view(-1, 1).long(), 1)\n",
    "        if self.ls_eps > 0:\n",
    "            one_hot = (1 - self.ls_eps) * one_hot + self.ls_eps / self.out_features\n",
    "\n",
    "        output = (one_hot * phi) + ((1.0 - one_hot) * cosine)\n",
    "        output *= self.scale\n",
    "\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "68fa5fc3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoConfig, AutoTokenizer, AutoModelForMaskedLM, AutoModel\n",
    "\n",
    "\n",
    "class TrainingModel(nn.Module):\n",
    "    def __init__(self, model, num_classes):\n",
    "        super(TrainingModel, self).__init__()\n",
    "        self.model = model\n",
    "        self.arcface = ArcMarginProduct(768, num_classes, scale=20.0, margin=0.1)\n",
    "\n",
    "    def forward(self, ids, mask, label):\n",
    "        x = self.model(ids, mask)\n",
    "        return self.arcface(x, label)\n",
    "\n",
    "    \n",
    "class VecModel(nn.Module):\n",
    "    def __init__(self, model_name):\n",
    "        super(VecModel, self).__init__()\n",
    "        self.backbone = AutoModel.from_pretrained(model_name)\n",
    "\n",
    "    def forward(self, ids, mask):\n",
    "        out = self.backbone(ids, mask)[0]\n",
    "        out = (out[:, 1:MAX_LEN//2, :]*mask[:, 1:MAX_LEN//2, None]).mean(axis=1)\n",
    "        \n",
    "        return F.normalize(out)\n",
    "\n",
    "    def save(self, path):\n",
    "        torch.save(self.state_dict(), path)\n",
    "\n",
    "    def load(self, path):\n",
    "        self.load_state_dict(torch.load(path, map_location='cpu'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9b2207e2",
   "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>topic_id</th>\n",
       "      <th>content_ids</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>t_00004da3a1b2</td>\n",
       "      <td>c_1108dd0c7a5d</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>t_00004da3a1b2</td>\n",
       "      <td>c_376c5a8eb028</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>t_00004da3a1b2</td>\n",
       "      <td>c_5bc0e1e2cba0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>t_00004da3a1b2</td>\n",
       "      <td>c_76231f9d0b5e</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>t_00068291e9a4</td>\n",
       "      <td>c_639ea2ef9c95</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>279914</th>\n",
       "      <td>t_fff9e5407d13</td>\n",
       "      <td>c_d64037a72376</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>279915</th>\n",
       "      <td>t_fffbe1d5d43c</td>\n",
       "      <td>c_46f852a49c08</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>279916</th>\n",
       "      <td>t_fffbe1d5d43c</td>\n",
       "      <td>c_6659207b25d5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>279917</th>\n",
       "      <td>t_fffe14f1be1e</td>\n",
       "      <td>c_cece166bad6a</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>279918</th>\n",
       "      <td>t_fffe811a6da9</td>\n",
       "      <td>c_92b8fad372ee</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>279919 rows × 2 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "              topic_id     content_ids\n",
       "0       t_00004da3a1b2  c_1108dd0c7a5d\n",
       "1       t_00004da3a1b2  c_376c5a8eb028\n",
       "2       t_00004da3a1b2  c_5bc0e1e2cba0\n",
       "3       t_00004da3a1b2  c_76231f9d0b5e\n",
       "4       t_00068291e9a4  c_639ea2ef9c95\n",
       "...                ...             ...\n",
       "279914  t_fff9e5407d13  c_d64037a72376\n",
       "279915  t_fffbe1d5d43c  c_46f852a49c08\n",
       "279916  t_fffbe1d5d43c  c_6659207b25d5\n",
       "279917  t_fffe14f1be1e  c_cece166bad6a\n",
       "279918  t_fffe811a6da9  c_92b8fad372ee\n",
       "\n",
       "[279919 rows x 2 columns]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "corr_df = pd.read_csv(\"data/correlations.csv\")\n",
    "\n",
    "corr_df[\"content_ids\"] = corr_df[\"content_ids\"].apply(lambda x: x.split())\n",
    "\n",
    "corr_df = corr_df.explode('content_ids').reset_index(drop=True)\n",
    "corr_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "1e3d6939",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(61517, 9)\n",
      "(61517, 10)\n",
      "(61517, 12)\n",
      "(61517, 14)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>id</th>\n",
       "      <th>title</th>\n",
       "      <th>description</th>\n",
       "      <th>channel</th>\n",
       "      <th>category</th>\n",
       "      <th>level</th>\n",
       "      <th>language</th>\n",
       "      <th>parent</th>\n",
       "      <th>has_content</th>\n",
       "      <th>parent_title</th>\n",
       "      <th>grandpa</th>\n",
       "      <th>grandpa_title</th>\n",
       "      <th>ggrandpa</th>\n",
       "      <th>ggrandpa_title</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>t_00004da3a1b2</td>\n",
       "      <td>Откриването на резисторите @ Открития и проект...</td>\n",
       "      <td>Изследване на материали, които предизвикват на...</td>\n",
       "      <td>000cf7</td>\n",
       "      <td>source</td>\n",
       "      <td>4</td>\n",
       "      <td>bg</td>\n",
       "      <td>t_16e29365b50d</td>\n",
       "      <td>True</td>\n",
       "      <td>Открития и проекти</td>\n",
       "      <td>t_c85886762db7</td>\n",
       "      <td>Физика</td>\n",
       "      <td>t_27a0a3c7ee9a</td>\n",
       "      <td>Наука</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>t_00068291e9a4</td>\n",
       "      <td>Entradas e saídas de uma função @ Álgebra: fun...</td>\n",
       "      <td>Entenda um pouco mais sobre funções.</td>\n",
       "      <td>8e286a</td>\n",
       "      <td>source</td>\n",
       "      <td>4</td>\n",
       "      <td>pt</td>\n",
       "      <td>t_d14b6c2a2b70</td>\n",
       "      <td>True</td>\n",
       "      <td>Álgebra: funções</td>\n",
       "      <td>t_461c9a8dd863</td>\n",
       "      <td>9º Ano</td>\n",
       "      <td>t_a6420bcff569</td>\n",
       "      <td>Matemática por ano (Alinhada à BNCC)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>t_00069b63a70a</td>\n",
       "      <td>Transcripts @ Flow Charts: Logical Thinking? @...</td>\n",
       "      <td>NaN</td>\n",
       "      <td>6e3ba4</td>\n",
       "      <td>source</td>\n",
       "      <td>3</td>\n",
       "      <td>en</td>\n",
       "      <td>t_4054df11a74e</td>\n",
       "      <td>True</td>\n",
       "      <td>Flow Charts: Logical Thinking?</td>\n",
       "      <td>t_acbbd893e6af</td>\n",
       "      <td>Engineering</td>\n",
       "      <td>t_9c5f09334d83</td>\n",
       "      <td>MIT Blossoms</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>t_0006d41a73a8</td>\n",
       "      <td>Графики на експоненциални функции (Алгебра 2 н...</td>\n",
       "      <td>Научи повече за графиките на сложните показате...</td>\n",
       "      <td>000cf7</td>\n",
       "      <td>source</td>\n",
       "      <td>4</td>\n",
       "      <td>bg</td>\n",
       "      <td>t_e2452e21d252</td>\n",
       "      <td>True</td>\n",
       "      <td>Показателни и логаритмични функции</td>\n",
       "      <td>t_5f4cc8e02423</td>\n",
       "      <td>Алгебра (цялото съдържание)</td>\n",
       "      <td>t_afd8db5e4f44</td>\n",
       "      <td>Математика</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>t_0008768bdee6</td>\n",
       "      <td>100 સુધીનો સરવાળો @ 100 સુધીના સરવાળા અને બાદબ...</td>\n",
       "      <td>37 અને 49 જેવી બે-અંકની સંખ્યાઓ ઉમેરતા શીખો.</td>\n",
       "      <td>5223e0</td>\n",
       "      <td>supplemental</td>\n",
       "      <td>4</td>\n",
       "      <td>gu</td>\n",
       "      <td>t_0da7a331d666</td>\n",
       "      <td>True</td>\n",
       "      <td>100 સુધીના સરવાળા અને બાદબાકી</td>\n",
       "      <td>t_d20fc01fdf38</td>\n",
       "      <td>પ્રારંભિક ગણિત</td>\n",
       "      <td>t_38bca5075c51</td>\n",
       "      <td>ધોરણ ૩</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "               id                                              title  \\\n",
       "0  t_00004da3a1b2  Откриването на резисторите @ Открития и проект...   \n",
       "1  t_00068291e9a4  Entradas e saídas de uma função @ Álgebra: fun...   \n",
       "2  t_00069b63a70a  Transcripts @ Flow Charts: Logical Thinking? @...   \n",
       "3  t_0006d41a73a8  Графики на експоненциални функции (Алгебра 2 н...   \n",
       "4  t_0008768bdee6  100 સુધીનો સરવાળો @ 100 સુધીના સરવાળા અને બાદબ...   \n",
       "\n",
       "                                         description channel      category  \\\n",
       "0  Изследване на материали, които предизвикват на...  000cf7        source   \n",
       "1               Entenda um pouco mais sobre funções.  8e286a        source   \n",
       "2                                                NaN  6e3ba4        source   \n",
       "3  Научи повече за графиките на сложните показате...  000cf7        source   \n",
       "4       37 અને 49 જેવી બે-અંકની સંખ્યાઓ ઉમેરતા શીખો.  5223e0  supplemental   \n",
       "\n",
       "   level language          parent  has_content  \\\n",
       "0      4       bg  t_16e29365b50d         True   \n",
       "1      4       pt  t_d14b6c2a2b70         True   \n",
       "2      3       en  t_4054df11a74e         True   \n",
       "3      4       bg  t_e2452e21d252         True   \n",
       "4      4       gu  t_0da7a331d666         True   \n",
       "\n",
       "                         parent_title         grandpa  \\\n",
       "0                  Открития и проекти  t_c85886762db7   \n",
       "1                    Álgebra: funções  t_461c9a8dd863   \n",
       "2      Flow Charts: Logical Thinking?  t_acbbd893e6af   \n",
       "3  Показателни и логаритмични функции  t_5f4cc8e02423   \n",
       "4       100 સુધીના સરવાળા અને બાદબાકી  t_d20fc01fdf38   \n",
       "\n",
       "                 grandpa_title        ggrandpa  \\\n",
       "0                       Физика  t_27a0a3c7ee9a   \n",
       "1                       9º Ano  t_a6420bcff569   \n",
       "2                  Engineering  t_9c5f09334d83   \n",
       "3  Алгебра (цялото съдържание)  t_afd8db5e4f44   \n",
       "4               પ્રારંભિક ગણિત  t_38bca5075c51   \n",
       "\n",
       "                         ggrandpa_title  \n",
       "0                                 Наука  \n",
       "1  Matemática por ano (Alinhada à BNCC)  \n",
       "2                          MIT Blossoms  \n",
       "3                            Математика  \n",
       "4                                ધોરણ ૩  "
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "topic_df = pd.read_csv(\"data/topics.csv\")\n",
    "topic_df[\"title\"] = topic_df[\"title\"].fillna(\"\")\n",
    "\n",
    "title_map = topic_df.set_index(\"id\")[\"title\"].to_dict()\n",
    "parent_map = topic_df.set_index(\"id\")[\"parent\"].to_dict()\n",
    "\n",
    "topic_df = topic_df[topic_df[\"has_content\"]].reset_index(drop=True)\n",
    "print(topic_df.shape)\n",
    "\n",
    "topic_df[\"parent_title\"] = topic_df[\"parent\"].apply(lambda x: title_map.get(x, \"\"))\n",
    "print(topic_df.shape)\n",
    "\n",
    "topic_df[\"grandpa\"] = topic_df[\"parent\"].apply(lambda x: parent_map.get(x))\n",
    "topic_df[\"grandpa_title\"] = topic_df[\"grandpa\"].apply(lambda x: title_map.get(x, \"\"))\n",
    "print(topic_df.shape)\n",
    "\n",
    "topic_df[\"ggrandpa\"] = topic_df[\"grandpa\"].apply(lambda x: parent_map.get(x))\n",
    "topic_df[\"ggrandpa_title\"] = topic_df[\"ggrandpa\"].apply(lambda x: title_map.get(x, \"\"))\n",
    "print(topic_df.shape)\n",
    "\n",
    "topic_df[\"title\"] = topic_df[\"title\"] + \" @ \" + topic_df[\"parent_title\"] + \" @ \" + topic_df[\"grandpa_title\"] + \" @ \" + topic_df[\"ggrandpa_title\"]\n",
    "\n",
    "topic_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "9bb133c0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0    20506\n",
       "1    20506\n",
       "2    20505\n",
       "Name: fold, dtype: int64"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from sklearn.model_selection import KFold\n",
    "\n",
    "\n",
    "N_FOLDS = 3\n",
    "\n",
    "splits = KFold(N_FOLDS, shuffle=True, random_state=0).split(topic_df, topic_df[\"has_content\"])\n",
    "\n",
    "topic_df[\"fold\"] = 0\n",
    "\n",
    "for f, (train_ind, val_ind) in enumerate(splits):\n",
    "    topic_df.loc[val_ind, \"fold\"] = f\n",
    "\n",
    "topic_df[\"fold\"].value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "1a458bdd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1    53269\n",
       "0     8248\n",
       "Name: fold, dtype: int64"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "topic_df.loc[(topic_df[\"category\"] == \"source\") | (topic_df[\"fold\"] > 0), \"fold\"] = 1\n",
    "topic_df[\"fold\"].value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "e8477884",
   "metadata": {},
   "outputs": [],
   "source": [
    "le_topic = LabelEncoder()\n",
    "topic_df[\"t\"] = topic_df[\"title\"].fillna(\"\") + \" | \" + topic_df[\"description\"].fillna(\"\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "6590e6a7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((50096, 17), (8248, 17))"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "TEST_FOLD = 0\n",
    "\n",
    "topic_df  = topic_df[topic_df[\"language\"] != \"ar\"]\n",
    "\n",
    "topic_df[\"sw\"] = 1\n",
    "\n",
    "topic_df_test = topic_df[topic_df[\"fold\"] == TEST_FOLD].reset_index(drop=True)\n",
    "topic_df_train = topic_df[topic_df[\"fold\"] != TEST_FOLD].reset_index(drop=True)\n",
    "\n",
    "\n",
    "topic_df_train.shape, topic_df_test.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "522f61d0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(146629, 9)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>id</th>\n",
       "      <th>title</th>\n",
       "      <th>description</th>\n",
       "      <th>kind</th>\n",
       "      <th>text</th>\n",
       "      <th>language</th>\n",
       "      <th>copyright_holder</th>\n",
       "      <th>license</th>\n",
       "      <th>t</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>c_00002381196d</td>\n",
       "      <td>Sumar números de varios dígitos: 48,029+233,930</td>\n",
       "      <td>Suma 48,029+233,930 mediante el algoritmo está...</td>\n",
       "      <td>video</td>\n",
       "      <td>NaN</td>\n",
       "      <td>es</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>Sumar números de varios dígitos: 48,029+233,93...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>c_000087304a9e</td>\n",
       "      <td>Trovare i fattori di un numero</td>\n",
       "      <td>Sal trova i fattori di 120.\\n\\n</td>\n",
       "      <td>video</td>\n",
       "      <td>NaN</td>\n",
       "      <td>it</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>Trovare i fattori di un numero | video | Sal t...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>c_0000ad142ddb</td>\n",
       "      <td>Sumar curvas de demanda</td>\n",
       "      <td>Cómo añadir curvas de demanda\\n\\n</td>\n",
       "      <td>video</td>\n",
       "      <td>NaN</td>\n",
       "      <td>es</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>Sumar curvas de demanda | video | Cómo añadir ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>c_0000c03adc8d</td>\n",
       "      <td>Nado de aproximação</td>\n",
       "      <td>Neste vídeo você vai aprender o nado de aproxi...</td>\n",
       "      <td>document</td>\n",
       "      <td>\\nNado de aproximação\\nSaber nadar nas ondas ...</td>\n",
       "      <td>pt</td>\n",
       "      <td>Sikana Education</td>\n",
       "      <td>CC BY-NC-ND</td>\n",
       "      <td>Nado de aproximação | document | Neste vídeo v...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>c_00016694ea2a</td>\n",
       "      <td>geometry-m3-topic-a-overview.pdf</td>\n",
       "      <td>Estándares Comunes del Estado de Nueva York\\n\\...</td>\n",
       "      <td>document</td>\n",
       "      <td>Estándares Comunes del Estado de Nueva York\\n\\...</td>\n",
       "      <td>es</td>\n",
       "      <td>Engage NY</td>\n",
       "      <td>CC BY-NC-SA</td>\n",
       "      <td>geometry-m3-topic-a-overview.pdf | document | ...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "               id                                             title  \\\n",
       "0  c_00002381196d  Sumar números de varios dígitos: 48,029+233,930    \n",
       "1  c_000087304a9e                    Trovare i fattori di un numero   \n",
       "2  c_0000ad142ddb                           Sumar curvas de demanda   \n",
       "3  c_0000c03adc8d                               Nado de aproximação   \n",
       "4  c_00016694ea2a                  geometry-m3-topic-a-overview.pdf   \n",
       "\n",
       "                                         description      kind  \\\n",
       "0  Suma 48,029+233,930 mediante el algoritmo está...     video   \n",
       "1                    Sal trova i fattori di 120.\\n\\n     video   \n",
       "2                  Cómo añadir curvas de demanda\\n\\n     video   \n",
       "3  Neste vídeo você vai aprender o nado de aproxi...  document   \n",
       "4  Estándares Comunes del Estado de Nueva York\\n\\...  document   \n",
       "\n",
       "                                                text language  \\\n",
       "0                                                NaN       es   \n",
       "1                                                NaN       it   \n",
       "2                                                NaN       es   \n",
       "3  \\nNado de aproximação\\nSaber nadar nas ondas ...       pt   \n",
       "4  Estándares Comunes del Estado de Nueva York\\n\\...       es   \n",
       "\n",
       "   copyright_holder      license  \\\n",
       "0               NaN          NaN   \n",
       "1               NaN          NaN   \n",
       "2               NaN          NaN   \n",
       "3  Sikana Education  CC BY-NC-ND   \n",
       "4         Engage NY  CC BY-NC-SA   \n",
       "\n",
       "                                                   t  \n",
       "0  Sumar números de varios dígitos: 48,029+233,93...  \n",
       "1  Trovare i fattori di un numero | video | Sal t...  \n",
       "2  Sumar curvas de demanda | video | Cómo añadir ...  \n",
       "3  Nado de aproximação | document | Neste vídeo v...  \n",
       "4  geometry-m3-topic-a-overview.pdf | document | ...  "
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "content_df = pd.read_csv(\"data/content.csv\")\n",
    "\n",
    "content_df  = content_df[content_df[\"language\"] != \"ar\"]\n",
    "\n",
    "content_df.loc[content_df[\"title\"] == content_df[\"description\"], \"description\"] = None\n",
    "content_df[\"description\"].fillna(content_df[\"text\"].fillna(\"\").apply(lambda x: x[:256]), inplace=True)\n",
    "content_df[\"t\"] = content_df[\"title\"].fillna(\"\") + \" | \" + content_df[\"kind\"].fillna(\"\") + \" | \" + content_df[\"description\"].fillna(\"\")\n",
    "\n",
    "print(content_df.shape)\n",
    "content_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "bc1b4b03",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(133704, 3)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_contents = corr_df.merge(topic_df_train[[\"id\"]], left_on=\"topic_id\", right_on=\"id\")\n",
    "train_contents[\"sw\"] = 1\n",
    "train_contents[\"sw\"] /= np.sqrt(train_contents.groupby(\"topic_id\")[\"id\"].transform(\"count\"))\n",
    "\n",
    "train_contents = train_contents.groupby(\"content_ids\")[\"sw\"].sum().reset_index()\n",
    "train_contents[\"sw\"] = train_contents[\"sw\"].clip(0.1, 10)\n",
    "\n",
    "train_contents.rename(columns={\"content_ids\": \"id\"}, inplace=True)\n",
    "train_contents[\"is_train\"] = True\n",
    "train_contents[\"sw\"] = train_contents.shape[0]*train_contents[\"sw\"]/train_contents[\"sw\"].sum()\n",
    "\n",
    "train_contents[\"sw\"] = 1\n",
    "train_contents.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "d3745f50",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.9115250052854483"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "content_df = content_df.merge(train_contents, on=\"id\", how=\"left\")\n",
    "content_df[\"is_train\"].fillna(False, inplace=True)\n",
    "content_df[\"sw\"].fillna(1.0, inplace=True)\n",
    "\n",
    "content_df[\"is_train\"].mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "fe22c4dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "content_df[\"label\"] = 0\n",
    "content_df.loc[content_df[\"is_train\"], \"label\"] = le_topic.fit_transform(content_df.loc[content_df[\"is_train\"], \"t\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "2be1c7b1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>id</th>\n",
       "      <th>title</th>\n",
       "      <th>description</th>\n",
       "      <th>channel</th>\n",
       "      <th>category</th>\n",
       "      <th>level</th>\n",
       "      <th>language</th>\n",
       "      <th>parent</th>\n",
       "      <th>has_content</th>\n",
       "      <th>parent_title</th>\n",
       "      <th>grandpa</th>\n",
       "      <th>grandpa_title</th>\n",
       "      <th>ggrandpa</th>\n",
       "      <th>ggrandpa_title</th>\n",
       "      <th>fold</th>\n",
       "      <th>t</th>\n",
       "      <th>sw</th>\n",
       "      <th>label</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>t_00004da3a1b2</td>\n",
       "      <td>Откриването на резисторите @ Открития и проект...</td>\n",
       "      <td>Изследване на материали, които предизвикват на...</td>\n",
       "      <td>000cf7</td>\n",
       "      <td>source</td>\n",
       "      <td>4</td>\n",
       "      <td>bg</td>\n",
       "      <td>t_16e29365b50d</td>\n",
       "      <td>True</td>\n",
       "      <td>Открития и проекти</td>\n",
       "      <td>t_c85886762db7</td>\n",
       "      <td>Физика</td>\n",
       "      <td>t_27a0a3c7ee9a</td>\n",
       "      <td>Наука</td>\n",
       "      <td>1</td>\n",
       "      <td>Откриването на резисторите @ Открития и проект...</td>\n",
       "      <td>1</td>\n",
       "      <td>[111544, 113249, 114120, 114624]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>t_00068291e9a4</td>\n",
       "      <td>Entradas e saídas de uma função @ Álgebra: fun...</td>\n",
       "      <td>Entenda um pouco mais sobre funções.</td>\n",
       "      <td>8e286a</td>\n",
       "      <td>source</td>\n",
       "      <td>4</td>\n",
       "      <td>pt</td>\n",
       "      <td>t_d14b6c2a2b70</td>\n",
       "      <td>True</td>\n",
       "      <td>Álgebra: funções</td>\n",
       "      <td>t_461c9a8dd863</td>\n",
       "      <td>9º Ano</td>\n",
       "      <td>t_a6420bcff569</td>\n",
       "      <td>Matemática por ano (Alinhada à BNCC)</td>\n",
       "      <td>1</td>\n",
       "      <td>Entradas e saídas de uma função @ Álgebra: fun...</td>\n",
       "      <td>1</td>\n",
       "      <td>[89691, 35260, 35259, 89692]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>t_00069b63a70a</td>\n",
       "      <td>Transcripts @ Flow Charts: Logical Thinking? @...</td>\n",
       "      <td>NaN</td>\n",
       "      <td>6e3ba4</td>\n",
       "      <td>source</td>\n",
       "      <td>3</td>\n",
       "      <td>en</td>\n",
       "      <td>t_4054df11a74e</td>\n",
       "      <td>True</td>\n",
       "      <td>Flow Charts: Logical Thinking?</td>\n",
       "      <td>t_acbbd893e6af</td>\n",
       "      <td>Engineering</td>\n",
       "      <td>t_9c5f09334d83</td>\n",
       "      <td>MIT Blossoms</td>\n",
       "      <td>1</td>\n",
       "      <td>Transcripts @ Flow Charts: Logical Thinking? @...</td>\n",
       "      <td>1</td>\n",
       "      <td>[48571]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>t_0006d41a73a8</td>\n",
       "      <td>Графики на експоненциални функции (Алгебра 2 н...</td>\n",
       "      <td>Научи повече за графиките на сложните показате...</td>\n",
       "      <td>000cf7</td>\n",
       "      <td>source</td>\n",
       "      <td>4</td>\n",
       "      <td>bg</td>\n",
       "      <td>t_e2452e21d252</td>\n",
       "      <td>True</td>\n",
       "      <td>Показателни и логаритмични функции</td>\n",
       "      <td>t_5f4cc8e02423</td>\n",
       "      <td>Алгебра (цялото съдържание)</td>\n",
       "      <td>t_afd8db5e4f44</td>\n",
       "      <td>Математика</td>\n",
       "      <td>1</td>\n",
       "      <td>Графики на експоненциални функции (Алгебра 2 н...</td>\n",
       "      <td>1</td>\n",
       "      <td>[111488, 116870, 111475, 116276, 116277]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>t_0008768bdee6</td>\n",
       "      <td>100 સુધીનો સરવાળો @ 100 સુધીના સરવાળા અને બાદબ...</td>\n",
       "      <td>37 અને 49 જેવી બે-અંકની સંખ્યાઓ ઉમેરતા શીખો.</td>\n",
       "      <td>5223e0</td>\n",
       "      <td>supplemental</td>\n",
       "      <td>4</td>\n",
       "      <td>gu</td>\n",
       "      <td>t_0da7a331d666</td>\n",
       "      <td>True</td>\n",
       "      <td>100 સુધીના સરવાળા અને બાદબાકી</td>\n",
       "      <td>t_d20fc01fdf38</td>\n",
       "      <td>પ્રારંભિક ગણિત</td>\n",
       "      <td>t_38bca5075c51</td>\n",
       "      <td>ધોરણ ૩</td>\n",
       "      <td>1</td>\n",
       "      <td>100 સુધીનો સરવાળો @ 100 સુધીના સરવાળા અને બાદબ...</td>\n",
       "      <td>1</td>\n",
       "      <td>[126376, 126369, 126548]</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",
       "      <td>...</td>\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",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>50091</th>\n",
       "      <td>t_fff7f2dd208b</td>\n",
       "      <td>Fatoração de polinômios encontrando fatores co...</td>\n",
       "      <td>Aprenda a fatorar utilizando fatores comuns.</td>\n",
       "      <td>8e286a</td>\n",
       "      <td>source</td>\n",
       "      <td>4</td>\n",
       "      <td>pt</td>\n",
       "      <td>t_28dfc9e80110</td>\n",
       "      <td>True</td>\n",
       "      <td>Álgebra: expressões algébricas</td>\n",
       "      <td>t_461c9a8dd863</td>\n",
       "      <td>9º Ano</td>\n",
       "      <td>t_a6420bcff569</td>\n",
       "      <td>Matemática por ano (Alinhada à BNCC)</td>\n",
       "      <td>1</td>\n",
       "      <td>Fatoração de polinômios encontrando fatores co...</td>\n",
       "      <td>1</td>\n",
       "      <td>[47584, 47552, 34508, 34509, 34510, 47572]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>50092</th>\n",
       "      <td>t_fff830472691</td>\n",
       "      <td>Scalar Projections @ Vector Analysis @ Analysi...</td>\n",
       "      <td>NaN</td>\n",
       "      <td>fef095</td>\n",
       "      <td>source</td>\n",
       "      <td>4</td>\n",
       "      <td>en</td>\n",
       "      <td>t_c75d6acecf78</td>\n",
       "      <td>True</td>\n",
       "      <td>Vector Analysis</td>\n",
       "      <td>t_de58fdbcad60</td>\n",
       "      <td>Analysis</td>\n",
       "      <td>t_9d8388491acd</td>\n",
       "      <td>Math</td>\n",
       "      <td>1</td>\n",
       "      <td>Scalar Projections @ Vector Analysis @ Analysi...</td>\n",
       "      <td>1</td>\n",
       "      <td>[91761, 91762]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>50093</th>\n",
       "      <td>t_fff9e5407d13</td>\n",
       "      <td>NA_U06 - El periódico @ Lengua española @ PF (...</td>\n",
       "      <td>NaN</td>\n",
       "      <td>71fd51</td>\n",
       "      <td>supplemental</td>\n",
       "      <td>2</td>\n",
       "      <td>es</td>\n",
       "      <td>t_5bd8f6ae9f7d</td>\n",
       "      <td>True</td>\n",
       "      <td>Lengua española</td>\n",
       "      <td>t_3eca90fe0fe1</td>\n",
       "      <td>PF (Español)</td>\n",
       "      <td>NaN</td>\n",
       "      <td></td>\n",
       "      <td>1</td>\n",
       "      <td>NA_U06 - El periódico @ Lengua española @ PF (...</td>\n",
       "      <td>1</td>\n",
       "      <td>[22630, 41353, 82122, 63443, 41684, 62933, 640...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>50094</th>\n",
       "      <td>t_fffbe1d5d43c</td>\n",
       "      <td>Inscribed shapes problem solving @ Miduara @ C...</td>\n",
       "      <td>Use properties of inscribed angles to prove pr...</td>\n",
       "      <td>0c929f</td>\n",
       "      <td>source</td>\n",
       "      <td>4</td>\n",
       "      <td>sw</td>\n",
       "      <td>t_50145b9bab3f</td>\n",
       "      <td>True</td>\n",
       "      <td>Miduara</td>\n",
       "      <td>t_c520677cef1d</td>\n",
       "      <td>Class 9 (India)</td>\n",
       "      <td>t_4853606228df</td>\n",
       "      <td>Hisabati</td>\n",
       "      <td>1</td>\n",
       "      <td>Inscribed shapes problem solving @ Miduara @ C...</td>\n",
       "      <td>1</td>\n",
       "      <td>[18867, 85597]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>50095</th>\n",
       "      <td>t_fffe14f1be1e</td>\n",
       "      <td>Lección 7 @ Unidad 4 @ Español Actividades 1ro...</td>\n",
       "      <td>NaN</td>\n",
       "      <td>6e90a7</td>\n",
       "      <td>aligned</td>\n",
       "      <td>6</td>\n",
       "      <td>es</td>\n",
       "      <td>t_d448c707984d</td>\n",
       "      <td>True</td>\n",
       "      <td>Unidad 4</td>\n",
       "      <td>t_8c734f803ed6</td>\n",
       "      <td>Español Actividades 1ro</td>\n",
       "      <td>t_7eb0dc459971</td>\n",
       "      <td>01 Primero</td>\n",
       "      <td>1</td>\n",
       "      <td>Lección 7 @ Unidad 4 @ Español Actividades 1ro...</td>\n",
       "      <td>1</td>\n",
       "      <td>[61581]</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>50096 rows × 18 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                   id                                              title  \\\n",
       "0      t_00004da3a1b2  Откриването на резисторите @ Открития и проект...   \n",
       "1      t_00068291e9a4  Entradas e saídas de uma função @ Álgebra: fun...   \n",
       "2      t_00069b63a70a  Transcripts @ Flow Charts: Logical Thinking? @...   \n",
       "3      t_0006d41a73a8  Графики на експоненциални функции (Алгебра 2 н...   \n",
       "4      t_0008768bdee6  100 સુધીનો સરવાળો @ 100 સુધીના સરવાળા અને બાદબ...   \n",
       "...               ...                                                ...   \n",
       "50091  t_fff7f2dd208b  Fatoração de polinômios encontrando fatores co...   \n",
       "50092  t_fff830472691  Scalar Projections @ Vector Analysis @ Analysi...   \n",
       "50093  t_fff9e5407d13  NA_U06 - El periódico @ Lengua española @ PF (...   \n",
       "50094  t_fffbe1d5d43c  Inscribed shapes problem solving @ Miduara @ C...   \n",
       "50095  t_fffe14f1be1e  Lección 7 @ Unidad 4 @ Español Actividades 1ro...   \n",
       "\n",
       "                                             description channel  \\\n",
       "0      Изследване на материали, които предизвикват на...  000cf7   \n",
       "1                   Entenda um pouco mais sobre funções.  8e286a   \n",
       "2                                                    NaN  6e3ba4   \n",
       "3      Научи повече за графиките на сложните показате...  000cf7   \n",
       "4           37 અને 49 જેવી બે-અંકની સંખ્યાઓ ઉમેરતા શીખો.  5223e0   \n",
       "...                                                  ...     ...   \n",
       "50091       Aprenda a fatorar utilizando fatores comuns.  8e286a   \n",
       "50092                                                NaN  fef095   \n",
       "50093                                                NaN  71fd51   \n",
       "50094  Use properties of inscribed angles to prove pr...  0c929f   \n",
       "50095                                                NaN  6e90a7   \n",
       "\n",
       "           category  level language          parent  has_content  \\\n",
       "0            source      4       bg  t_16e29365b50d         True   \n",
       "1            source      4       pt  t_d14b6c2a2b70         True   \n",
       "2            source      3       en  t_4054df11a74e         True   \n",
       "3            source      4       bg  t_e2452e21d252         True   \n",
       "4      supplemental      4       gu  t_0da7a331d666         True   \n",
       "...             ...    ...      ...             ...          ...   \n",
       "50091        source      4       pt  t_28dfc9e80110         True   \n",
       "50092        source      4       en  t_c75d6acecf78         True   \n",
       "50093  supplemental      2       es  t_5bd8f6ae9f7d         True   \n",
       "50094        source      4       sw  t_50145b9bab3f         True   \n",
       "50095       aligned      6       es  t_d448c707984d         True   \n",
       "\n",
       "                             parent_title         grandpa  \\\n",
       "0                      Открития и проекти  t_c85886762db7   \n",
       "1                        Álgebra: funções  t_461c9a8dd863   \n",
       "2          Flow Charts: Logical Thinking?  t_acbbd893e6af   \n",
       "3      Показателни и логаритмични функции  t_5f4cc8e02423   \n",
       "4           100 સુધીના સરવાળા અને બાદબાકી  t_d20fc01fdf38   \n",
       "...                                   ...             ...   \n",
       "50091      Álgebra: expressões algébricas  t_461c9a8dd863   \n",
       "50092                     Vector Analysis  t_de58fdbcad60   \n",
       "50093                     Lengua española  t_3eca90fe0fe1   \n",
       "50094                             Miduara  t_c520677cef1d   \n",
       "50095                            Unidad 4  t_8c734f803ed6   \n",
       "\n",
       "                     grandpa_title        ggrandpa  \\\n",
       "0                           Физика  t_27a0a3c7ee9a   \n",
       "1                           9º Ano  t_a6420bcff569   \n",
       "2                      Engineering  t_9c5f09334d83   \n",
       "3      Алгебра (цялото съдържание)  t_afd8db5e4f44   \n",
       "4                   પ્રારંભિક ગણિત  t_38bca5075c51   \n",
       "...                            ...             ...   \n",
       "50091                       9º Ano  t_a6420bcff569   \n",
       "50092                     Analysis  t_9d8388491acd   \n",
       "50093                 PF (Español)             NaN   \n",
       "50094              Class 9 (India)  t_4853606228df   \n",
       "50095      Español Actividades 1ro  t_7eb0dc459971   \n",
       "\n",
       "                             ggrandpa_title  fold  \\\n",
       "0                                     Наука     1   \n",
       "1      Matemática por ano (Alinhada à BNCC)     1   \n",
       "2                              MIT Blossoms     1   \n",
       "3                                Математика     1   \n",
       "4                                    ધોરણ ૩     1   \n",
       "...                                     ...   ...   \n",
       "50091  Matemática por ano (Alinhada à BNCC)     1   \n",
       "50092                                  Math     1   \n",
       "50093                                           1   \n",
       "50094                              Hisabati     1   \n",
       "50095                            01 Primero     1   \n",
       "\n",
       "                                                       t  sw  \\\n",
       "0      Откриването на резисторите @ Открития и проект...   1   \n",
       "1      Entradas e saídas de uma função @ Álgebra: fun...   1   \n",
       "2      Transcripts @ Flow Charts: Logical Thinking? @...   1   \n",
       "3      Графики на експоненциални функции (Алгебра 2 н...   1   \n",
       "4      100 સુધીનો સરવાળો @ 100 સુધીના સરવાળા અને બાદબ...   1   \n",
       "...                                                  ...  ..   \n",
       "50091  Fatoração de polinômios encontrando fatores co...   1   \n",
       "50092  Scalar Projections @ Vector Analysis @ Analysi...   1   \n",
       "50093  NA_U06 - El periódico @ Lengua española @ PF (...   1   \n",
       "50094  Inscribed shapes problem solving @ Miduara @ C...   1   \n",
       "50095  Lección 7 @ Unidad 4 @ Español Actividades 1ro...   1   \n",
       "\n",
       "                                                   label  \n",
       "0                       [111544, 113249, 114120, 114624]  \n",
       "1                           [89691, 35260, 35259, 89692]  \n",
       "2                                                [48571]  \n",
       "3               [111488, 116870, 111475, 116276, 116277]  \n",
       "4                               [126376, 126369, 126548]  \n",
       "...                                                  ...  \n",
       "50091         [47584, 47552, 34508, 34509, 34510, 47572]  \n",
       "50092                                     [91761, 91762]  \n",
       "50093  [22630, 41353, 82122, 63443, 41684, 62933, 640...  \n",
       "50094                                     [18867, 85597]  \n",
       "50095                                            [61581]  \n",
       "\n",
       "[50096 rows x 18 columns]"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#corr_df[\"label\"] = le_topic.transform(corr_df[\"topic_id\"])\n",
    "corr_df = corr_df.merge(content_df[content_df[\"is_train\"]][[\"id\", \"label\"]].rename(columns={\"id\": \"content_ids\"}), on=\"content_ids\")\n",
    "\n",
    "\n",
    "content_labels = corr_df.groupby(\"topic_id\").agg({\"label\": list}).reset_index()\n",
    "content_labels = content_labels.rename(columns={\"topic_id\": \"id\"})\n",
    "content_labels[\"label\"] = content_labels[\"label\"].apply(lambda x: list(set(x)))\n",
    "\n",
    "topic_df_train = topic_df_train.merge(content_labels, on=\"id\")\n",
    "\n",
    "topic_df_train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "c3a040ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "content_df[\"label\"] = content_df[\"label\"].apply(lambda x: [x])\n",
    "\n",
    "topic_df_test[\"label\"] = 0\n",
    "topic_df_test[\"label\"] = topic_df_test[\"label\"].apply(lambda x: [x])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "ec897042",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(131987, 133656)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(le_topic.classes_), content_df[\"is_train\"].sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "d2996f47",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ebc77a5848d44670bdca99964621afbe",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading:   0%|          | 0.00/28.0 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "42853d3c921a47f8af649f876e2c60e4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading:   0%|          | 0.00/625 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e61ffa9f978e42b7b935f61af8bc16a2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading:   0%|          | 0.00/851k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "565c0fa88b4444df9ba5017a1949d676",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading:   0%|          | 0.00/1.64M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "MODEL_NAME = \"bert-base-multilingual-uncased\"\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, add_prefix_space=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "90020221",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "313489877e304278867074460c035814",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading:   0%|          | 0.00/641M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at bert-base-multilingual-uncased were not used when initializing BertModel: ['cls.predictions.bias', 'cls.predictions.transform.dense.weight', 'cls.predictions.transform.LayerNorm.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.seq_relationship.weight', 'cls.seq_relationship.bias', 'cls.predictions.decoder.weight', 'cls.predictions.transform.dense.bias']\n",
      "- This IS expected if you are initializing BertModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    }
   ],
   "source": [
    "NW = 8\n",
    "BS = 64\n",
    "\n",
    "class LECRDataset(Dataset):\n",
    "    def __init__(self, df, tokenizer, max_len=512, aug=False):\n",
    "        self.df = df.reset_index(drop=True)\n",
    "        self.tokenizer = tokenizer\n",
    "        self.max_len = max_len\n",
    "        self.mask_token = self.tokenizer.convert_tokens_to_ids(self.tokenizer.mask_token)\n",
    "        self.aug = aug\n",
    "        self.num_classes = len(le_topic.classes_)\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.df.shape[0]\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        row = self.df.iloc[index]\n",
    "        encoding = self.tokenizer(str(row.t), return_offsets_mapping=True,\n",
    "                                  padding='max_length', truncation=True, max_length=self.max_len)\n",
    "\n",
    "        encoding = {key: torch.as_tensor(val) for key, val in encoding.items()}\n",
    "\n",
    "        if self.aug:\n",
    "            ix = torch.rand(size=(self.max_len,)) < 0.15\n",
    "            encoding[\"input_ids\"][ix] = self.mask_token\n",
    "            \n",
    "        target = torch.zeros(self.num_classes, dtype=torch.float32)\n",
    "        for l in row.label:\n",
    "            target[l] = 1/len(row.label)\n",
    "        \n",
    "        return encoding, target, torch.FloatTensor([row.sw])\n",
    "    \n",
    "    \n",
    "\n",
    "\n",
    "\n",
    "\n",
    "vec_model = VecModel(MODEL_NAME)\n",
    "model = TrainingModel(vec_model, num_classes=len(le_topic.classes_))\n",
    "model = model.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "ba250761",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.neighbors import NearestNeighbors\n",
    "\n",
    "\n",
    "def to_vec(model, ds):\n",
    "    val_loader = DataLoader(ds, batch_size=BS, shuffle=False, num_workers=NW,\n",
    "                             pin_memory=False, drop_last=False)\n",
    "    \n",
    "    model.eval()\n",
    "\n",
    "    tbar = tqdm(val_loader, file=sys.stdout)\n",
    "\n",
    "    vectors = []\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for idx, (data, target, _) in enumerate(tbar):\n",
    "            ids, mask, target = data[\"input_ids\"].cuda(), data[\"attention_mask\"].cuda(), target.cuda()\n",
    "\n",
    "            vec = model.model(ids, mask)\n",
    "            vectors.append(vec.detach().cpu().numpy())\n",
    "\n",
    "    V = np.concatenate(vectors)\n",
    "    \n",
    "    return V\n",
    "\n",
    "\n",
    "\n",
    "def get_matches(V_topic, V_content, topic_ids, content_ids, n_neighbors=5):\n",
    "    \n",
    "    neighbors_model = NearestNeighbors(n_neighbors=n_neighbors, metric='cosine', n_jobs=-1)\n",
    "    neighbors_model.fit(V_content)\n",
    "    dists, indices = neighbors_model.kneighbors(V_topic)\n",
    "    \n",
    "    res_df = pd.DataFrame({\"topic_id\": np.repeat(topic_ids, n_neighbors),\n",
    "                           'content_id': content_ids[indices.ravel()],\n",
    "                           'vec_dist': dists.ravel()\n",
    "                          })\n",
    "    return res_df\n",
    "\n",
    "\n",
    "def f2(recall, precision):\n",
    "    return 5*precision*recall/(4*precision + recall)\n",
    "\n",
    "def metric(pred_df):\n",
    "    scores = np.zeros(pred_df.shape[0])\n",
    "\n",
    "    for i, (gt, pred) in enumerate(zip(pred_df[\"content_ids\"].values, pred_df[\"pred\"].values)):\n",
    "        hits = 0\n",
    "\n",
    "        gt, pred = gt.split(), pred.split()\n",
    "\n",
    "        for l in gt:\n",
    "            if l in pred:\n",
    "                hits += 1\n",
    "\n",
    "        if hits > 0:           \n",
    "            recall = hits / len(gt)\n",
    "            precision = hits / len(pred)\n",
    "\n",
    "            scores[i] = f2(recall, precision)\n",
    "\n",
    "    return scores.mean()\n",
    "    \n",
    "\n",
    "def evaluate(topic_df, content_df):    \n",
    "    topic_ds = LECRDataset(topic_df, tokenizer, MAX_LEN, aug=False)\n",
    "    content_ds = LECRDataset(content_df, tokenizer, MAX_LEN, aug=False)\n",
    "\n",
    "\n",
    "    V_topic = to_vec(model, topic_ds)\n",
    "    V_content = to_vec(model, content_ds)\n",
    "\n",
    "    res_df = get_matches(V_topic, V_content, topic_df[\"id\"].values, content_df[\"id\"].values)\n",
    "\n",
    "    res_df[\"rank\"] = res_df.groupby(\"topic_id\")[\"vec_dist\"].rank(method=\"first\", ascending=True)\n",
    "\n",
    "    pred_df = res_df.copy() #res_df[(res_df[\"rank\"] == 1) | (res_df[\"vec_dist\"] < 0.1)]\n",
    "    pred_df = pred_df.groupby(\"topic_id\")[\"content_id\"].apply(lambda x: \" \".join(list(x)))\n",
    "    pred_df = pred_df.reset_index().rename(columns={\"content_id\": 'pred'})\n",
    "    pred_df = pred_df.merge(pd.read_csv(\"data/correlations.csv\"), on=\"topic_id\")\n",
    "\n",
    "    return metric(pred_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "71338fef",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_33343/362973530.py:2: 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",
      "  cls_df[\"ix\"] = le_topic.transform(cls_df[\"t\"])\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████| 2063/2063 [01:27<00:00, 23.64it/s]\n"
     ]
    }
   ],
   "source": [
    "cls_df = content_df[content_df[\"is_train\"]]\n",
    "cls_df[\"ix\"] = le_topic.transform(cls_df[\"t\"])\n",
    "cls_df = cls_df[[\"ix\", \"t\"]].drop_duplicates().sort_values(\"ix\").reset_index(drop=True)\n",
    "\n",
    "cls_df[\"label\"] = cls_df[\"ix\"].apply(lambda x: [x])\n",
    "cls_df[\"sw\"] = 1\n",
    "\n",
    "ds = LECRDataset(cls_df, tokenizer, MAX_LEN, aug=False)\n",
    "\n",
    "V_cls = to_vec(model, ds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "d2cd6d02",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TrainingModel(\n",
       "  (model): VecModel(\n",
       "    (backbone): BertModel(\n",
       "      (embeddings): BertEmbeddings(\n",
       "        (word_embeddings): Embedding(105879, 768, padding_idx=0)\n",
       "        (position_embeddings): Embedding(512, 768)\n",
       "        (token_type_embeddings): Embedding(2, 768)\n",
       "        (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "        (dropout): Dropout(p=0.1, inplace=False)\n",
       "      )\n",
       "      (encoder): BertEncoder(\n",
       "        (layer): ModuleList(\n",
       "          (0): BertLayer(\n",
       "            (attention): BertAttention(\n",
       "              (self): BertSelfAttention(\n",
       "                (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (output): BertSelfOutput(\n",
       "                (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "                (dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (intermediate): BertIntermediate(\n",
       "              (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            )\n",
       "            (output): BertOutput(\n",
       "              (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (1): BertLayer(\n",
       "            (attention): BertAttention(\n",
       "              (self): BertSelfAttention(\n",
       "                (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (output): BertSelfOutput(\n",
       "                (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "                (dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (intermediate): BertIntermediate(\n",
       "              (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            )\n",
       "            (output): BertOutput(\n",
       "              (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (2): BertLayer(\n",
       "            (attention): BertAttention(\n",
       "              (self): BertSelfAttention(\n",
       "                (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (output): BertSelfOutput(\n",
       "                (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "                (dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (intermediate): BertIntermediate(\n",
       "              (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            )\n",
       "            (output): BertOutput(\n",
       "              (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (3): BertLayer(\n",
       "            (attention): BertAttention(\n",
       "              (self): BertSelfAttention(\n",
       "                (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (output): BertSelfOutput(\n",
       "                (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "                (dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (intermediate): BertIntermediate(\n",
       "              (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            )\n",
       "            (output): BertOutput(\n",
       "              (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (4): BertLayer(\n",
       "            (attention): BertAttention(\n",
       "              (self): BertSelfAttention(\n",
       "                (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (output): BertSelfOutput(\n",
       "                (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "                (dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (intermediate): BertIntermediate(\n",
       "              (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            )\n",
       "            (output): BertOutput(\n",
       "              (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (5): BertLayer(\n",
       "            (attention): BertAttention(\n",
       "              (self): BertSelfAttention(\n",
       "                (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (output): BertSelfOutput(\n",
       "                (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "                (dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (intermediate): BertIntermediate(\n",
       "              (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            )\n",
       "            (output): BertOutput(\n",
       "              (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (6): BertLayer(\n",
       "            (attention): BertAttention(\n",
       "              (self): BertSelfAttention(\n",
       "                (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (output): BertSelfOutput(\n",
       "                (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "                (dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (intermediate): BertIntermediate(\n",
       "              (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            )\n",
       "            (output): BertOutput(\n",
       "              (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (7): BertLayer(\n",
       "            (attention): BertAttention(\n",
       "              (self): BertSelfAttention(\n",
       "                (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (output): BertSelfOutput(\n",
       "                (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "                (dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (intermediate): BertIntermediate(\n",
       "              (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            )\n",
       "            (output): BertOutput(\n",
       "              (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (8): BertLayer(\n",
       "            (attention): BertAttention(\n",
       "              (self): BertSelfAttention(\n",
       "                (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (output): BertSelfOutput(\n",
       "                (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "                (dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (intermediate): BertIntermediate(\n",
       "              (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            )\n",
       "            (output): BertOutput(\n",
       "              (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (9): BertLayer(\n",
       "            (attention): BertAttention(\n",
       "              (self): BertSelfAttention(\n",
       "                (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (output): BertSelfOutput(\n",
       "                (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "                (dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (intermediate): BertIntermediate(\n",
       "              (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            )\n",
       "            (output): BertOutput(\n",
       "              (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (10): BertLayer(\n",
       "            (attention): BertAttention(\n",
       "              (self): BertSelfAttention(\n",
       "                (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (output): BertSelfOutput(\n",
       "                (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "                (dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (intermediate): BertIntermediate(\n",
       "              (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            )\n",
       "            (output): BertOutput(\n",
       "              (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (11): BertLayer(\n",
       "            (attention): BertAttention(\n",
       "              (self): BertSelfAttention(\n",
       "                (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (output): BertSelfOutput(\n",
       "                (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "                (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "                (dropout): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "            )\n",
       "            (intermediate): BertIntermediate(\n",
       "              (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            )\n",
       "            (output): BertOutput(\n",
       "              (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (pooler): BertPooler(\n",
       "        (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "        (activation): Tanh()\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (arcface): ArcMarginProduct()\n",
       ")"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.arcface.arc_weight = nn.Parameter(torch.FloatTensor(V_cls))\n",
    "model.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "7ddb6091",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_33343/2716447418.py:7: FutureWarning: The frame.append method is deprecated and will be removed from pandas in a future version. Use pandas.concat instead.\n",
      "  train_df = topic_df_train.append(content_df[content_df[\"is_train\"]]).reset_index(drop=True)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1 Loss: 5.5447 Acc: 0.05 lr: 4e-06: 100%|█| 2871/2871 [06:26<00:00,  7.42i\n",
      "100%|█████████████████████████████████████████| 129/129 [00:05<00:00, 21.93it/s]\n",
      "100%|███████████████████████████████████████| 2292/2292 [01:37<00:00, 23.50it/s]\n",
      "Val score: 0.23389063847525834\n",
      "\n",
      "Epoch 2 Loss: 3.6191 Acc: 0.15 lr: 4e-06: 100%|█| 2871/2871 [06:17<00:00,  7.61i\n",
      "100%|█████████████████████████████████████████| 129/129 [00:05<00:00, 22.10it/s]\n",
      "100%|███████████████████████████████████████| 2292/2292 [01:36<00:00, 23.68it/s]\n",
      "Val score: 0.2589142255279675\n",
      "\n",
      "Epoch 3 Loss: 6.9317 Acc: 0.06 lr: 8e-05: 100%|█| 2871/2871 [06:17<00:00,  7.61i\n",
      "100%|█████████████████████████████████████████| 129/129 [00:05<00:00, 22.08it/s]\n",
      "100%|███████████████████████████████████████| 2292/2292 [01:36<00:00, 23.66it/s]\n",
      "Val score: 0.38302543724467397\n",
      "\n",
      "Epoch 4 Loss: 4.3301 Acc: 0.13 lr: 8e-05: 100%|█| 2871/2871 [06:17<00:00,  7.61i\n",
      "100%|█████████████████████████████████████████| 129/129 [00:05<00:00, 22.04it/s]\n",
      "100%|███████████████████████████████████████| 2292/2292 [01:36<00:00, 23.64it/s]\n",
      "Val score: 0.417302609818477\n",
      "\n",
      "Epoch 5 Loss: 2.7004 Acc: 0.28 lr: 8e-05: 100%|█| 2871/2871 [06:17<00:00,  7.60i\n",
      "100%|█████████████████████████████████████████| 129/129 [00:05<00:00, 22.03it/s]\n",
      "100%|███████████████████████████████████████| 2292/2292 [01:36<00:00, 23.64it/s]\n",
      "Val score: 0.44034597391716196\n",
      "\n",
      "Epoch 6 Loss: 1.9683 Acc: 0.4 lr: 8e-05: 100%|█| 2871/2871 [06:17<00:00,  7.61it\n",
      "100%|█████████████████████████████████████████| 129/129 [00:05<00:00, 21.87it/s]\n",
      "100%|███████████████████████████████████████| 2292/2292 [01:37<00:00, 23.62it/s]\n",
      "Val score: 0.45753204372748\n",
      "\n",
      "Epoch 7 Loss: 1.5738 Acc: 0.49 lr: 8e-05: 100%|█| 2871/2871 [06:17<00:00,  7.61i\n",
      "100%|█████████████████████████████████████████| 129/129 [00:05<00:00, 21.98it/s]\n",
      "100%|███████████████████████████████████████| 2292/2292 [01:37<00:00, 23.61it/s]\n",
      "Val score: 0.4771557165637657\n",
      "\n",
      "Epoch 8 Loss: 1.3307 Acc: 0.55 lr: 8e-05: 100%|█| 2871/2871 [06:17<00:00,  7.61i\n",
      "100%|█████████████████████████████████████████| 129/129 [00:05<00:00, 22.09it/s]\n",
      "100%|███████████████████████████████████████| 2292/2292 [01:36<00:00, 23.63it/s]\n",
      "Val score: 0.48787664542981457\n",
      "\n",
      "Epoch 9 Loss: 1.1703 Acc: 0.6 lr: 8e-05: 100%|█| 2871/2871 [06:17<00:00,  7.61it\n",
      "100%|█████████████████████████████████████████| 129/129 [00:05<00:00, 22.00it/s]\n",
      "100%|███████████████████████████████████████| 2292/2292 [01:37<00:00, 23.59it/s]\n",
      "Val score: 0.5002505853453436\n",
      "\n",
      "Epoch 10 Loss: 1.0655 Acc: 0.63 lr: 8e-05: 100%|█| 2871/2871 [06:19<00:00,  7.56\n",
      "100%|█████████████████████████████████████████| 129/129 [00:05<00:00, 21.62it/s]\n",
      "100%|███████████████████████████████████████| 2292/2292 [01:38<00:00, 23.26it/s]\n",
      "Val score: 0.5077235359515816\n",
      "\n",
      "Epoch 11 Loss: 0.9941 Acc: 0.65 lr: 8e-05: 100%|█| 2871/2871 [06:20<00:00,  7.54\n",
      "100%|█████████████████████████████████████████| 129/129 [00:05<00:00, 22.07it/s]\n",
      "100%|███████████████████████████████████████| 2292/2292 [01:37<00:00, 23.62it/s]\n",
      "Val score: 0.5118959212732301\n",
      "\n",
      "Epoch 12 Loss: 0.9454 Acc: 0.67 lr: 8e-05: 100%|█| 2871/2871 [06:17<00:00,  7.61\n",
      "100%|█████████████████████████████████████████| 129/129 [00:05<00:00, 21.82it/s]\n",
      "100%|███████████████████████████████████████| 2292/2292 [01:37<00:00, 23.61it/s]\n",
      "Val score: 0.5186763150115845\n",
      "\n",
      "Epoch 13 Loss: 0.914 Acc: 0.67 lr: 8e-05: 100%|█| 2871/2871 [06:17<00:00,  7.61i\n",
      "100%|█████████████████████████████████████████| 129/129 [00:05<00:00, 21.99it/s]\n",
      "100%|███████████████████████████████████████| 2292/2292 [01:37<00:00, 23.57it/s]\n",
      "Val score: 0.5219527988478166\n",
      "\n",
      "Epoch 14 Loss: 0.8925 Acc: 0.68 lr: 8e-05: 100%|█| 2871/2871 [06:22<00:00,  7.50\n",
      "100%|█████████████████████████████████████████| 129/129 [00:05<00:00, 21.75it/s]\n",
      "100%|███████████████████████████████████████| 2292/2292 [01:38<00:00, 23.29it/s]\n",
      "Val score: 0.5282788632524141\n",
      "\n",
      "Epoch 15 Loss: 0.8813 Acc: 0.68 lr: 8e-05: 100%|█| 2871/2871 [06:27<00:00,  7.40\n",
      "100%|█████████████████████████████████████████| 129/129 [00:05<00:00, 21.75it/s]\n",
      "100%|███████████████████████████████████████| 2292/2292 [01:37<00:00, 23.45it/s]\n",
      "Val score: 0.5286691738882909\n",
      "\n",
      "Epoch 16 Loss: 0.881 Acc: 0.67 lr: 8e-05: 100%|█| 2871/2871 [06:17<00:00,  7.60i\n",
      "100%|█████████████████████████████████████████| 129/129 [00:05<00:00, 22.00it/s]\n",
      "100%|███████████████████████████████████████| 2292/2292 [01:37<00:00, 23.59it/s]\n",
      "Val score: 0.532278723433796\n",
      "\n",
      "Epoch 17 Loss: 0.8855 Acc: 0.66 lr: 8e-05: 100%|█| 2871/2871 [06:26<00:00,  7.43\n",
      "100%|█████████████████████████████████████████| 129/129 [00:05<00:00, 22.03it/s]\n",
      "100%|███████████████████████████████████████| 2292/2292 [01:37<00:00, 23.59it/s]\n",
      "Val score: 0.532335546297434\n",
      "\n",
      "Epoch 18 Loss: 0.8988 Acc: 0.65 lr: 8e-05: 100%|█| 2871/2871 [06:17<00:00,  7.61\n",
      "100%|█████████████████████████████████████████| 129/129 [00:05<00:00, 21.97it/s]\n",
      "100%|███████████████████████████████████████| 2292/2292 [01:37<00:00, 23.57it/s]\n",
      "Val score: 0.5350320149303943\n",
      "\n",
      "Epoch 19 Loss: 0.9213 Acc: 0.64 lr: 8e-05: 100%|█| 2871/2871 [06:20<00:00,  7.54\n",
      "100%|█████████████████████████████████████████| 129/129 [00:05<00:00, 21.66it/s]\n",
      "100%|███████████████████████████████████████| 2292/2292 [01:38<00:00, 23.25it/s]\n",
      "Val score: 0.5367193504564579\n",
      "\n",
      "Epoch 20 Loss: 0.9511 Acc: 0.61 lr: 8e-05: 100%|█| 2871/2871 [06:20<00:00,  7.55\n",
      "100%|█████████████████████████████████████████| 129/129 [00:05<00:00, 21.97it/s]\n",
      "100%|███████████████████████████████████████| 2292/2292 [01:37<00:00, 23.58it/s]\n",
      "Val score: 0.5374650859690486\n",
      "\n"
     ]
    }
   ],
   "source": [
    "def cross_entropy(pred, soft_targets, sw):\n",
    "    logsoftmax = nn.LogSoftmax(dim=1)\n",
    "    return torch.mean(torch.sum(- soft_targets * logsoftmax(pred), 1)*sw)\n",
    "\n",
    "\n",
    "def train(model, epochs=2, n_batches=1, dynamic_margin=True):\n",
    "    train_df = topic_df_train.append(content_df[content_df[\"is_train\"]]).reset_index(drop=True)\n",
    "    train_ds = LECRDataset(train_df, tokenizer, MAX_LEN, aug=False)\n",
    "\n",
    "    train_loader = DataLoader(train_ds, batch_size=BS, shuffle=True, num_workers=NW,\n",
    "                               pin_memory=False, drop_last=True)\n",
    "\n",
    "    optimizer = get_optimizer(model)\n",
    "\n",
    "    scaler = torch.cuda.amp.GradScaler()\n",
    "\n",
    "    for e in range(epochs):\n",
    "        if dynamic_margin:\n",
    "            model.arcface.set_margin(0.1 + e*0.02)\n",
    "        model.cuda()\n",
    "        model.train()\n",
    "        tbar = tqdm(train_loader, file=sys.stdout)\n",
    "\n",
    "        lr = adjust_lr(optimizer, e)\n",
    "\n",
    "        loss_list, acc_list = [], []\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        for idx, (data, label, sw) in enumerate(tbar):\n",
    "            ids, mask, label = data[\"input_ids\"].cuda(), data[\"attention_mask\"].cuda(), label.cuda()\n",
    "\n",
    "            with torch.cuda.amp.autocast():\n",
    "                preds = model(ids, mask, label)\n",
    "                loss = cross_entropy(preds, label, sw.cuda())\n",
    "\n",
    "            loss = loss / n_batches\n",
    "            scaler.scale(loss).backward()\n",
    "\n",
    "            if ((idx + 1) % n_batches) == 0:\n",
    "                scaler.unscale_(optimizer)\n",
    "                nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",
    "                scaler.step(optimizer)\n",
    "                scaler.update()\n",
    "                optimizer.zero_grad()\n",
    "                \n",
    "            acc = (preds.softmax(dim=1)*(label > 0)).sum(axis=1).mean()\n",
    "\n",
    "            loss_list.append(loss.detach().cpu().item())\n",
    "            acc_list.append(acc.detach().cpu().item())\n",
    "            avg_loss = np.round(n_batches * np.mean(loss_list), 4)\n",
    "            avg_acc = np.round(np.mean(acc_list), 2)\n",
    "\n",
    "            tbar.set_description(f\"Epoch {e + 1} Loss: {avg_loss} Acc: {avg_acc} lr: {lr}\")\n",
    "\n",
    "        torch.save(model.state_dict(), \"models/tmp.pth\")\n",
    "        log_text = f\"Epoch {e+1}\\n\"\n",
    "        logfile = open(f\"logs/log.txt\", 'a')\n",
    "        logfile.write(log_text)\n",
    "        logfile.close()\n",
    "        \n",
    "        val_score = evaluate(topic_df_test, content_df)\n",
    "        print(\"Val score:\", val_score)\n",
    "        print()\n",
    "        \n",
    "        \n",
    "torch.save(model.state_dict(), \"models/tmp.pth\")\n",
    "log_text = f\"Epoch 0\\n\"\n",
    "logfile = open(f\"logs/log.txt\", 'a')\n",
    "logfile.write(log_text)\n",
    "logfile.close()\n",
    "\n",
    "train(model, epochs=20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "98dc7ac4",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_33343/2716447418.py:7: FutureWarning: The frame.append method is deprecated and will be removed from pandas in a future version. Use pandas.concat instead.\n",
      "  train_df = topic_df_train.append(content_df[content_df[\"is_train\"]]).reset_index(drop=True)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1 Loss: 0.7917 Acc: 0.69 lr: 2e-05: 100%|█| 2871/2871 [06:17<00:00,  7.61i\n",
      "100%|█████████████████████████████████████████| 129/129 [00:05<00:00, 21.99it/s]\n",
      "100%|███████████████████████████████████████| 2292/2292 [01:37<00:00, 23.62it/s]\n",
      "Val score: 0.5397609663533129\n",
      "\n",
      "Epoch 2 Loss: 0.6483 Acc: 0.77 lr: 8e-06: 100%|█| 2871/2871 [06:17<00:00,  7.60i\n",
      "100%|█████████████████████████████████████████| 129/129 [00:05<00:00, 21.98it/s]\n",
      "100%|███████████████████████████████████████| 2292/2292 [01:37<00:00, 23.60it/s]\n",
      "Val score: 0.541537604325936\n",
      "\n"
     ]
    }
   ],
   "source": [
    "def adjust_lr(optimizer, epoch):\n",
    "    if epoch >= 1:\n",
    "        lr = 8e-6\n",
    "    else:\n",
    "        lr = 2e-5\n",
    "\n",
    "    optimizer.param_groups[0]['lr'] = lr\n",
    "    optimizer.param_groups[1]['lr'] = 100*lr     \n",
    "    \n",
    "    return lr\n",
    "\n",
    "model.arcface.set_margin(0.5)\n",
    "train(model, epochs=2, dynamic_margin=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "b01ee5af",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.model.save(\"models/vec_model_v36bert.pth\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "038fce91",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████| 129/129 [00:05<00:00, 22.12it/s]\n",
      "100%|███████████████████████████████████████| 2292/2292 [01:37<00:00, 23.60it/s]\n"
     ]
    }
   ],
   "source": [
    "topic_ds = LECRDataset(topic_df_test, tokenizer, MAX_LEN, aug=False)\n",
    "content_ds = LECRDataset(content_df, tokenizer, MAX_LEN, aug=False)\n",
    "\n",
    "\n",
    "V_topic = to_vec(model, topic_ds)\n",
    "V_content = to_vec(model, content_ds)\n",
    "\n",
    "res_df = get_matches(V_topic, V_content, topic_df_test[\"id\"].values, content_df[\"id\"].values,\n",
    "                    n_neighbors=20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "1c3788f9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.8546878354504416"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def best_case_metric(pred_df):\n",
    "    scores = np.zeros(pred_df.shape[0])\n",
    "\n",
    "    for i, (gt, pred) in enumerate(zip(pred_df[\"content_ids\"].values, pred_df[\"pred\"].values)):\n",
    "        hits = 0\n",
    "\n",
    "        gt, pred = gt.split(), pred.split()\n",
    "\n",
    "        for l in gt:\n",
    "            if l in pred:\n",
    "                hits += 1\n",
    "\n",
    "        if hits > 0:           \n",
    "            recall = hits / len(gt)\n",
    "            precision = 1 #hits / len(pred)\n",
    "\n",
    "            scores[i] = f2(recall, precision)\n",
    "\n",
    "    return scores.mean()\n",
    "\n",
    "\n",
    "pred_df = res_df.groupby(\"topic_id\")[\"content_id\"].apply(lambda x: \" \".join(list(x)))\n",
    "pred_df = pred_df.reset_index().rename(columns={\"content_id\": 'pred'})\n",
    "pred_df = pred_df.merge(pd.read_csv(\"data/correlations.csv\"), on=\"topic_id\")\n",
    "\n",
    "best_case_metric(pred_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "b317fe32",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.5137049682496062"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res_df[\"rank\"] = res_df.groupby(\"topic_id\")[\"vec_dist\"].rank(method=\"dense\")\n",
    "\n",
    "\n",
    "pred_df = res_df[(res_df[\"rank\"] <= 2) | (res_df[\"vec_dist\"] < 0.3)]\n",
    "pred_df = pred_df.groupby(\"topic_id\")[\"content_id\"].apply(lambda x: \" \".join(list(x)))\n",
    "pred_df = pred_df.reset_index().rename(columns={\"content_id\": 'pred'})\n",
    "pred_df = pred_df.merge(pd.read_csv(\"data/correlations.csv\"), on=\"topic_id\")\n",
    "\n",
    "metric(pred_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "655829d7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<AxesSubplot:>"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAYMAAAD8CAYAAACVZ8iyAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAAAYDUlEQVR4nO3df5Dc9X3f8ec7qPwwCkhActVIakRqxa6M3AlcgY5nMieTggDXIhPiiiFGchRrEst2ptWMgbgZZfjRyg2EwtSmVkBFeFxkQtxBjXGoirXxkIkwEDBCYMwZZCOVHw4Scg/8o0fe/WM/B8txp9vbvdv7ID0fMzv6fj/fz3f3tXunfe/n+/l+9yIzkSQd2X5upgNIkmaexUCSZDGQJFkMJElYDCRJWAwkSbRRDCJic0S8FBGPj7FtfURkRJxS1iMiboqIwYh4LCJOb+m7KiKeLrdVLe1nRMSuss9NERFT9eQkSe1pZ2RwG7B8dGNELATOBX7Q0nw+sLjc1gI3l74nARuAs4AzgQ0RMbfsczPw8Zb93vZYkqTpNWExyMxvAvvH2HQD8Bmg9aq1FcDt2bQTmBMR84DzgO2ZuT8zDwDbgeVl2wmZuTObV7/dDlzU1TOSJE3arE52iogVwL7M/Paoozrzgeda1veWtkO17x2jfUKnnHJKLlq0qK28r776Kscff3xbfXup1lxQb7Zac0G92WrNBfVmqzUXdJ/t4Ycf/vvM/IXR7ZMuBhHxLuAPaR4i6qmIWEvz8BN9fX1cd911be03NDTE7NmzpzNaR2rNBfVmqzUX1Jut1lxQb7Zac0H32ZYtW/b9MTdk5oQ3YBHweFleCrwE7Cm3YZrzBv8Y+CJwSct+TwHzgEuAL7a0f7G0zQO+09L+ln6Hup1xxhnZrh07drTdt5dqzZVZb7Zac2XWm63WXJn1Zqs1V2b32YCHcoz31EmfWpqZuzLzFzNzUWYuonlo5/TMfAHYBlxWzio6GziYmc8D9wLnRsTcMnF8LnBv2fajiDi7nEV0GXD3ZDNJkrrTzqmldwB/C7wnIvZGxJpDdL8HeAYYBP4M+ARAZu4HrgYeLLerShulzy1ln+8BX+/sqUiSOjXhnEFmXjLB9kUtywmsG6ffZmDzGO0PAadNlEOSNH28AlmSZDGQJFkMJElYDCRJWAwkSXT4dRSS1EuLrvjamO17Nl7Y4ySHL0cGkiSLgSTJYiBJwjkDSRUZb25A08+RgSTJYiBJ8jCRpHcwTzmdOo4MJEkWA0mSxUCShMVAkoQTyJJmwMjE7/qlw6z22oIqODKQJFkMJEkWA0kSFgNJEm0Ug4jYHBEvRcTjLW1/EhHfiYjHIuJ/RMSclm1XRsRgRDwVEee1tC8vbYMRcUVL+6kR8UBp/0pEHD2Fz0+S1IZ2Rga3ActHtW0HTsvM9wPfBa4EiIglwErgfWWfL0TEURFxFPB54HxgCXBJ6QvwOeCGzHw3cABY09UzkiRN2oTFIDO/Cewf1fa/MnO4rO4EFpTlFcDWzPxpZj4LDAJnlttgZj6TmT8DtgIrIiKADwJ3lf23ABd195QkSZM1FXMGvwN8vSzPB55r2ba3tI3XfjLwSkthGWmXJPVQVxedRcRngWHgy1MTZ8LHWwusBejr66PRaLS139DQUNt9e6nWXFBvtlpzQb3Zasy1fmnz81/fcW8uT6Vun2+Nr9mI6crWcTGIiNXAh4BzMjNL8z5gYUu3BaWNcdpfBuZExKwyOmjt/zaZuQnYBNDf358DAwNtZW00GrTbt5dqzQX1Zqs1F9SbrcZcq1uuQL5+19R/EcKeSwe62r/G12zEdGXr6DBRRCwHPgN8ODNfa9m0DVgZEcdExKnAYuBbwIPA4nLm0NE0J5m3lSKyA7i47L8KuLuzpyJJ6tSEJTki7gAGgFMiYi+wgebZQ8cA25tzwOzMzN/LzN0RcSfwBM3DR+sy8/VyP58E7gWOAjZn5u7yEJcDWyPiGuAR4NYpfH6SjkCH+lvK/uGbsU1YDDLzkjGax33DzsxrgWvHaL8HuGeM9mdonm0kSZohXoEsSbIYSJIsBpIkLAaSJCwGkiQsBpIkLAaSJCwGkiQsBpIkLAaSJCwGkiQsBpIkuvzjNpJ0KIf69lDVxZGBJMliIEmyGEiSsBhIkrAYSJKwGEiSsBhIkrAYSJKwGEiSaKMYRMTmiHgpIh5vaTspIrZHxNPl37mlPSLipogYjIjHIuL0ln1Wlf5PR8SqlvYzImJX2eemiIipfpKSpENrZ2RwG7B8VNsVwH2ZuRi4r6wDnA8sLre1wM3QLB7ABuAs4Exgw0gBKX0+3rLf6MeSJE2zCYtBZn4T2D+qeQWwpSxvAS5qab89m3YCcyJiHnAesD0z92fmAWA7sLxsOyEzd2ZmAre33JckqUc6nTPoy8zny/ILQF9Zng8819Jvb2k7VPveMdolST3U9beWZmZGRE5FmIlExFqah5/o6+uj0Wi0td/Q0FDbfXup1lxQb7Zac0G92WYy1/qlw4fc3nfcxH2mWjuvRa0/S5i+bJ0WgxcjYl5mPl8O9bxU2vcBC1v6LSht+4CBUe2N0r5gjP5jysxNwCaA/v7+HBgYGK/rWzQaDdrt20u15oJ6s9WaC+rNNpO5Vk/wFdbrlw5z/a7efpP+nksHJuxT688Spi9bp4eJtgEjZwStAu5uab+snFV0NnCwHE66Fzg3IuaWieNzgXvLth9FxNnlLKLLWu5LktQjE5bkiLiD5qf6UyJiL82zgjYCd0bEGuD7wEdK93uAC4BB4DXgYwCZuT8irgYeLP2uysyRSelP0Dxj6Tjg6+UmSeqhCYtBZl4yzqZzxuibwLpx7mczsHmM9oeA0ybKIUmaPl6BLEmyGEiSpuDUUkl6J1k0zhlOezZe2OMkdXFkIEmyGEiSPEwkaQqMd+hF7xyODCRJFgNJksVAkoTFQJKExUCShMVAkoTFQJKExUCShMVAkoTFQJKExUCShMVAkoTFQJKExUCShMVAkoTFQJJEl8UgIv5tROyOiMcj4o6IODYiTo2IByJiMCK+EhFHl77HlPXBsn1Ry/1cWdqfiojzunxOkqRJ6rgYRMR84NNAf2aeBhwFrAQ+B9yQme8GDgBryi5rgAOl/YbSj4hYUvZ7H7Ac+EJEHNVpLknS5HV7mGgWcFxEzALeBTwPfBC4q2zfAlxUlleUdcr2cyIiSvvWzPxpZj4LDAJndplLkjQJHReDzNwHXAf8gGYROAg8DLySmcOl215gflmeDzxX9h0u/U9ubR9jH0lSD8zqdMeImEvzU/2pwCvAn9M8zDNtImItsBagr6+PRqPR1n5DQ0Nt9+2lWnNBvdlqzQX1ZutFrvVLhyfuNIa+4zrfd6q1vka1/ixh+rJ1XAyAXweezcwfAkTEV4EPAHMiYlb59L8A2Ff67wMWAnvLYaUTgZdb2ke07vMWmbkJ2ATQ39+fAwMDbQVtNBq027eXas0F9WarNRfUm60XuVZf8bWO9lu/dJjrd3XzNjR19lw68MZyrT9LmL5s3cwZ/AA4OyLeVY79nwM8AewALi59VgF3l+VtZZ2y/RuZmaV9ZTnb6FRgMfCtLnJJkiap45KcmQ9ExF3A3wHDwCM0P7V/DdgaEdeUtlvLLrcCX4qIQWA/zTOIyMzdEXEnzUIyDKzLzNc7zSVJmryuxmeZuQHYMKr5GcY4GygzfwL81jj3cy1wbTdZJEmd8wpkSVJ3IwNJR5ZFHU4Uq36ODCRJFgNJksVAkoTFQJKExUCShMVAkoTFQJKExUCShMVAkoRXIEsS8Narq9cvHX7ja7n3bLxwpiL1lCMDSZLFQJJkMZAkYTGQJGExkCRhMZAkYTGQJGExkCThRWeSxuCftzzyODKQJHVXDCJiTkTcFRHfiYgnI+JfRsRJEbE9Ip4u/84tfSMiboqIwYh4LCJOb7mfVaX/0xGxqtsnJUmanG5HBjcCf5WZ7wX+OfAkcAVwX2YuBu4r6wDnA4vLbS1wM0BEnARsAM4CzgQ2jBQQSVJvdFwMIuJE4NeAWwEy82eZ+QqwAthSum0BLirLK4Dbs2knMCci5gHnAdszc39mHgC2A8s7zSVJmrxuRganAj8E/ltEPBIRt0TE8UBfZj5f+rwA9JXl+cBzLfvvLW3jtUuSeiQys7MdI/qBncAHMvOBiLgR+BHwqcyc09LvQGbOjYi/BDZm5v2l/T7gcmAAODYzryntfwT8ODOvG+Mx19I8xERfX98ZW7dubSvr0NAQs2fP7uh5Tqdac0G92WrNBfVm6yTXrn0HpynNW/UdBy/+uCcPNSmtuZbOP3Fmw4zS7e/ZsmXLHs7M/tHt3ZxauhfYm5kPlPW7aM4PvBgR8zLz+XIY6KWyfR+wsGX/BaVtH82C0NreGOsBM3MTsAmgv78/BwYGxur2No1Gg3b79lKtuaDebLXmgnqzdZJrdY9OLV2/dJjrd9V3hntrrj2XDsxsmFGm6/es48NEmfkC8FxEvKc0nQM8AWwDRs4IWgXcXZa3AZeVs4rOBg6Ww0n3AudGxNwycXxuaZMk9Ui3JflTwJcj4mjgGeBjNAvMnRGxBvg+8JHS9x7gAmAQeK30JTP3R8TVwIOl31WZub/LXJKkSeiqGGTmo8Dbjj3RHCWM7pvAunHuZzOwuZsskqTOeQWyJMliIEmyGEiSsBhIkrAYSJKwGEiSsBhIkrAYSJLwz15KRzT/vKVGODKQJFkMJEkWA0kSFgNJEhYDSRIWA0kSFgNJEhYDSRIWA0kSXoEsSYc03lXaezZe2OMk08uRgSTJYiBJshhIkpiCOYOIOAp4CNiXmR+KiFOBrcDJwMPARzPzZxFxDHA7cAbwMvBvMnNPuY8rgTXA68CnM/PebnNJepPfTqqJTMXI4A+AJ1vWPwfckJnvBg7QfJOn/HugtN9Q+hERS4CVwPuA5cAXSoGRJPVIV8UgIhYAFwK3lPUAPgjcVbpsAS4qyyvKOmX7OaX/CmBrZv40M58FBoEzu8klSZqcbkcG/xn4DPAPZf1k4JXMHC7re4H5ZXk+8BxA2X6w9H+jfYx9JEk90PGcQUR8CHgpMx+OiIEpS3Tox1wLrAXo6+uj0Wi0td/Q0FDbfXup1lxQb7Zac0G92YaGhli/9PWZjjGmvuNg/dLhiTv2WDu5ZupnPV2/Z91MIH8A+HBEXAAcC5wA3AjMiYhZ5dP/AmBf6b8PWAjsjYhZwIk0J5JH2ke07vMWmbkJ2ATQ39+fAwMDbQVtNBq027eXas0F9WarNRfUm63RaHD9/a/OdIwxrV86zPW76rv2tZ1cey4d6E2YUabr96zjw0SZeWVmLsjMRTQngL+RmZcCO4CLS7dVwN1leVtZp2z/RmZmaV8ZEceUM5EWA9/qNJckafKmoyRfDmyNiGuAR4BbS/utwJciYhDYT7OAkJm7I+JO4AlgGFiXmXWOaSXpMDUlxSAzG0CjLD/DGGcDZeZPgN8aZ/9rgWunIoskafLqO1gnqWNjXVzWnAj1v7oOza+jkCRZDCRJFgNJEhYDSRIWA0kSFgNJEp5vJr3j+LcJNB0cGUiSLAaSJIuBJAmLgSQJi4EkCYuBJAmLgSQJi4EkCS86k6SOjHfx356NF/Y4ydSwGEiV8kpj9ZKHiSRJjgykmeYIQDVwZCBJshhIkrooBhGxMCJ2RMQTEbE7Iv6gtJ8UEdsj4uny79zSHhFxU0QMRsRjEXF6y32tKv2fjohV3T8tSdJkdDMyGAbWZ+YS4GxgXUQsAa4A7svMxcB9ZR3gfGBxua0FboZm8QA2AGcBZwIbRgqIJKk3Op5AzszngefL8v+NiCeB+cAKYKB02wI0gMtL++2ZmcDOiJgTEfNK3+2ZuR8gIrYDy4E7Os0m1ciJYtVsSuYMImIR8KvAA0BfKRQALwB9ZXk+8FzLbntL23jtkqQeieYH9S7uIGI28NfAtZn51Yh4JTPntGw/kJlzI+IvgY2ZeX9pv4/miGEAODYzryntfwT8ODOvG+Ox1tI8xERfX98ZW7dubSvj0NAQs2fP7uJZTo9ac0G92WrNBRNn27XvYA/TvKnvOHjxxzPy0BOqNVs3uZbOP3Fqw4zS7f+BZcuWPZyZ/aPbu7rOICL+EfAXwJcz86ul+cWImJeZz5fDQC+V9n3AwpbdF5S2fbx5WGmkvTHW42XmJmATQH9/fw4MDIzV7W0ajQbt9u2lWnNBvdlqzQUTZ1s9Q4eJ1i8d5vpddV5SVGu2bnLtuXRgasOMMl3/Bzr+KUREALcCT2bmn7Zs2gasAjaWf+9uaf9kRGylOVl8sBSMe4H/0DJpfC5wZae5pF4ZPQewfunwjL3hS93qpiR/APgosCsiHi1tf0izCNwZEWuA7wMfKdvuAS4ABoHXgI8BZOb+iLgaeLD0u2pkMlmS1BvdnE10PxDjbD5njP4JrBvnvjYDmzvNIknqjlcgS5L8ojppIl4foCOBxUAqfNPXkczDRJIkRwY68jgC0HR6p/45TEcGkiRHBjp8OQKQ2ufIQJJkMZAkeZhIhwEPB0ndsxjoHWHkDd8vg5Omh4eJJEkWA0mSh4lUGY//SzPDkYEkyZGBZoYjAKkuFgNNK9/0pXcGi4GmhG/60qHV/gV2zhlIkhwZaGy79h304i7pCOLIQJLkyOBIN95xzPVLexxE0oyqphhExHLgRuAo4JbM3DjDkd6RnMiV1IkqikFEHAV8HvhXwF7gwYjYlplP9DJHr99IR750rZOzCXzTlw4PtZxlVEUxAM4EBjPzGYCI2AqsAKalGNT2RlpbHklHnlqKwXzguZb1vcBZM5RFkmbceB8Sb1t+/LQ8XmTmtNzxpEJEXAwsz8zfLesfBc7KzE+O6rcWWFtW3wM81eZDnAL8/RTFnUq15oJ6s9WaC+rNVmsuqDdbrbmg+2y/lJm/MLqxlpHBPmBhy/qC0vYWmbkJ2DTZO4+IhzKzv/N406PWXFBvtlpzQb3Zas0F9WarNRdMX7ZarjN4EFgcEadGxNHASmDbDGeSpCNGFSODzByOiE8C99I8tXRzZu6e4ViSdMSoohgAZOY9wD3TdPeTPrTUI7Xmgnqz1ZoL6s1Way6oN1utuWCaslUxgSxJmlm1zBlIkmbQYVUMImJ5RDwVEYMRccUY24+JiK+U7Q9ExKJKcv1aRPxdRAyX02x7oo1c/y4inoiIxyLivoj4pYqy/V5E7IqIRyPi/ohYUkOuln6/GREZET07I6WN12x1RPywvGaPRsTv1pCr9PlI+V3bHRH/vRe52skWETe0vF7fjYhXKsn1TyJiR0Q8Uv5/XtD1g2bmYXGjOfH8PeCXgaOBbwNLRvX5BPBfy/JK4CuV5FoEvB+4Hbi4otdrGfCusvz7vXi9JpHthJblDwN/VUOu0u/ngW8CO4H+il6z1cB/6UWeSeZaDDwCzC3rv1hLtlH9P0Xz5JYZz0Vz3uD3y/ISYE+3j3s4jQze+EqLzPwZMPKVFq1WAFvK8l3AORERM50rM/dk5mPAP0xzlsnm2pGZr5XVnTSv/6gl249aVo8HejH51c7vGMDVwOeAn/Qg02Sz9Vo7uT4OfD4zDwBk5ksVZWt1CXBHJbkSOKEsnwj8n24f9HAqBmN9pcX88fpk5jBwEDi5glwzYbK51gBfn9ZEb2orW0Ssi4jvAf8J+HQNuSLidGBhZvb6C6fa/Xn+ZjmscFdELBxj+0zk+hXgVyLibyJiZ/kG415o+/9AOUR6KvCNSnL9MfDbEbGX5lmYn+r2QQ+nYqBpEhG/DfQDfzLTWVpl5ucz858ClwP/fqbzRMTPAX8KrJ/pLOP4n8CizHw/sJ03R8kzbRbNQ0UDND99/1lEzJnJQGNYCdyVma/PdJDiEuC2zFwAXAB8qfz+dexwKgbtfKXFG30iYhbN4dXLFeSaCW3liohfBz4LfDgzf1pTthZbgYumM1AxUa6fB04DGhGxBzgb2NajSeQJX7PMfLnlZ3gLcEYNuWh+8t2Wmf8vM58FvkuzONSQbcRKenOICNrLtQa4EyAz/xY4luZ3FnWuFxM1vbjR/HTxDM2h3Miky/tG9VnHWyeQ76whV0vf2+jdBHI7r9ev0pzIWlzhz3Jxy/K/Bh6qIdeo/g16N4Hczms2r2X5N4CdleRaDmwpy6fQPERycg3ZSr/3Anso12XVkIvmIdvVZfmf0Zwz6CrftD+xXt5oDpe+W97APlvarqL5qRaa1fPPgUHgW8AvV5LrX9D8dPQqzZHK7kpy/W/gReDRcttW0c/yRmB3ybXjUG/Kvcw1qm/PikGbr9l/LK/Zt8tr9t5KcgXNw2tPALuAlbW8ZmX9j4GNvcrU5mu2BPib8rN8FDi328f0CmRJ0mE1ZyBJ6pDFQJJkMZAkWQwkSVgMJElYDCRJWAwkSVgMJEnA/wcbiqJhrxhSYgAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "res_df[\"vec_dist\"].hist(bins=50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "fdbd4bf6",
   "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>topic_id</th>\n",
       "      <th>content_id</th>\n",
       "      <th>vec_dist</th>\n",
       "      <th>rank</th>\n",
       "      <th>id_x</th>\n",
       "      <th>language_x</th>\n",
       "      <th>id_y</th>\n",
       "      <th>language_y</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>t_000d1fb3f2f5</td>\n",
       "      <td>c_acaa5b1ab542</td>\n",
       "      <td>0.377952</td>\n",
       "      <td>1.0</td>\n",
       "      <td>c_acaa5b1ab542</td>\n",
       "      <td>en</td>\n",
       "      <td>t_000d1fb3f2f5</td>\n",
       "      <td>en</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>t_000d1fb3f2f5</td>\n",
       "      <td>c_1ec97b588bce</td>\n",
       "      <td>0.394667</td>\n",
       "      <td>2.0</td>\n",
       "      <td>c_1ec97b588bce</td>\n",
       "      <td>en</td>\n",
       "      <td>t_000d1fb3f2f5</td>\n",
       "      <td>en</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>t_000d1fb3f2f5</td>\n",
       "      <td>c_b8d730238789</td>\n",
       "      <td>0.399614</td>\n",
       "      <td>3.0</td>\n",
       "      <td>c_b8d730238789</td>\n",
       "      <td>en</td>\n",
       "      <td>t_000d1fb3f2f5</td>\n",
       "      <td>en</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>t_000d1fb3f2f5</td>\n",
       "      <td>c_7868f655c31e</td>\n",
       "      <td>0.439027</td>\n",
       "      <td>4.0</td>\n",
       "      <td>c_7868f655c31e</td>\n",
       "      <td>en</td>\n",
       "      <td>t_000d1fb3f2f5</td>\n",
       "      <td>en</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>t_000d1fb3f2f5</td>\n",
       "      <td>c_7cc189e7acb0</td>\n",
       "      <td>0.477397</td>\n",
       "      <td>5.0</td>\n",
       "      <td>c_7cc189e7acb0</td>\n",
       "      <td>en</td>\n",
       "      <td>t_000d1fb3f2f5</td>\n",
       "      <td>en</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>164955</th>\n",
       "      <td>t_c988e92bc1e5</td>\n",
       "      <td>c_83a15becefaf</td>\n",
       "      <td>0.640286</td>\n",
       "      <td>16.0</td>\n",
       "      <td>c_83a15becefaf</td>\n",
       "      <td>sw</td>\n",
       "      <td>t_c988e92bc1e5</td>\n",
       "      <td>sw</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>164956</th>\n",
       "      <td>t_c988e92bc1e5</td>\n",
       "      <td>c_31c0b5a57cb9</td>\n",
       "      <td>0.640469</td>\n",
       "      <td>17.0</td>\n",
       "      <td>c_31c0b5a57cb9</td>\n",
       "      <td>sw</td>\n",
       "      <td>t_c988e92bc1e5</td>\n",
       "      <td>sw</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>164957</th>\n",
       "      <td>t_c988e92bc1e5</td>\n",
       "      <td>c_3e7250b6b5dd</td>\n",
       "      <td>0.641504</td>\n",
       "      <td>18.0</td>\n",
       "      <td>c_3e7250b6b5dd</td>\n",
       "      <td>sw</td>\n",
       "      <td>t_c988e92bc1e5</td>\n",
       "      <td>sw</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>164958</th>\n",
       "      <td>t_c988e92bc1e5</td>\n",
       "      <td>c_ad037f7f0f48</td>\n",
       "      <td>0.644516</td>\n",
       "      <td>19.0</td>\n",
       "      <td>c_ad037f7f0f48</td>\n",
       "      <td>sw</td>\n",
       "      <td>t_c988e92bc1e5</td>\n",
       "      <td>sw</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>164959</th>\n",
       "      <td>t_c988e92bc1e5</td>\n",
       "      <td>c_6eef87ee03d7</td>\n",
       "      <td>0.653560</td>\n",
       "      <td>20.0</td>\n",
       "      <td>c_6eef87ee03d7</td>\n",
       "      <td>sw</td>\n",
       "      <td>t_c988e92bc1e5</td>\n",
       "      <td>sw</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>164960 rows × 8 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "              topic_id      content_id  vec_dist  rank            id_x  \\\n",
       "0       t_000d1fb3f2f5  c_acaa5b1ab542  0.377952   1.0  c_acaa5b1ab542   \n",
       "1       t_000d1fb3f2f5  c_1ec97b588bce  0.394667   2.0  c_1ec97b588bce   \n",
       "2       t_000d1fb3f2f5  c_b8d730238789  0.399614   3.0  c_b8d730238789   \n",
       "3       t_000d1fb3f2f5  c_7868f655c31e  0.439027   4.0  c_7868f655c31e   \n",
       "4       t_000d1fb3f2f5  c_7cc189e7acb0  0.477397   5.0  c_7cc189e7acb0   \n",
       "...                ...             ...       ...   ...             ...   \n",
       "164955  t_c988e92bc1e5  c_83a15becefaf  0.640286  16.0  c_83a15becefaf   \n",
       "164956  t_c988e92bc1e5  c_31c0b5a57cb9  0.640469  17.0  c_31c0b5a57cb9   \n",
       "164957  t_c988e92bc1e5  c_3e7250b6b5dd  0.641504  18.0  c_3e7250b6b5dd   \n",
       "164958  t_c988e92bc1e5  c_ad037f7f0f48  0.644516  19.0  c_ad037f7f0f48   \n",
       "164959  t_c988e92bc1e5  c_6eef87ee03d7  0.653560  20.0  c_6eef87ee03d7   \n",
       "\n",
       "       language_x            id_y language_y  \n",
       "0              en  t_000d1fb3f2f5         en  \n",
       "1              en  t_000d1fb3f2f5         en  \n",
       "2              en  t_000d1fb3f2f5         en  \n",
       "3              en  t_000d1fb3f2f5         en  \n",
       "4              en  t_000d1fb3f2f5         en  \n",
       "...           ...             ...        ...  \n",
       "164955         sw  t_c988e92bc1e5         sw  \n",
       "164956         sw  t_c988e92bc1e5         sw  \n",
       "164957         sw  t_c988e92bc1e5         sw  \n",
       "164958         sw  t_c988e92bc1e5         sw  \n",
       "164959         sw  t_c988e92bc1e5         sw  \n",
       "\n",
       "[164960 rows x 8 columns]"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = res_df.merge(content_df[[\"id\", \"language\"]], left_on=\"content_id\", right_on=\"id\")\n",
    "x = x.merge(topic_df[[\"id\", \"language\"]], left_on=\"topic_id\", right_on=\"id\")\n",
    "x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "420793e2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.9537039282250243, 0.967804784566738)"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x[\"same_lang\"] = x[\"language_x\"] == x[\"language_y\"]\n",
    "x[\"same_lang\"].mean(), x[x[\"vec_dist\"] < 0.5][\"same_lang\"].mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "8c7c086e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.5953627609674871"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pred_df = x[(x[\"rank\"] <= 1) | (x[\"same_lang\"] & (x[\"vec_dist\"] < 0.45))]\n",
    "pred_df = pred_df.groupby(\"topic_id\")[\"content_id\"].apply(lambda x: \" \".join(list(x)))\n",
    "pred_df = pred_df.reset_index().rename(columns={\"content_id\": 'pred'})\n",
    "pred_df = pred_df.merge(pd.read_csv(\"data/correlations.csv\"), on=\"topic_id\")\n",
    "\n",
    "metric(pred_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "b36b3f8d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.9047889836231436"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res_df = get_matches(V_topic, V_content, topic_df_test[\"id\"].values, content_df[\"id\"].values,\n",
    "                    n_neighbors=50)\n",
    "\n",
    "pred_df = res_df.groupby(\"topic_id\")[\"content_id\"].apply(lambda x: \" \".join(list(x)))\n",
    "pred_df = pred_df.reset_index().rename(columns={\"content_id\": 'pred'})\n",
    "pred_df = pred_df.merge(pd.read_csv(\"data/correlations.csv\"), on=\"topic_id\")\n",
    "\n",
    "best_case_metric(pred_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f21d9af",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "964c1b1b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}