examples/00_quick_start/nrms_ebnerd.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Getting started\n",
    "\n",
    "In this notebook, we illustrate how to use the Neural News Recommendation with Multi-Head Self-Attention ([NRMS](https://aclanthology.org/D19-1671/)). The implementation is taken from the [recommenders](https://github.com/recommenders-team/recommenders) repository. We have simply stripped the model to keep it cleaner.\n",
    "\n",
    "We use a small dataset, which is downloaded from [recsys.eb.dk](https://recsys.eb.dk/). All the datasets are stored in the folder path ```~/ebnerd_data/*```."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load functionality"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer, AutoModel\n",
    "from pathlib import Path\n",
    "import tensorflow as tf\n",
    "import polars as pl\n",
    "\n",
    "from ebrec.utils._constants import (\n",
    "    DEFAULT_HISTORY_ARTICLE_ID_COL,\n",
    "    DEFAULT_CLICKED_ARTICLES_COL,\n",
    "    DEFAULT_INVIEW_ARTICLES_COL,\n",
    "    DEFAULT_IMPRESSION_ID_COL,\n",
    "    DEFAULT_SUBTITLE_COL,\n",
    "    DEFAULT_LABELS_COL,\n",
    "    DEFAULT_TITLE_COL,\n",
    "    DEFAULT_USER_COL,\n",
    ")\n",
    "\n",
    "from ebrec.utils._behaviors import (\n",
    "    create_binary_labels_column,\n",
    "    sampling_strategy_wu2019,\n",
    "    add_known_user_column,\n",
    "    add_prediction_scores,\n",
    "    truncate_history,\n",
    ")\n",
    "from ebrec.evaluation import MetricEvaluator, AucScore, NdcgScore, MrrScore\n",
    "from ebrec.utils._articles import convert_text2encoding_with_transformers\n",
    "from ebrec.utils._polars import concat_str_columns, slice_join_dataframes\n",
    "from ebrec.utils._articles import create_article_id_to_value_mapping\n",
    "from ebrec.utils._nlp import get_transformers_word_embeddings\n",
    "from ebrec.utils._python import write_submission_file, rank_predictions_by_score\n",
    "\n",
    "from ebrec.models.newsrec.dataloader import NRMSDataLoader\n",
    "from ebrec.models.newsrec.model_config import hparams_nrms\n",
    "from ebrec.models.newsrec import NRMSModel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ebnerd_from_path(path: Path, history_size: int = 30) -> pl.DataFrame:\n",
    "    \"\"\"\n",
    "    Load ebnerd - function\n",
    "    \"\"\"\n",
    "    df_history = (\n",
    "        pl.scan_parquet(path.joinpath(\"history.parquet\"))\n",
    "        .select(DEFAULT_USER_COL, DEFAULT_HISTORY_ARTICLE_ID_COL)\n",
    "        .pipe(\n",
    "            truncate_history,\n",
    "            column=DEFAULT_HISTORY_ARTICLE_ID_COL,\n",
    "            history_size=history_size,\n",
    "            padding_value=0,\n",
    "            enable_warning=False,\n",
    "        )\n",
    "    )\n",
    "    df_behaviors = (\n",
    "        pl.scan_parquet(path.joinpath(\"behaviors.parquet\"))\n",
    "        .collect()\n",
    "        .pipe(\n",
    "            slice_join_dataframes,\n",
    "            df2=df_history.collect(),\n",
    "            on=DEFAULT_USER_COL,\n",
    "            how=\"left\",\n",
    "        )\n",
    "    )\n",
    "    return df_behaviors"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate labels\n",
    "We sample a few just to get started. For testset we just make up a dummy column with 0 and 1 - this is not the true labels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "PATH = Path(\"~/ebnerd_data\")\n",
    "DATASPLIT = \"ebnerd_small\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this example we sample the dataset, just to keep it smaller. Also, one can simply add the testset similary to the validation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (2, 6)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>user_id</th><th>article_id_fixed</th><th>article_ids_inview</th><th>article_ids_clicked</th><th>impression_id</th><th>labels</th></tr><tr><td>u32</td><td>list[i32]</td><td>list[i64]</td><td>list[i64]</td><td>u32</td><td>list[i8]</td></tr></thead><tbody><tr><td>765999</td><td>[9764822, 9764640, … 9747803]</td><td>[9770792, 9767697, … 9770729]</td><td>[9770729]</td><td>184706144</td><td>[0, 0, … 1]</td></tr><tr><td>292310</td><td>[9747757, 9769981, … 9735909]</td><td>[9756785, 9772635, … 9756785]</td><td>[9755712]</td><td>346725041</td><td>[0, 0, … 0]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (2, 6)\n",
       "┌─────────┬───────────────────┬───────────────────┬──────────────────┬───────────────┬─────────────┐\n",
       "│ user_id ┆ article_id_fixed  ┆ article_ids_invie ┆ article_ids_clic ┆ impression_id ┆ labels      │\n",
       "│ ---     ┆ ---               ┆ w                 ┆ ked              ┆ ---           ┆ ---         │\n",
       "│ u32     ┆ list[i32]         ┆ ---               ┆ ---              ┆ u32           ┆ list[i8]    │\n",
       "│         ┆                   ┆ list[i64]         ┆ list[i64]        ┆               ┆             │\n",
       "╞═════════╪═══════════════════╪═══════════════════╪══════════════════╪═══════════════╪═════════════╡\n",
       "│ 765999  ┆ [9764822,         ┆ [9770792,         ┆ [9770729]        ┆ 184706144     ┆ [0, 0, … 1] │\n",
       "│         ┆ 9764640, …        ┆ 9767697, …        ┆                  ┆               ┆             │\n",
       "│         ┆ 9747803]          ┆ 9770729]          ┆                  ┆               ┆             │\n",
       "│ 292310  ┆ [9747757,         ┆ [9756785,         ┆ [9755712]        ┆ 346725041     ┆ [0, 0, … 0] │\n",
       "│         ┆ 9769981, …        ┆ 9772635, …        ┆                  ┆               ┆             │\n",
       "│         ┆ 9735909]          ┆ 9756785]          ┆                  ┆               ┆             │\n",
       "└─────────┴───────────────────┴───────────────────┴──────────────────┴───────────────┴─────────────┘"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "COLUMNS = [\n",
    "    DEFAULT_USER_COL,\n",
    "    DEFAULT_HISTORY_ARTICLE_ID_COL,\n",
    "    DEFAULT_INVIEW_ARTICLES_COL,\n",
    "    DEFAULT_CLICKED_ARTICLES_COL,\n",
    "    DEFAULT_IMPRESSION_ID_COL,\n",
    "]\n",
    "HISTORY_SIZE = 10\n",
    "FRACTION = 0.01\n",
    "\n",
    "df_train = (\n",
    "    ebnerd_from_path(PATH.joinpath(DATASPLIT, \"train\"), history_size=HISTORY_SIZE)\n",
    "    .select(COLUMNS)\n",
    "    .pipe(\n",
    "        sampling_strategy_wu2019,\n",
    "        npratio=4,\n",
    "        shuffle=True,\n",
    "        with_replacement=True,\n",
    "        seed=123,\n",
    "    )\n",
    "    .pipe(create_binary_labels_column)\n",
    "    .sample(fraction=FRACTION)\n",
    ")\n",
    "# =>\n",
    "df_validation = (\n",
    "    ebnerd_from_path(PATH.joinpath(DATASPLIT, \"validation\"), history_size=HISTORY_SIZE)\n",
    "    .select(COLUMNS)\n",
    "    .pipe(create_binary_labels_column)\n",
    "    .sample(fraction=FRACTION)\n",
    ")\n",
    "df_train.head(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load articles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (2, 21)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>article_id</th><th>title</th><th>subtitle</th><th>last_modified_time</th><th>premium</th><th>body</th><th>published_time</th><th>image_ids</th><th>article_type</th><th>url</th><th>ner_clusters</th><th>entity_groups</th><th>topics</th><th>category</th><th>subcategory</th><th>category_str</th><th>total_inviews</th><th>total_pageviews</th><th>total_read_time</th><th>sentiment_score</th><th>sentiment_label</th></tr><tr><td>i32</td><td>str</td><td>str</td><td>datetime[μs]</td><td>bool</td><td>str</td><td>datetime[μs]</td><td>list[i64]</td><td>str</td><td>str</td><td>list[str]</td><td>list[str]</td><td>list[str]</td><td>i16</td><td>list[i16]</td><td>str</td><td>i32</td><td>i32</td><td>f32</td><td>f32</td><td>str</td></tr></thead><tbody><tr><td>3000022</td><td>&quot;Hanks beskyldt…</td><td>&quot;Tom Hanks har …</td><td>2023-06-29 06:20:32</td><td>false</td><td>&quot;Tom Hanks skul…</td><td>2006-09-20 09:24:18</td><td>[3518381]</td><td>&quot;article_defaul…</td><td>&quot;https://ekstra…</td><td>[&quot;David Gardner&quot;]</td><td>[&quot;PER&quot;]</td><td>[&quot;Kriminalitet&quot;, &quot;Kendt&quot;, … &quot;Litteratur&quot;]</td><td>414</td><td>[432]</td><td>&quot;underholdning&quot;</td><td>null</td><td>null</td><td>null</td><td>0.9911</td><td>&quot;Negative&quot;</td></tr><tr><td>3000063</td><td>&quot;Bostrups aske …</td><td>&quot;Studieværten b…</td><td>2023-06-29 06:20:32</td><td>false</td><td>&quot;Strålende sens…</td><td>2006-09-24 07:45:30</td><td>[3170935, 3170939]</td><td>&quot;article_defaul…</td><td>&quot;https://ekstra…</td><td>[]</td><td>[]</td><td>[&quot;Kendt&quot;, &quot;Underholdning&quot;, … &quot;Personlig begivenhed&quot;]</td><td>118</td><td>[133]</td><td>&quot;nyheder&quot;</td><td>null</td><td>null</td><td>null</td><td>0.5155</td><td>&quot;Neutral&quot;</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (2, 21)\n",
       "┌───────────┬───────────┬───────────┬───────────┬───┬───────────┬───────────┬───────────┬──────────┐\n",
       "│ article_i ┆ title     ┆ subtitle  ┆ last_modi ┆ … ┆ total_pag ┆ total_rea ┆ sentiment ┆ sentimen │\n",
       "│ d         ┆ ---       ┆ ---       ┆ fied_time ┆   ┆ eviews    ┆ d_time    ┆ _score    ┆ t_label  │\n",
       "│ ---       ┆ str       ┆ str       ┆ ---       ┆   ┆ ---       ┆ ---       ┆ ---       ┆ ---      │\n",
       "│ i32       ┆           ┆           ┆ datetime[ ┆   ┆ i32       ┆ f32       ┆ f32       ┆ str      │\n",
       "│           ┆           ┆           ┆ μs]       ┆   ┆           ┆           ┆           ┆          │\n",
       "╞═══════════╪═══════════╪═══════════╪═══════════╪═══╪═══════════╪═══════════╪═══════════╪══════════╡\n",
       "│ 3000022   ┆ Hanks     ┆ Tom Hanks ┆ 2023-06-2 ┆ … ┆ null      ┆ null      ┆ 0.9911    ┆ Negative │\n",
       "│           ┆ beskyldt  ┆ har angiv ┆ 9         ┆   ┆           ┆           ┆           ┆          │\n",
       "│           ┆ for misha ┆ eligt     ┆ 06:20:32  ┆   ┆           ┆           ┆           ┆          │\n",
       "│           ┆ ndling    ┆ mishand…  ┆           ┆   ┆           ┆           ┆           ┆          │\n",
       "│ 3000063   ┆ Bostrups  ┆ Studievær ┆ 2023-06-2 ┆ … ┆ null      ┆ null      ┆ 0.5155    ┆ Neutral  │\n",
       "│           ┆ aske      ┆ ten blev  ┆ 9         ┆   ┆           ┆           ┆           ┆          │\n",
       "│           ┆ spredt i  ┆ mindet    ┆ 06:20:32  ┆   ┆           ┆           ┆           ┆          │\n",
       "│           ┆ Furesøen  ┆ med gla…  ┆           ┆   ┆           ┆           ┆           ┆          │\n",
       "└───────────┴───────────┴───────────┴───────────┴───┴───────────┴───────────┴───────────┴──────────┘"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_articles = pl.read_parquet(PATH.joinpath(\"articles.parquet\"))\n",
    "df_articles.head(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Init model using HuggingFace's tokenizer and wordembedding\n",
    "In the original implementation, they use the GloVe embeddings and tokenizer. To get going fast, we'll use a multilingual LLM from Hugging Face. \n",
    "Utilizing the tokenizer to tokenize the articles and the word-embedding to init NRMS.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "TRANSFORMER_MODEL_NAME = \"FacebookAI/xlm-roberta-base\"\n",
    "TEXT_COLUMNS_TO_USE = [DEFAULT_SUBTITLE_COL, DEFAULT_TITLE_COL]\n",
    "MAX_TITLE_LENGTH = 30\n",
    "\n",
    "# LOAD HUGGINGFACE:\n",
    "transformer_model = AutoModel.from_pretrained(TRANSFORMER_MODEL_NAME)\n",
    "transformer_tokenizer = AutoTokenizer.from_pretrained(TRANSFORMER_MODEL_NAME)\n",
    "\n",
    "# We'll init the word embeddings using the\n",
    "word2vec_embedding = get_transformers_word_embeddings(transformer_model)\n",
    "#\n",
    "df_articles, cat_cal = concat_str_columns(df_articles, columns=TEXT_COLUMNS_TO_USE)\n",
    "df_articles, token_col_title = convert_text2encoding_with_transformers(\n",
    "    df_articles, transformer_tokenizer, cat_cal, max_length=MAX_TITLE_LENGTH\n",
    ")\n",
    "# =>\n",
    "article_mapping = create_article_id_to_value_mapping(\n",
    "    df=df_articles, value_col=token_col_title\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Initiate the dataloaders\n",
    "In the implementations we have disconnected the models and data. Hence, you should built a dataloader that fits your needs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataloader = NRMSDataLoader(\n",
    "    behaviors=df_train,\n",
    "    article_dict=article_mapping,\n",
    "    unknown_representation=\"zeros\",\n",
    "    history_column=DEFAULT_HISTORY_ARTICLE_ID_COL,\n",
    "    eval_mode=False,\n",
    "    batch_size=64,\n",
    ")\n",
    "val_dataloader = NRMSDataLoader(\n",
    "    behaviors=df_validation,\n",
    "    article_dict=article_mapping,\n",
    "    unknown_representation=\"zeros\",\n",
    "    history_column=DEFAULT_HISTORY_ARTICLE_ID_COL,\n",
    "    eval_mode=True,\n",
    "    batch_size=32,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train the model\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37/37 [==============================] - ETA: 0s - loss: 1.6025"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-05-23 22:39:00.135465: W external/local_tsl/tsl/framework/cpu_allocator_impl.cc:83] Allocation of 46080368640 exceeds 10% of free system memory.\n",
      "2024-05-23 22:39:12.235572: W external/local_tsl/tsl/framework/cpu_allocator_impl.cc:83] Allocation of 46080368640 exceeds 10% of free system memory.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 1: val_loss improved from inf to 0.00000, saving model to downloads/data/state_dict/NRMS/weights\n",
      "37/37 [==============================] - 254s 7s/step - loss: 1.6025 - val_loss: 0.0000e+00\n"
     ]
    }
   ],
   "source": [
    "MODEL_NAME = \"NRMS\"\n",
    "LOG_DIR = f\"downloads/runs/{MODEL_NAME}\"\n",
    "MODEL_WEIGHTS = f\"downloads/data/state_dict/{MODEL_NAME}/weights\"\n",
    "\n",
    "# CALLBACKS\n",
    "tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=LOG_DIR, histogram_freq=1)\n",
    "early_stopping = tf.keras.callbacks.EarlyStopping(monitor=\"val_loss\", patience=2)\n",
    "modelcheckpoint = tf.keras.callbacks.ModelCheckpoint(\n",
    "    filepath=MODEL_WEIGHTS, save_best_only=True, save_weights_only=True, verbose=1\n",
    ")\n",
    "\n",
    "hparams_nrms.history_size = HISTORY_SIZE\n",
    "model = NRMSModel(\n",
    "    hparams=hparams_nrms,\n",
    "    word2vec_embedding=word2vec_embedding,\n",
    "    seed=42,\n",
    ")\n",
    "hist = model.model.fit(\n",
    "    train_dataloader,\n",
    "    validation_data=val_dataloader,\n",
    "    epochs=1,\n",
    "    callbacks=[tensorboard_callback, early_stopping, modelcheckpoint],\n",
    ")\n",
    "_ = model.model.load_weights(filepath=MODEL_WEIGHTS)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example how to compute some metrics:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "77/77 [==============================] - 132s 2s/step\n"
     ]
    }
   ],
   "source": [
    "pred_validation = model.scorer.predict(val_dataloader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Add the predictions to the dataframe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (2, 8)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>user_id</th><th>article_id_fixed</th><th>article_ids_inview</th><th>article_ids_clicked</th><th>impression_id</th><th>labels</th><th>scores</th><th>is_known_user</th></tr><tr><td>u32</td><td>list[i32]</td><td>list[i32]</td><td>list[i32]</td><td>u32</td><td>list[i8]</td><td>list[f64]</td><td>bool</td></tr></thead><tbody><tr><td>1148312</td><td>[9777704, 9775596, … 9779423]</td><td>[9785888, 9695098, … 9782808]</td><td>[9785404]</td><td>240468712</td><td>[0, 0, … 0]</td><td>[0.523742, 0.291834, … 0.428764]</td><td>false</td></tr><tr><td>2550273</td><td>[9778623, 9778444, … 9779648]</td><td>[9787176, 9786563, … 9787185]</td><td>[9787176]</td><td>259925945</td><td>[1, 0, … 0]</td><td>[0.431534, 0.44851, … 0.445333]</td><td>false</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (2, 8)\n",
       "┌─────────┬────────────┬────────────┬────────────┬────────────┬────────────┬───────────┬───────────┐\n",
       "│ user_id ┆ article_id ┆ article_id ┆ article_id ┆ impression ┆ labels     ┆ scores    ┆ is_known_ │\n",
       "│ ---     ┆ _fixed     ┆ s_inview   ┆ s_clicked  ┆ _id        ┆ ---        ┆ ---       ┆ user      │\n",
       "│ u32     ┆ ---        ┆ ---        ┆ ---        ┆ ---        ┆ list[i8]   ┆ list[f64] ┆ ---       │\n",
       "│         ┆ list[i32]  ┆ list[i32]  ┆ list[i32]  ┆ u32        ┆            ┆           ┆ bool      │\n",
       "╞═════════╪════════════╪════════════╪════════════╪════════════╪════════════╪═══════════╪═══════════╡\n",
       "│ 1148312 ┆ [9777704,  ┆ [9785888,  ┆ [9785404]  ┆ 240468712  ┆ [0, 0, …   ┆ [0.523742 ┆ false     │\n",
       "│         ┆ 9775596, … ┆ 9695098, … ┆            ┆            ┆ 0]         ┆ ,         ┆           │\n",
       "│         ┆ 9779423]   ┆ 9782808]   ┆            ┆            ┆            ┆ 0.291834, ┆           │\n",
       "│         ┆            ┆            ┆            ┆            ┆            ┆ …         ┆           │\n",
       "│         ┆            ┆            ┆            ┆            ┆            ┆ 0.428764] ┆           │\n",
       "│ 2550273 ┆ [9778623,  ┆ [9787176,  ┆ [9787176]  ┆ 259925945  ┆ [1, 0, …   ┆ [0.431534 ┆ false     │\n",
       "│         ┆ 9778444, … ┆ 9786563, … ┆            ┆            ┆ 0]         ┆ ,         ┆           │\n",
       "│         ┆ 9779648]   ┆ 9787185]   ┆            ┆            ┆            ┆ 0.44851,  ┆           │\n",
       "│         ┆            ┆            ┆            ┆            ┆            ┆ …         ┆           │\n",
       "│         ┆            ┆            ┆            ┆            ┆            ┆ 0.445333] ┆           │\n",
       "└─────────┴────────────┴────────────┴────────────┴────────────┴────────────┴───────────┴───────────┘"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_validation = add_prediction_scores(df_validation, pred_validation.tolist()).pipe(\n",
    "    add_known_user_column, known_users=df_train[DEFAULT_USER_COL]\n",
    ")\n",
    "df_validation.head(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compute metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<MetricEvaluator class>: \n",
       " {\n",
       "    \"auc\": 0.5298784949833476,\n",
       "    \"mrr\": 0.3242617779212621,\n",
       "    \"ndcg@5\": 0.36254779534285553,\n",
       "    \"ndcg@10\": 0.4420484974243122\n",
       "}"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "metrics = MetricEvaluator(\n",
    "    labels=df_validation[\"labels\"].to_list(),\n",
    "    predictions=df_validation[\"scores\"].to_list(),\n",
    "    metric_functions=[AucScore(), MrrScore(), NdcgScore(k=5), NdcgScore(k=10)],\n",
    ")\n",
    "metrics.evaluate()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Make submission file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (2, 9)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>user_id</th><th>article_id_fixed</th><th>article_ids_inview</th><th>article_ids_clicked</th><th>impression_id</th><th>labels</th><th>scores</th><th>is_known_user</th><th>ranked_scores</th></tr><tr><td>u32</td><td>list[i32]</td><td>list[i32]</td><td>list[i32]</td><td>u32</td><td>list[i8]</td><td>list[f64]</td><td>bool</td><td>list[i64]</td></tr></thead><tbody><tr><td>1148312</td><td>[9777704, 9775596, … 9779423]</td><td>[9785888, 9695098, … 9782808]</td><td>[9785404]</td><td>240468712</td><td>[0, 0, … 0]</td><td>[0.523742, 0.291834, … 0.428764]</td><td>false</td><td>[3, 29, … 17]</td></tr><tr><td>2550273</td><td>[9778623, 9778444, … 9779648]</td><td>[9787176, 9786563, … 9787185]</td><td>[9787176]</td><td>259925945</td><td>[1, 0, … 0]</td><td>[0.431534, 0.44851, … 0.445333]</td><td>false</td><td>[5, 2, … 3]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (2, 9)\n",
       "┌─────────┬────────────┬───────────┬───────────┬───┬───────────┬───────────┬───────────┬───────────┐\n",
       "│ user_id ┆ article_id ┆ article_i ┆ article_i ┆ … ┆ labels    ┆ scores    ┆ is_known_ ┆ ranked_sc │\n",
       "│ ---     ┆ _fixed     ┆ ds_inview ┆ ds_clicke ┆   ┆ ---       ┆ ---       ┆ user      ┆ ores      │\n",
       "│ u32     ┆ ---        ┆ ---       ┆ d         ┆   ┆ list[i8]  ┆ list[f64] ┆ ---       ┆ ---       │\n",
       "│         ┆ list[i32]  ┆ list[i32] ┆ ---       ┆   ┆           ┆           ┆ bool      ┆ list[i64] │\n",
       "│         ┆            ┆           ┆ list[i32] ┆   ┆           ┆           ┆           ┆           │\n",
       "╞═════════╪════════════╪═══════════╪═══════════╪═══╪═══════════╪═══════════╪═══════════╪═══════════╡\n",
       "│ 1148312 ┆ [9777704,  ┆ [9785888, ┆ [9785404] ┆ … ┆ [0, 0, …  ┆ [0.523742 ┆ false     ┆ [3, 29, … │\n",
       "│         ┆ 9775596, … ┆ 9695098,  ┆           ┆   ┆ 0]        ┆ ,         ┆           ┆ 17]       │\n",
       "│         ┆ 9779423]   ┆ …         ┆           ┆   ┆           ┆ 0.291834, ┆           ┆           │\n",
       "│         ┆            ┆ 9782808]  ┆           ┆   ┆           ┆ …         ┆           ┆           │\n",
       "│         ┆            ┆           ┆           ┆   ┆           ┆ 0.428764] ┆           ┆           │\n",
       "│ 2550273 ┆ [9778623,  ┆ [9787176, ┆ [9787176] ┆ … ┆ [1, 0, …  ┆ [0.431534 ┆ false     ┆ [5, 2, …  │\n",
       "│         ┆ 9778444, … ┆ 9786563,  ┆           ┆   ┆ 0]        ┆ ,         ┆           ┆ 3]        │\n",
       "│         ┆ 9779648]   ┆ …         ┆           ┆   ┆           ┆ 0.44851,  ┆           ┆           │\n",
       "│         ┆            ┆ 9787185]  ┆           ┆   ┆           ┆ …         ┆           ┆           │\n",
       "│         ┆            ┆           ┆           ┆   ┆           ┆ 0.445333] ┆           ┆           │\n",
       "└─────────┴────────────┴───────────┴───────────┴───┴───────────┴───────────┴───────────┴───────────┘"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_validation = df_validation.with_columns(\n",
    "    pl.col(\"scores\")\n",
    "    .map_elements(lambda x: list(rank_predictions_by_score(x)))\n",
    "    .alias(\"ranked_scores\")\n",
    ")\n",
    "df_validation.head(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is using the validation, simply add the testset to your flow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2446it [00:00, 13918.82it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Zipping downloads/predictions.txt to downloads/predictions.zip\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "write_submission_file(\n",
    "    impression_ids=df_validation[DEFAULT_IMPRESSION_ID_COL],\n",
    "    prediction_scores=df_validation[\"ranked_scores\"],\n",
    "    path=\"downloads/predictions.txt\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# DONE 🚀"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}