prepare-dataset.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import utils\n",
    "import polars as pl\n",
    "\n",
    "comments_df = utils.augmented_comments()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "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: (24_134_648, 10)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>id</th><th>by</th><th>time</th><th>text</th><th>parent</th><th>top_level_parent</th><th>kids</th><th>siblings_count</th><th>sibling_rank</th><th>nested_level</th></tr><tr><td>i64</td><td>str</td><td>datetime[μs]</td><td>str</td><td>i64</td><td>i64</td><td>list[i64]</td><td>u32</td><td>i64</td><td>i32</td></tr></thead><tbody><tr><td>194</td><td>&quot;jmzachary&quot;</td><td>2007-02-20 22:33:51</td><td>&quot;Thanks for the rationale. I&#x27;m …</td><td>189</td><td>189</td><td>[205, 422, 199]</td><td>20</td><td>13</td><td>0</td></tr><tr><td>195</td><td>&quot;jdroid&quot;</td><td>2007-02-20 22:36:52</td><td>&quot;You&#x27;ve filled a hole reddit wa…</td><td>189</td><td>189</td><td>[259]</td><td>20</td><td>3</td><td>0</td></tr><tr><td>199</td><td>&quot;Zak&quot;</td><td>2007-02-20 22:48:33</td><td>&quot;I don&#x27;t think the fact that th…</td><td>194</td><td>189</td><td>[1644, 1897]</td><td>3</td><td>3</td><td>1</td></tr><tr><td>205</td><td>&quot;ninwa&quot;</td><td>2007-02-20 23:30:23</td><td>&quot;Really? I was most interested …</td><td>194</td><td>189</td><td>[210, 209]</td><td>3</td><td>1</td><td>1</td></tr><tr><td>209</td><td>&quot;ninwa&quot;</td><td>2007-02-20 23:41:34</td><td>&quot;This comment added through the…</td><td>205</td><td>189</td><td>null</td><td>2</td><td>2</td><td>2</td></tr><tr><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td></tr><tr><td>41813369</td><td>&quot;whall6&quot;</td><td>2024-10-11 20:25:08</td><td>&quot;If you truly only have 5 minut…</td><td>41812596</td><td>41811263</td><td>null</td><td>1</td><td>1</td><td>3</td></tr><tr><td>41813371</td><td>&quot;julianeon&quot;</td><td>2024-10-11 20:25:25</td><td>&quot;But consider the tradeoff: it&#x27;…</td><td>41811539</td><td>41811263</td><td>null</td><td>13</td><td>1</td><td>1</td></tr><tr><td>41813375</td><td>&quot;throw0101c&quot;</td><td>2024-10-11 20:25:51</td><td>&quot;&gt; &lt;i&gt;There is so much coal... …</td><td>41811328</td><td>41807681</td><td>null</td><td>1</td><td>1</td><td>3</td></tr><tr><td>41813381</td><td>&quot;marcosdumay&quot;</td><td>2024-10-11 20:26:21</td><td>&quot;Add forced sedentarism into th…</td><td>41812891</td><td>41811263</td><td>null</td><td>1</td><td>1</td><td>2</td></tr><tr><td>41813383</td><td>&quot;davio&quot;</td><td>2024-10-11 20:26:28</td><td>&quot;hims sells the generic version…</td><td>41813102</td><td>41811263</td><td>null</td><td>1</td><td>1</td><td>3</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (24_134_648, 10)\n",
       "┌──────────┬───────────┬───────────┬───────────┬───┬───────────┬───────────┬───────────┬───────────┐\n",
       "│ id       ┆ by        ┆ time      ┆ text      ┆ … ┆ kids      ┆ siblings_ ┆ sibling_r ┆ nested_le │\n",
       "│ ---      ┆ ---       ┆ ---       ┆ ---       ┆   ┆ ---       ┆ count     ┆ ank       ┆ vel       │\n",
       "│ i64      ┆ str       ┆ datetime[ ┆ str       ┆   ┆ list[i64] ┆ ---       ┆ ---       ┆ ---       │\n",
       "│          ┆           ┆ μs]       ┆           ┆   ┆           ┆ u32       ┆ i64       ┆ i32       │\n",
       "╞══════════╪═══════════╪═══════════╪═══════════╪═══╪═══════════╪═══════════╪═══════════╪═══════════╡\n",
       "│ 194      ┆ jmzachary ┆ 2007-02-2 ┆ Thanks    ┆ … ┆ [205,     ┆ 20        ┆ 13        ┆ 0         │\n",
       "│          ┆           ┆ 0         ┆ for the   ┆   ┆ 422, 199] ┆           ┆           ┆           │\n",
       "│          ┆           ┆ 22:33:51  ┆ rationale ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ . I'm …   ┆   ┆           ┆           ┆           ┆           │\n",
       "│ 195      ┆ jdroid    ┆ 2007-02-2 ┆ You've    ┆ … ┆ [259]     ┆ 20        ┆ 3         ┆ 0         │\n",
       "│          ┆           ┆ 0         ┆ filled a  ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆ 22:36:52  ┆ hole      ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ reddit    ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ wa…       ┆   ┆           ┆           ┆           ┆           │\n",
       "│ 199      ┆ Zak       ┆ 2007-02-2 ┆ I don't   ┆ … ┆ [1644,    ┆ 3         ┆ 3         ┆ 1         │\n",
       "│          ┆           ┆ 0         ┆ think the ┆   ┆ 1897]     ┆           ┆           ┆           │\n",
       "│          ┆           ┆ 22:48:33  ┆ fact that ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ th…       ┆   ┆           ┆           ┆           ┆           │\n",
       "│ 205      ┆ ninwa     ┆ 2007-02-2 ┆ Really? I ┆ … ┆ [210,     ┆ 3         ┆ 1         ┆ 1         │\n",
       "│          ┆           ┆ 0         ┆ was most  ┆   ┆ 209]      ┆           ┆           ┆           │\n",
       "│          ┆           ┆ 23:30:23  ┆ intereste ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ d …       ┆   ┆           ┆           ┆           ┆           │\n",
       "│ 209      ┆ ninwa     ┆ 2007-02-2 ┆ This      ┆ … ┆ null      ┆ 2         ┆ 2         ┆ 2         │\n",
       "│          ┆           ┆ 0         ┆ comment   ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆ 23:41:34  ┆ added     ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ through   ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ the…      ┆   ┆           ┆           ┆           ┆           │\n",
       "│ …        ┆ …         ┆ …         ┆ …         ┆ … ┆ …         ┆ …         ┆ …         ┆ …         │\n",
       "│ 41813369 ┆ whall6    ┆ 2024-10-1 ┆ If you    ┆ … ┆ null      ┆ 1         ┆ 1         ┆ 3         │\n",
       "│          ┆           ┆ 1         ┆ truly     ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆ 20:25:08  ┆ only have ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ 5 minut…  ┆   ┆           ┆           ┆           ┆           │\n",
       "│ 41813371 ┆ julianeon ┆ 2024-10-1 ┆ But       ┆ … ┆ null      ┆ 13        ┆ 1         ┆ 1         │\n",
       "│          ┆           ┆ 1         ┆ consider  ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆ 20:25:25  ┆ the       ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ tradeoff: ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ it'…      ┆   ┆           ┆           ┆           ┆           │\n",
       "│ 41813375 ┆ throw0101 ┆ 2024-10-1 ┆ >         ┆ … ┆ null      ┆ 1         ┆ 1         ┆ 3         │\n",
       "│          ┆ c         ┆ 1         ┆ <i>There  ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆ 20:25:51  ┆ is so     ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ much      ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ coal... … ┆   ┆           ┆           ┆           ┆           │\n",
       "│ 41813381 ┆ marcosdum ┆ 2024-10-1 ┆ Add       ┆ … ┆ null      ┆ 1         ┆ 1         ┆ 2         │\n",
       "│          ┆ ay        ┆ 1         ┆ forced    ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆ 20:26:21  ┆ sedentari ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ sm into   ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ th…       ┆   ┆           ┆           ┆           ┆           │\n",
       "│ 41813383 ┆ davio     ┆ 2024-10-1 ┆ hims      ┆ … ┆ null      ┆ 1         ┆ 1         ┆ 3         │\n",
       "│          ┆           ┆ 1         ┆ sells the ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆ 20:26:28  ┆ generic   ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ version…  ┆   ┆           ┆           ┆           ┆           │\n",
       "└──────────┴───────────┴───────────┴───────────┴───┴───────────┴───────────┴───────────┴───────────┘"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = utils.full_dataset()\n",
    "\n",
    "stories_with_100_points = df.filter(pl.col(\"score\") >= 100)\n",
    "\n",
    "# filter comments_df to only include comments whose top_level_parent is in the stories_with_100_points dataframe\n",
    "comments_df = comments_df.filter(\n",
    "    pl.col(\"top_level_parent\").is_in(stories_with_100_points[\"id\"])\n",
    ")\n",
    "\n",
    "comments_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Iteration 1, current length: 533314\n",
      "Iteration 2, current length: 1614377\n",
      "Iteration 3, current length: 2828828\n",
      "Iteration 4, current length: 3874258\n",
      "Iteration 5, current length: 4636895\n",
      "Iteration 6, current length: 5139157\n"
     ]
    },
    {
     "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: (5_445_792, 10)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>id</th><th>by</th><th>time</th><th>text</th><th>parent</th><th>top_level_parent</th><th>kids</th><th>siblings_count</th><th>sibling_rank</th><th>nested_level</th></tr><tr><td>i64</td><td>str</td><td>datetime[μs]</td><td>str</td><td>i64</td><td>i64</td><td>list[i64]</td><td>u32</td><td>i64</td><td>i32</td></tr></thead><tbody><tr><td>195</td><td>&quot;jdroid&quot;</td><td>2007-02-20 22:36:52</td><td>&quot;You&#x27;ve filled a hole reddit wa…</td><td>189</td><td>189</td><td>[259]</td><td>20</td><td>3</td><td>0</td></tr><tr><td>259</td><td>&quot;whatsreal&quot;</td><td>2007-02-21 06:18:50</td><td>&quot;HaHa!  Yes thank you Paul, I w…</td><td>195</td><td>189</td><td>[453, 7198]</td><td>1</td><td>1</td><td>1</td></tr><tr><td>287</td><td>&quot;ced&quot;</td><td>2007-02-21 09:08:16</td><td>&quot;Since community-building is pa…</td><td>189</td><td>189</td><td>[365]</td><td>20</td><td>2</td><td>0</td></tr><tr><td>353</td><td>&quot;jhenzie&quot;</td><td>2007-02-21 18:47:34</td><td>&quot;Paul, \r\n",
       "\r\n",
       "can we play with arc…</td><td>189</td><td>189</td><td>[7199]</td><td>20</td><td>1</td><td>0</td></tr><tr><td>365</td><td>&quot;jdroid&quot;</td><td>2007-02-21 19:25:14</td><td>&quot;If I bought yspace.com would y…</td><td>287</td><td>189</td><td>null</td><td>1</td><td>1</td><td>1</td></tr><tr><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td></tr><tr><td>41813352</td><td>&quot;dullcrisp&quot;</td><td>2024-10-11 20:22:59</td><td>&quot;A &lt;i&gt;scrappy&lt;/i&gt; marketing tea…</td><td>41813266</td><td>41811263</td><td>null</td><td>1</td><td>1</td><td>3</td></tr><tr><td>41813354</td><td>&quot;ragnese&quot;</td><td>2024-10-11 20:23:06</td><td>&quot;It was only true that Chrome w…</td><td>41813112</td><td>41809698</td><td>null</td><td>1</td><td>1</td><td>4</td></tr><tr><td>41813362</td><td>&quot;layer8&quot;</td><td>2024-10-11 20:23:58</td><td>&quot;It’s not simply a limitation i…</td><td>41812523</td><td>41812523</td><td>null</td><td>31</td><td>3</td><td>0</td></tr><tr><td>41813363</td><td>&quot;nightski&quot;</td><td>2024-10-11 20:24:07</td><td>&quot;Nothing &quot;needs&quot; to happen. Peo…</td><td>41812696</td><td>41811263</td><td>null</td><td>1</td><td>1</td><td>5</td></tr><tr><td>41813368</td><td>null</td><td>2024-10-11 20:24:58</td><td>null</td><td>41813116</td><td>41811263</td><td>null</td><td>1</td><td>1</td><td>3</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5_445_792, 10)\n",
       "┌──────────┬───────────┬───────────┬───────────┬───┬───────────┬───────────┬───────────┬───────────┐\n",
       "│ id       ┆ by        ┆ time      ┆ text      ┆ … ┆ kids      ┆ siblings_ ┆ sibling_r ┆ nested_le │\n",
       "│ ---      ┆ ---       ┆ ---       ┆ ---       ┆   ┆ ---       ┆ count     ┆ ank       ┆ vel       │\n",
       "│ i64      ┆ str       ┆ datetime[ ┆ str       ┆   ┆ list[i64] ┆ ---       ┆ ---       ┆ ---       │\n",
       "│          ┆           ┆ μs]       ┆           ┆   ┆           ┆ u32       ┆ i64       ┆ i32       │\n",
       "╞══════════╪═══════════╪═══════════╪═══════════╪═══╪═══════════╪═══════════╪═══════════╪═══════════╡\n",
       "│ 195      ┆ jdroid    ┆ 2007-02-2 ┆ You've    ┆ … ┆ [259]     ┆ 20        ┆ 3         ┆ 0         │\n",
       "│          ┆           ┆ 0         ┆ filled a  ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆ 22:36:52  ┆ hole      ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ reddit    ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ wa…       ┆   ┆           ┆           ┆           ┆           │\n",
       "│ 259      ┆ whatsreal ┆ 2007-02-2 ┆ HaHa!     ┆ … ┆ [453,     ┆ 1         ┆ 1         ┆ 1         │\n",
       "│          ┆           ┆ 1         ┆ Yes thank ┆   ┆ 7198]     ┆           ┆           ┆           │\n",
       "│          ┆           ┆ 06:18:50  ┆ you Paul, ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ I w…      ┆   ┆           ┆           ┆           ┆           │\n",
       "│ 287      ┆ ced       ┆ 2007-02-2 ┆ Since com ┆ … ┆ [365]     ┆ 20        ┆ 2         ┆ 0         │\n",
       "│          ┆           ┆ 1         ┆ munity-bu ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆ 09:08:16  ┆ ilding is ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ pa…       ┆   ┆           ┆           ┆           ┆           │\n",
       "│ 353      ┆ jhenzie   ┆ 2007-02-2 ┆ Paul, \n",
       "   ┆ … ┆ [7199]    ┆ 20        ┆ 1         ┆ 0         │\n",
       "│          ┆           ┆ 1         ┆ \n",
       "         ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆ 18:47:34  ┆ can we    ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ play with ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ arc…      ┆   ┆           ┆           ┆           ┆           │\n",
       "│ 365      ┆ jdroid    ┆ 2007-02-2 ┆ If I      ┆ … ┆ null      ┆ 1         ┆ 1         ┆ 1         │\n",
       "│          ┆           ┆ 1         ┆ bought    ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆ 19:25:14  ┆ yspace.co ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ m would   ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ y…        ┆   ┆           ┆           ┆           ┆           │\n",
       "│ …        ┆ …         ┆ …         ┆ …         ┆ … ┆ …         ┆ …         ┆ …         ┆ …         │\n",
       "│ 41813352 ┆ dullcrisp ┆ 2024-10-1 ┆ A <i>scra ┆ … ┆ null      ┆ 1         ┆ 1         ┆ 3         │\n",
       "│          ┆           ┆ 1         ┆ ppy</i>   ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆ 20:22:59  ┆ marketing ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ tea…      ┆   ┆           ┆           ┆           ┆           │\n",
       "│ 41813354 ┆ ragnese   ┆ 2024-10-1 ┆ It was    ┆ … ┆ null      ┆ 1         ┆ 1         ┆ 4         │\n",
       "│          ┆           ┆ 1         ┆ only true ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆ 20:23:06  ┆ that      ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ Chrome w… ┆   ┆           ┆           ┆           ┆           │\n",
       "│ 41813362 ┆ layer8    ┆ 2024-10-1 ┆ It’s not  ┆ … ┆ null      ┆ 31        ┆ 3         ┆ 0         │\n",
       "│          ┆           ┆ 1         ┆ simply a  ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆ 20:23:58  ┆ limitatio ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ n i…      ┆   ┆           ┆           ┆           ┆           │\n",
       "│ 41813363 ┆ nightski  ┆ 2024-10-1 ┆ Nothing   ┆ … ┆ null      ┆ 1         ┆ 1         ┆ 5         │\n",
       "│          ┆           ┆ 1         ┆ \"needs\"   ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆ 20:24:07  ┆ to        ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ happen.   ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ Peo…      ┆   ┆           ┆           ┆           ┆           │\n",
       "│ 41813368 ┆ null      ┆ 2024-10-1 ┆ null      ┆ … ┆ null      ┆ 1         ┆ 1         ┆ 3         │\n",
       "│          ┆           ┆ 1         ┆           ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆ 20:24:58  ┆           ┆   ┆           ┆           ┆           ┆           │\n",
       "└──────────┴───────────┴───────────┴───────────┴───┴───────────┴───────────┴───────────┴───────────┘"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Let's limit to comments for which all of their parents have a sibling rank <=\n",
    "# 3. If a comment isn't in the top 3 of its parent it's less likely to have been\n",
    "# read and upvoted by enough people to form a representative sample.\n",
    "\n",
    "# We'll iteratively add comments with sibling_rank <= 3 up to 6 iterations.\n",
    "\n",
    "filtered_comments_df = comments_df.filter(\n",
    "    (pl.col(\"sibling_rank\") <= 3) & (pl.col(\"parent\").eq(pl.col(\"top_level_parent\")))\n",
    ")\n",
    "\n",
    "for iteration in range(6):\n",
    "    print(f\"Iteration {iteration + 1}, current length: {len(filtered_comments_df)}\")\n",
    "    new_filtered_comments_df = comments_df.filter(\n",
    "        (\n",
    "            (pl.col(\"sibling_rank\") <= 3)\n",
    "            & (pl.col(\"parent\").is_in(filtered_comments_df[\"id\"]))\n",
    "        )\n",
    "        | (pl.col(\"id\").is_in(filtered_comments_df[\"id\"]))\n",
    "    )\n",
    "\n",
    "    if len(new_filtered_comments_df) == len(filtered_comments_df):\n",
    "        break\n",
    "\n",
    "    filtered_comments_df = new_filtered_comments_df\n",
    "\n",
    "comments_df = filtered_comments_df\n",
    "\n",
    "del filtered_comments_df\n",
    "comments_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "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: (1_109_206, 10)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>id</th><th>by</th><th>time</th><th>text</th><th>parent</th><th>top_level_parent</th><th>kids</th><th>siblings_count</th><th>sibling_rank</th><th>nested_level</th></tr><tr><td>i64</td><td>str</td><td>datetime[μs]</td><td>str</td><td>i64</td><td>i64</td><td>list[i64]</td><td>u32</td><td>i64</td><td>i32</td></tr></thead><tbody><tr><td>195</td><td>&quot;jdroid&quot;</td><td>2007-02-20 22:36:52</td><td>&quot;You&#x27;ve filled a hole reddit wa…</td><td>189</td><td>189</td><td>[259]</td><td>20</td><td>3</td><td>0</td></tr><tr><td>287</td><td>&quot;ced&quot;</td><td>2007-02-21 09:08:16</td><td>&quot;Since community-building is pa…</td><td>189</td><td>189</td><td>[365]</td><td>20</td><td>2</td><td>0</td></tr><tr><td>353</td><td>&quot;jhenzie&quot;</td><td>2007-02-21 18:47:34</td><td>&quot;Paul, \r\n",
       "\r\n",
       "can we play with arc…</td><td>189</td><td>189</td><td>[7199]</td><td>20</td><td>1</td><td>0</td></tr><tr><td>27552</td><td>&quot;pg&quot;</td><td>2007-06-12 19:02:46</td><td>&quot;Don&#x27;t worry, I&#x27;m not going to …</td><td>27550</td><td>27550</td><td>[28315, 27559, … 28523]</td><td>7</td><td>1</td><td>0</td></tr><tr><td>27560</td><td>&quot;danielha&quot;</td><td>2007-06-12 19:15:56</td><td>&quot;Prediction: this submission re…</td><td>27550</td><td>27550</td><td>[28098]</td><td>7</td><td>2</td><td>0</td></tr><tr><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td></tr><tr><td>41812977</td><td>&quot;shadowgovt&quot;</td><td>2024-10-11 19:50:42</td><td>&quot;Their incentive is really to m…</td><td>41810118</td><td>41809698</td><td>null</td><td>7</td><td>-1</td><td>2</td></tr><tr><td>41813078</td><td>&quot;itronitron&quot;</td><td>2024-10-11 19:59:06</td><td>&quot;This is the type of comment th…</td><td>41812493</td><td>41811263</td><td>[41813116, 41813266]</td><td>9</td><td>3</td><td>1</td></tr><tr><td>41813079</td><td>&quot;ragnese&quot;</td><td>2024-10-11 19:59:16</td><td>&quot;The vast, &lt;i&gt;vast&lt;/i&gt;, majorit…</td><td>41809962</td><td>41809698</td><td>null</td><td>12</td><td>-1</td><td>1</td></tr><tr><td>41813239</td><td>&quot;glenstein&quot;</td><td>2024-10-11 20:13:25</td><td>&quot;I think this is such a helpful…</td><td>41812493</td><td>41811263</td><td>null</td><td>9</td><td>2</td><td>1</td></tr><tr><td>41813277</td><td>&quot;ohples&quot;</td><td>2024-10-11 20:16:05</td><td>&quot;But will it support Login.gov?&quot;</td><td>41737653</td><td>41737653</td><td>null</td><td>11</td><td>-1</td><td>0</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (1_109_206, 10)\n",
       "┌──────────┬───────────┬───────────┬───────────┬───┬───────────┬───────────┬───────────┬───────────┐\n",
       "│ id       ┆ by        ┆ time      ┆ text      ┆ … ┆ kids      ┆ siblings_ ┆ sibling_r ┆ nested_le │\n",
       "│ ---      ┆ ---       ┆ ---       ┆ ---       ┆   ┆ ---       ┆ count     ┆ ank       ┆ vel       │\n",
       "│ i64      ┆ str       ┆ datetime[ ┆ str       ┆   ┆ list[i64] ┆ ---       ┆ ---       ┆ ---       │\n",
       "│          ┆           ┆ μs]       ┆           ┆   ┆           ┆ u32       ┆ i64       ┆ i32       │\n",
       "╞══════════╪═══════════╪═══════════╪═══════════╪═══╪═══════════╪═══════════╪═══════════╪═══════════╡\n",
       "│ 195      ┆ jdroid    ┆ 2007-02-2 ┆ You've    ┆ … ┆ [259]     ┆ 20        ┆ 3         ┆ 0         │\n",
       "│          ┆           ┆ 0         ┆ filled a  ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆ 22:36:52  ┆ hole      ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ reddit    ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ wa…       ┆   ┆           ┆           ┆           ┆           │\n",
       "│ 287      ┆ ced       ┆ 2007-02-2 ┆ Since com ┆ … ┆ [365]     ┆ 20        ┆ 2         ┆ 0         │\n",
       "│          ┆           ┆ 1         ┆ munity-bu ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆ 09:08:16  ┆ ilding is ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ pa…       ┆   ┆           ┆           ┆           ┆           │\n",
       "│ 353      ┆ jhenzie   ┆ 2007-02-2 ┆ Paul, \n",
       "   ┆ … ┆ [7199]    ┆ 20        ┆ 1         ┆ 0         │\n",
       "│          ┆           ┆ 1         ┆ \n",
       "         ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆ 18:47:34  ┆ can we    ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ play with ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ arc…      ┆   ┆           ┆           ┆           ┆           │\n",
       "│ 27552    ┆ pg        ┆ 2007-06-1 ┆ Don't     ┆ … ┆ [28315,   ┆ 7         ┆ 1         ┆ 0         │\n",
       "│          ┆           ┆ 2         ┆ worry,    ┆   ┆ 27559, …  ┆           ┆           ┆           │\n",
       "│          ┆           ┆ 19:02:46  ┆ I'm not   ┆   ┆ 28523]    ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ going to  ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ …         ┆   ┆           ┆           ┆           ┆           │\n",
       "│ 27560    ┆ danielha  ┆ 2007-06-1 ┆ Predictio ┆ … ┆ [28098]   ┆ 7         ┆ 2         ┆ 0         │\n",
       "│          ┆           ┆ 2         ┆ n: this   ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆ 19:15:56  ┆ submissio ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ n re…     ┆   ┆           ┆           ┆           ┆           │\n",
       "│ …        ┆ …         ┆ …         ┆ …         ┆ … ┆ …         ┆ …         ┆ …         ┆ …         │\n",
       "│ 41812977 ┆ shadowgov ┆ 2024-10-1 ┆ Their     ┆ … ┆ null      ┆ 7         ┆ -1        ┆ 2         │\n",
       "│          ┆ t         ┆ 1         ┆ incentive ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆ 19:50:42  ┆ is really ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ to m…     ┆   ┆           ┆           ┆           ┆           │\n",
       "│ 41813078 ┆ itronitro ┆ 2024-10-1 ┆ This is   ┆ … ┆ [41813116 ┆ 9         ┆ 3         ┆ 1         │\n",
       "│          ┆ n         ┆ 1         ┆ the type  ┆   ┆ ,         ┆           ┆           ┆           │\n",
       "│          ┆           ┆ 19:59:06  ┆ of        ┆   ┆ 41813266] ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ comment   ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ th…       ┆   ┆           ┆           ┆           ┆           │\n",
       "│ 41813079 ┆ ragnese   ┆ 2024-10-1 ┆ The vast, ┆ … ┆ null      ┆ 12        ┆ -1        ┆ 1         │\n",
       "│          ┆           ┆ 1         ┆ <i>vast</ ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆ 19:59:16  ┆ i>,       ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ majorit…  ┆   ┆           ┆           ┆           ┆           │\n",
       "│ 41813239 ┆ glenstein ┆ 2024-10-1 ┆ I think   ┆ … ┆ null      ┆ 9         ┆ 2         ┆ 1         │\n",
       "│          ┆           ┆ 1         ┆ this is   ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆ 20:13:25  ┆ such a    ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ helpful…  ┆   ┆           ┆           ┆           ┆           │\n",
       "│ 41813277 ┆ ohples    ┆ 2024-10-1 ┆ But will  ┆ … ┆ null      ┆ 11        ┆ -1        ┆ 0         │\n",
       "│          ┆           ┆ 1         ┆ it        ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆ 20:16:05  ┆ support   ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ Login.gov ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ ?         ┆   ┆           ┆           ┆           ┆           │\n",
       "└──────────┴───────────┴───────────┴───────────┴───┴───────────┴───────────┴───────────┴───────────┘"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# limit all comments to ones that have at least 5 siblings and fewer than 25\n",
    "# siblings. If there are more than 25 I don't trust the peoples' votes since\n",
    "# they probably haven't had time to read all of them.\n",
    "comments_df = comments_df.filter(\n",
    "    (pl.col(\"siblings_count\") >= 5) & (pl.col(\"siblings_count\") < 25)\n",
    ")\n",
    "\n",
    "comments_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "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: (1_104_093, 10)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>id</th><th>by</th><th>time</th><th>text</th><th>parent</th><th>top_level_parent</th><th>kids</th><th>siblings_count</th><th>sibling_rank</th><th>nested_level</th></tr><tr><td>i64</td><td>str</td><td>datetime[μs]</td><td>str</td><td>i64</td><td>i64</td><td>list[i64]</td><td>u32</td><td>i64</td><td>i32</td></tr></thead><tbody><tr><td>195</td><td>&quot;jdroid&quot;</td><td>2007-02-20 22:36:52</td><td>&quot;You&#x27;ve filled a hole reddit wa…</td><td>189</td><td>189</td><td>[259]</td><td>20</td><td>3</td><td>0</td></tr><tr><td>287</td><td>&quot;ced&quot;</td><td>2007-02-21 09:08:16</td><td>&quot;Since community-building is pa…</td><td>189</td><td>189</td><td>[365]</td><td>20</td><td>2</td><td>0</td></tr><tr><td>353</td><td>&quot;jhenzie&quot;</td><td>2007-02-21 18:47:34</td><td>&quot;Paul, \r\n",
       "\r\n",
       "can we play with arc…</td><td>189</td><td>189</td><td>[7199]</td><td>20</td><td>1</td><td>0</td></tr><tr><td>27552</td><td>&quot;pg&quot;</td><td>2007-06-12 19:02:46</td><td>&quot;Don&#x27;t worry, I&#x27;m not going to …</td><td>27550</td><td>27550</td><td>[28315, 27559, … 28523]</td><td>7</td><td>1</td><td>0</td></tr><tr><td>27560</td><td>&quot;danielha&quot;</td><td>2007-06-12 19:15:56</td><td>&quot;Prediction: this submission re…</td><td>27550</td><td>27550</td><td>[28098]</td><td>7</td><td>2</td><td>0</td></tr><tr><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td></tr><tr><td>41812565</td><td>&quot;vessenes&quot;</td><td>2024-10-11 19:14:29</td><td>&quot;Tirzepatide and Semaglutide ar…</td><td>41812493</td><td>41811263</td><td>[41812815]</td><td>9</td><td>1</td><td>1</td></tr><tr><td>41812643</td><td>&quot;squidlogic&quot;</td><td>2024-10-11 19:23:13</td><td>&quot;In my experience fitness is le…</td><td>41812339</td><td>41811263</td><td>[41812886, 41812696, … 41813092]</td><td>6</td><td>1</td><td>3</td></tr><tr><td>41812717</td><td>&quot;jrflowers&quot;</td><td>2024-10-11 19:29:53</td><td>&quot;&gt; People in the 50s weren&#x27;t sl…</td><td>41812339</td><td>41811263</td><td>[41813315, 41812727]</td><td>6</td><td>3</td><td>3</td></tr><tr><td>41813078</td><td>&quot;itronitron&quot;</td><td>2024-10-11 19:59:06</td><td>&quot;This is the type of comment th…</td><td>41812493</td><td>41811263</td><td>[41813116, 41813266]</td><td>9</td><td>3</td><td>1</td></tr><tr><td>41813239</td><td>&quot;glenstein&quot;</td><td>2024-10-11 20:13:25</td><td>&quot;I think this is such a helpful…</td><td>41812493</td><td>41811263</td><td>null</td><td>9</td><td>2</td><td>1</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (1_104_093, 10)\n",
       "┌──────────┬───────────┬───────────┬───────────┬───┬───────────┬───────────┬───────────┬───────────┐\n",
       "│ id       ┆ by        ┆ time      ┆ text      ┆ … ┆ kids      ┆ siblings_ ┆ sibling_r ┆ nested_le │\n",
       "│ ---      ┆ ---       ┆ ---       ┆ ---       ┆   ┆ ---       ┆ count     ┆ ank       ┆ vel       │\n",
       "│ i64      ┆ str       ┆ datetime[ ┆ str       ┆   ┆ list[i64] ┆ ---       ┆ ---       ┆ ---       │\n",
       "│          ┆           ┆ μs]       ┆           ┆   ┆           ┆ u32       ┆ i64       ┆ i32       │\n",
       "╞══════════╪═══════════╪═══════════╪═══════════╪═══╪═══════════╪═══════════╪═══════════╪═══════════╡\n",
       "│ 195      ┆ jdroid    ┆ 2007-02-2 ┆ You've    ┆ … ┆ [259]     ┆ 20        ┆ 3         ┆ 0         │\n",
       "│          ┆           ┆ 0         ┆ filled a  ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆ 22:36:52  ┆ hole      ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ reddit    ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ wa…       ┆   ┆           ┆           ┆           ┆           │\n",
       "│ 287      ┆ ced       ┆ 2007-02-2 ┆ Since com ┆ … ┆ [365]     ┆ 20        ┆ 2         ┆ 0         │\n",
       "│          ┆           ┆ 1         ┆ munity-bu ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆ 09:08:16  ┆ ilding is ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ pa…       ┆   ┆           ┆           ┆           ┆           │\n",
       "│ 353      ┆ jhenzie   ┆ 2007-02-2 ┆ Paul, \n",
       "   ┆ … ┆ [7199]    ┆ 20        ┆ 1         ┆ 0         │\n",
       "│          ┆           ┆ 1         ┆ \n",
       "         ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆ 18:47:34  ┆ can we    ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ play with ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ arc…      ┆   ┆           ┆           ┆           ┆           │\n",
       "│ 27552    ┆ pg        ┆ 2007-06-1 ┆ Don't     ┆ … ┆ [28315,   ┆ 7         ┆ 1         ┆ 0         │\n",
       "│          ┆           ┆ 2         ┆ worry,    ┆   ┆ 27559, …  ┆           ┆           ┆           │\n",
       "│          ┆           ┆ 19:02:46  ┆ I'm not   ┆   ┆ 28523]    ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ going to  ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ …         ┆   ┆           ┆           ┆           ┆           │\n",
       "│ 27560    ┆ danielha  ┆ 2007-06-1 ┆ Predictio ┆ … ┆ [28098]   ┆ 7         ┆ 2         ┆ 0         │\n",
       "│          ┆           ┆ 2         ┆ n: this   ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆ 19:15:56  ┆ submissio ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ n re…     ┆   ┆           ┆           ┆           ┆           │\n",
       "│ …        ┆ …         ┆ …         ┆ …         ┆ … ┆ …         ┆ …         ┆ …         ┆ …         │\n",
       "│ 41812565 ┆ vessenes  ┆ 2024-10-1 ┆ Tirzepati ┆ … ┆ [41812815 ┆ 9         ┆ 1         ┆ 1         │\n",
       "│          ┆           ┆ 1         ┆ de and    ┆   ┆ ]         ┆           ┆           ┆           │\n",
       "│          ┆           ┆ 19:14:29  ┆ Semagluti ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ de ar…    ┆   ┆           ┆           ┆           ┆           │\n",
       "│ 41812643 ┆ squidlogi ┆ 2024-10-1 ┆ In my exp ┆ … ┆ [41812886 ┆ 6         ┆ 1         ┆ 3         │\n",
       "│          ┆ c         ┆ 1         ┆ erience   ┆   ┆ ,         ┆           ┆           ┆           │\n",
       "│          ┆           ┆ 19:23:13  ┆ fitness   ┆   ┆ 41812696, ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ is le…    ┆   ┆ …         ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆           ┆   ┆ 4181309…  ┆           ┆           ┆           │\n",
       "│ 41812717 ┆ jrflowers ┆ 2024-10-1 ┆ > People  ┆ … ┆ [41813315 ┆ 6         ┆ 3         ┆ 3         │\n",
       "│          ┆           ┆ 1         ┆ in the    ┆   ┆ ,         ┆           ┆           ┆           │\n",
       "│          ┆           ┆ 19:29:53  ┆ 50s       ┆   ┆ 41812727] ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ weren't   ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ sl…       ┆   ┆           ┆           ┆           ┆           │\n",
       "│ 41813078 ┆ itronitro ┆ 2024-10-1 ┆ This is   ┆ … ┆ [41813116 ┆ 9         ┆ 3         ┆ 1         │\n",
       "│          ┆ n         ┆ 1         ┆ the type  ┆   ┆ ,         ┆           ┆           ┆           │\n",
       "│          ┆           ┆ 19:59:06  ┆ of        ┆   ┆ 41813266] ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ comment   ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ th…       ┆   ┆           ┆           ┆           ┆           │\n",
       "│ 41813239 ┆ glenstein ┆ 2024-10-1 ┆ I think   ┆ … ┆ null      ┆ 9         ┆ 2         ┆ 1         │\n",
       "│          ┆           ┆ 1         ┆ this is   ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆ 20:13:25  ┆ such a    ┆   ┆           ┆           ┆           ┆           │\n",
       "│          ┆           ┆           ┆ helpful…  ┆   ┆           ┆           ┆           ┆           │\n",
       "└──────────┴───────────┴───────────┴───────────┴───┴───────────┴───────────┴───────────┴───────────┘"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# filter comments that have non-null text and text length > 10\n",
    "comments_df = comments_df.filter(\n",
    "    (pl.col(\"text\").is_not_null())\n",
    "    & (pl.col(\"text\").str.len_chars() > 10)\n",
    "    & (pl.col(\"sibling_rank\") > 0)\n",
    ")\n",
    "\n",
    "comments_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Processing groups:   0%|          | 1044/1104093 [00:00<01:45, 10439.96it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Processing groups: 100%|██████████| 1104093/1104093 [00:01<00:00, 563691.17it/s]\n"
     ]
    },
    {
     "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: (371_411, 5)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>chosen</th><th>rejected</th><th>chosen_rank</th><th>rejected_rank</th><th>top_level_parent</th></tr><tr><td>i64</td><td>i64</td><td>i64</td><td>i64</td><td>i64</td></tr></thead><tbody><tr><td>353</td><td>287</td><td>1</td><td>2</td><td>189</td></tr><tr><td>353</td><td>195</td><td>1</td><td>3</td><td>189</td></tr><tr><td>27552</td><td>27560</td><td>1</td><td>2</td><td>27550</td></tr><tr><td>77373</td><td>77337</td><td>1</td><td>2</td><td>77246</td></tr><tr><td>77666</td><td>77438</td><td>1</td><td>2</td><td>77246</td></tr><tr><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td></tr><tr><td>41812258</td><td>41812168</td><td>1</td><td>3</td><td>41811263</td></tr><tr><td>41812339</td><td>41812380</td><td>1</td><td>2</td><td>41811263</td></tr><tr><td>41812339</td><td>41812405</td><td>1</td><td>3</td><td>41811263</td></tr><tr><td>41812643</td><td>41812457</td><td>1</td><td>2</td><td>41811263</td></tr><tr><td>41812643</td><td>41812717</td><td>1</td><td>3</td><td>41811263</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (371_411, 5)\n",
       "┌──────────┬──────────┬─────────────┬───────────────┬──────────────────┐\n",
       "│ chosen   ┆ rejected ┆ chosen_rank ┆ rejected_rank ┆ top_level_parent │\n",
       "│ ---      ┆ ---      ┆ ---         ┆ ---           ┆ ---              │\n",
       "│ i64      ┆ i64      ┆ i64         ┆ i64           ┆ i64              │\n",
       "╞══════════╪══════════╪═════════════╪═══════════════╪══════════════════╡\n",
       "│ 353      ┆ 287      ┆ 1           ┆ 2             ┆ 189              │\n",
       "│ 353      ┆ 195      ┆ 1           ┆ 3             ┆ 189              │\n",
       "│ 27552    ┆ 27560    ┆ 1           ┆ 2             ┆ 27550            │\n",
       "│ 77373    ┆ 77337    ┆ 1           ┆ 2             ┆ 77246            │\n",
       "│ 77666    ┆ 77438    ┆ 1           ┆ 2             ┆ 77246            │\n",
       "│ …        ┆ …        ┆ …           ┆ …             ┆ …                │\n",
       "│ 41812258 ┆ 41812168 ┆ 1           ┆ 3             ┆ 41811263         │\n",
       "│ 41812339 ┆ 41812380 ┆ 1           ┆ 2             ┆ 41811263         │\n",
       "│ 41812339 ┆ 41812405 ┆ 1           ┆ 3             ┆ 41811263         │\n",
       "│ 41812643 ┆ 41812457 ┆ 1           ┆ 2             ┆ 41811263         │\n",
       "│ 41812643 ┆ 41812717 ┆ 1           ┆ 3             ┆ 41811263         │\n",
       "└──────────┴──────────┴─────────────┴───────────────┴──────────────────┘"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "from typing import List, TypedDict, Dict\n",
    "from collections import defaultdict\n",
    "import datetime\n",
    "\n",
    "\n",
    "progress_bar = tqdm(total=len(comments_df), desc=\"Processing groups\")\n",
    "\n",
    "\n",
    "class GroupableComment(TypedDict):\n",
    "    id: int\n",
    "    parent: int\n",
    "    time: datetime.datetime\n",
    "    sibling_rank: int\n",
    "    top_level_parent: int\n",
    "\n",
    "\n",
    "class RewardPair(TypedDict):\n",
    "    chosen: int  # the id of the chosen comment\n",
    "    rejected: int  # the id of the rejected comment\n",
    "    chosen_rank: int\n",
    "    rejected_rank: int\n",
    "    top_level_parent: int\n",
    "\n",
    "\n",
    "def process_group(group: List[GroupableComment]) -> list[RewardPair]:\n",
    "    \"\"\"\n",
    "    The goal of this function is to create a reward dataset for Hacker News\n",
    "    comments. It takes in a dataframe containing all the comments within a given\n",
    "    hour, and then breaks them down by the parent id.\n",
    "\n",
    "    Within each parent id group, it orders the comments by sibling rank (that\n",
    "    is, their relative position in the thread). Lower sibling rank means the\n",
    "    comment is higher in the thread, which we can assume means it had more\n",
    "    upvotes and is better.\n",
    "\n",
    "    We take each comment and its directly following sibling and mark those as\n",
    "    \"chosen\" and \"rejected\", creating a pairwise dataset we can use for reward\n",
    "    modeling.\n",
    "    \"\"\"\n",
    "    if len(group) == 0:\n",
    "        return []\n",
    "\n",
    "    sorted_group = sorted(group, key=lambda x: x[\"sibling_rank\"])\n",
    "\n",
    "    top_comment = sorted_group[0]\n",
    "    other_comments = sorted_group[1:]\n",
    "\n",
    "    group_reward_pairs = []\n",
    "\n",
    "    for other_comment in other_comments:\n",
    "        # Make sure the comments were posted within 30 minutes of each other. If\n",
    "        # they were too far apart, they likely aren't comparable since the older\n",
    "        # one is likely to win.\n",
    "        if (other_comment[\"time\"] - top_comment[\"time\"]).total_seconds() > 1800:\n",
    "            break\n",
    "\n",
    "        group_reward_pairs.append(\n",
    "            {\n",
    "                \"chosen\": top_comment[\"id\"],\n",
    "                \"rejected\": other_comment[\"id\"],\n",
    "                \"chosen_rank\": top_comment[\"sibling_rank\"],\n",
    "                \"rejected_rank\": other_comment[\"sibling_rank\"],\n",
    "                \"top_level_parent\": top_comment[\"top_level_parent\"],\n",
    "            }\n",
    "        )\n",
    "\n",
    "    return group_reward_pairs\n",
    "\n",
    "\n",
    "reward_pairs: list[RewardPair] = []\n",
    "parent_group: list[GroupableComment] = []\n",
    "current_parent = None\n",
    "\n",
    "sorted_by_parent = comments_df.select(\n",
    "    pl.col(\"id\"),\n",
    "    pl.col(\"sibling_rank\"),\n",
    "    pl.col(\"parent\"),\n",
    "    pl.col(\"time\"),\n",
    "    pl.col(\"top_level_parent\"),\n",
    ").sort(\"parent\")\n",
    "\n",
    "for comment in sorted_by_parent.iter_rows(named=True):\n",
    "    comment = GroupableComment(**comment)\n",
    "\n",
    "    # We already sorted by parent, so if the current comment's parent isn't the\n",
    "    # same as the previous one, we know we've moved on to a new parent and can\n",
    "    # process the previous set of comments.\n",
    "    if comment[\"parent\"] != current_parent:\n",
    "        reward_pairs.extend(process_group(parent_group))\n",
    "\n",
    "        parent_group = []\n",
    "        current_parent = comment[\"parent\"]\n",
    "\n",
    "    parent_group.append(comment)\n",
    "    progress_bar.update(1)\n",
    "\n",
    "# Close the progress bar\n",
    "progress_bar.close()\n",
    "\n",
    "pairs_df = pl.DataFrame(reward_pairs)\n",
    "\n",
    "pairs_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Split counts:\n",
      "shape: (3, 2)\n",
      "┌───────┬────────┐\n",
      "│ split ┆ count  │\n",
      "│ ---   ┆ ---    │\n",
      "│ str   ┆ u32    │\n",
      "╞═══════╪════════╡\n",
      "│ train ┆ 297776 │\n",
      "│ test  ┆ 37610  │\n",
      "│ val   ┆ 36025  │\n",
      "└───────┴────────┘\n",
      "\n",
      "Split percentages:\n",
      "shape: (3,)\n",
      "Series: 'count' [f64]\n",
      "[\n",
      "\t80.17\n",
      "\t10.13\n",
      "\t9.7\n",
      "]\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "\n",
    "# Get unique top_level_parent values\n",
    "unique_parents = pairs_df[\"top_level_parent\"].unique()\n",
    "\n",
    "# Randomly assign splits\n",
    "np.random.seed(42)  # for reproducibility\n",
    "split_assignments = np.random.choice(\n",
    "    [\"train\", \"test\", \"val\"], size=len(unique_parents), p=[0.8, 0.1, 0.1]\n",
    ")\n",
    "\n",
    "# Create a dictionary mapping top_level_parent to split\n",
    "parent_to_split = dict(zip(unique_parents, split_assignments))\n",
    "\n",
    "# Add the new 'split' column\n",
    "pairs_df = pairs_df.with_columns(\n",
    "    pl.col(\"top_level_parent\").replace_strict(parent_to_split).alias(\"split\")\n",
    ")\n",
    "\n",
    "# Verify the split proportions\n",
    "split_counts = pairs_df[\"split\"].value_counts().sort(\"count\", descending=True)\n",
    "print(\"Split counts:\")\n",
    "print(split_counts)\n",
    "\n",
    "# Calculate percentages\n",
    "total = split_counts[\"count\"].sum()\n",
    "percentages = (split_counts[\"count\"] / total * 100).round(2)\n",
    "print(\"\\nSplit percentages:\")\n",
    "print(percentages)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_pairs = pairs_df.sample(n=100000, seed=42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Building prompts:   0%|          | 0/100000 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Building prompts: 100%|██████████| 100000/100000 [00:54<00:00, 1840.04it/s]\n",
      "Building prompts: 100%|██████████| 100000/100000 [00:49<00:00, 2005.72it/s]\n",
      "Building prompts: 100%|██████████| 100000/100000 [00:53<00:00, 1868.94it/s]\n",
      "Building prompts: 100%|██████████| 100000/100000 [00:49<00:00, 2026.30it/s]\n"
     ]
    },
    {
     "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: (100_000, 6)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>chosen</th><th>rejected</th><th>chosen_rank</th><th>rejected_rank</th><th>top_level_parent</th><th>split</th></tr><tr><td>i64</td><td>i64</td><td>i64</td><td>i64</td><td>i64</td><td>str</td></tr></thead><tbody><tr><td>8883038</td><td>8881764</td><td>1</td><td>3</td><td>8878903</td><td>&quot;val&quot;</td></tr><tr><td>30012617</td><td>30012962</td><td>1</td><td>2</td><td>30011382</td><td>&quot;train&quot;</td></tr><tr><td>8525987</td><td>8524085</td><td>1</td><td>3</td><td>8523421</td><td>&quot;train&quot;</td></tr><tr><td>15055589</td><td>15055588</td><td>1</td><td>2</td><td>15054903</td><td>&quot;train&quot;</td></tr><tr><td>19970816</td><td>19969988</td><td>1</td><td>2</td><td>19968496</td><td>&quot;train&quot;</td></tr><tr><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td></tr><tr><td>1713593</td><td>1713509</td><td>1</td><td>2</td><td>1713276</td><td>&quot;train&quot;</td></tr><tr><td>7564905</td><td>7564825</td><td>1</td><td>2</td><td>7564680</td><td>&quot;test&quot;</td></tr><tr><td>27519642</td><td>27519251</td><td>1</td><td>3</td><td>27514437</td><td>&quot;train&quot;</td></tr><tr><td>4986067</td><td>4985725</td><td>1</td><td>3</td><td>4985517</td><td>&quot;train&quot;</td></tr><tr><td>28607000</td><td>28606640</td><td>1</td><td>3</td><td>28606075</td><td>&quot;test&quot;</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (100_000, 6)\n",
       "┌──────────┬──────────┬─────────────┬───────────────┬──────────────────┬───────┐\n",
       "│ chosen   ┆ rejected ┆ chosen_rank ┆ rejected_rank ┆ top_level_parent ┆ split │\n",
       "│ ---      ┆ ---      ┆ ---         ┆ ---           ┆ ---              ┆ ---   │\n",
       "│ i64      ┆ i64      ┆ i64         ┆ i64           ┆ i64              ┆ str   │\n",
       "╞══════════╪══════════╪═════════════╪═══════════════╪══════════════════╪═══════╡\n",
       "│ 8883038  ┆ 8881764  ┆ 1           ┆ 3             ┆ 8878903          ┆ val   │\n",
       "│ 30012617 ┆ 30012962 ┆ 1           ┆ 2             ┆ 30011382         ┆ train │\n",
       "│ 8525987  ┆ 8524085  ┆ 1           ┆ 3             ┆ 8523421          ┆ train │\n",
       "│ 15055589 ┆ 15055588 ┆ 1           ┆ 2             ┆ 15054903         ┆ train │\n",
       "│ 19970816 ┆ 19969988 ┆ 1           ┆ 2             ┆ 19968496         ┆ train │\n",
       "│ …        ┆ …        ┆ …           ┆ …             ┆ …                ┆ …     │\n",
       "│ 1713593  ┆ 1713509  ┆ 1           ┆ 2             ┆ 1713276          ┆ train │\n",
       "│ 7564905  ┆ 7564825  ┆ 1           ┆ 2             ┆ 7564680          ┆ test  │\n",
       "│ 27519642 ┆ 27519251 ┆ 1           ┆ 3             ┆ 27514437         ┆ train │\n",
       "│ 4986067  ┆ 4985725  ┆ 1           ┆ 3             ┆ 4985517          ┆ train │\n",
       "│ 28607000 ┆ 28606640 ┆ 1           ┆ 3             ┆ 28606075         ┆ test  │\n",
       "└──────────┴──────────┴─────────────┴───────────────┴──────────────────┴───────┘"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from utils import build_all_prompts\n",
    "\n",
    "chosen_prompts_v1 = build_all_prompts(sample_pairs[\"chosen\"], version=\"v1\")\n",
    "chosen_prompts_v2 = build_all_prompts(sample_pairs[\"chosen\"], version=\"v2\")\n",
    "rejected_prompts_v1 = build_all_prompts(sample_pairs[\"rejected\"], version=\"v1\")\n",
    "rejected_prompts_v2 = build_all_prompts(sample_pairs[\"rejected\"], version=\"v2\")\n",
    "\n",
    "sample_pairs_v1 = sample_pairs.with_columns(\n",
    "    pl.Series(\"chosen_prompt\", chosen_prompts_v1),\n",
    "    pl.Series(\"rejected_prompt\", rejected_prompts_v1),\n",
    ")\n",
    "\n",
    "sample_pairs_v2 = sample_pairs.with_columns(\n",
    "    pl.Series(\"chosen_prompt\", chosen_prompts_v2),\n",
    "    pl.Series(\"rejected_prompt\", rejected_prompts_v2),\n",
    ")\n",
    "\n",
    "sample_pairs\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "27adc5af54c54eefbcf65417d1b63a77",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Saving the dataset (0/1 shards):   0%|          | 0/30000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "62e554f2a48b4ea4aa76f79a59261f9f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Saving the dataset (0/1 shards):   0%|          | 0/500 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "404f171035824f0584193265654dabb2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Saving the dataset (0/1 shards):   0%|          | 0/1000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6d9f0f50cba549fd84c6fcbf155ad72a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Saving the dataset (0/1 shards):   0%|          | 0/30000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7fd3a86e2f104d9e90de8335989fbed8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Saving the dataset (0/1 shards):   0%|          | 0/500 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3aa42bc4c68041c992aa3979a2647a54",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Saving the dataset (0/1 shards):   0%|          | 0/1000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import numpy as np\n",
    "from datasets import Dataset, DatasetDict\n",
    "\n",
    "\n",
    "# Create a DatasetDict to hold all splits\n",
    "dataset_dict_v1 = DatasetDict(\n",
    "    {\n",
    "        \"train\": Dataset.from_polars(\n",
    "            sample_pairs_v1.filter(pl.col(\"split\") == \"train\").sample(n=30000, seed=42)\n",
    "        ),\n",
    "        \"validation\": Dataset.from_polars(\n",
    "            sample_pairs_v1.filter(pl.col(\"split\") == \"val\").sample(n=500, seed=42)\n",
    "        ),\n",
    "        \"test\": Dataset.from_polars(\n",
    "            sample_pairs_v1.filter(pl.col(\"split\") == \"test\").sample(n=1000, seed=42)\n",
    "        ),\n",
    "    }\n",
    ")\n",
    "\n",
    "dataset_dict_v1.save_to_disk(\"./data/sample_pairs_v1\")\n",
    "\n",
    "dataset_dict_v2 = DatasetDict(\n",
    "    {\n",
    "        \"train\": Dataset.from_polars(\n",
    "            sample_pairs_v2.filter(pl.col(\"split\") == \"train\").sample(n=30000, seed=42)\n",
    "        ),\n",
    "        \"validation\": Dataset.from_polars(\n",
    "            sample_pairs_v2.filter(pl.col(\"split\") == \"val\").sample(n=500, seed=42)\n",
    "        ),\n",
    "        \"test\": Dataset.from_polars(\n",
    "            sample_pairs_v2.filter(pl.col(\"split\") == \"test\").sample(n=1000, seed=42)\n",
    "        ),\n",
    "    }\n",
    ")\n",
    "\n",
    "dataset_dict_v2.save_to_disk(\"./data/sample_pairs_v2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "95dbfdbadbd94357afb19adbde712461",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Uploading the dataset shards:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "274ea27dd9924bf8a9582123eef1b27b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Creating parquet from Arrow format:   0%|          | 0/30 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9ba640c5d470430a8d74f6d11e705a48",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/workspace/best-hn/.venv/lib/python3.12/site-packages/huggingface_hub/lfs.py:336: UserWarning: hf_transfer is enabled but does not support uploading from bytes or BinaryIO, falling back to regular upload\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fb2913f1c58f4672af1c43bf24cb977a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Uploading the dataset shards:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "184d6934064f485ca15af8b9437fc068",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Creating parquet from Arrow format:   0%|          | 0/1 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "15769fd91ce74d5283a31d391c65540b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a795274e73484af3992075b88ab9a2e8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Uploading the dataset shards:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "02918ee592ee4dbd86c64a89395a565a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Creating parquet from Arrow format:   0%|          | 0/1 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "331a1a089fd04040be8c8351291ac971",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "62bf150af2bf40fe9c1c3d92a91b721e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Uploading the dataset shards:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d1ad47d93a7a45d9a0b2ffe9000836f0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Creating parquet from Arrow format:   0%|          | 0/30 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bd33415e5de24cd3b7a41ee64d032440",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d630b2d83577489e88a834a3051769fb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Uploading the dataset shards:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2cfac82b96d44482a2aa572b96ce62d5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Creating parquet from Arrow format:   0%|          | 0/1 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "285cf84079864bcf82b8564653a001ba",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a916daacb06744e6b4eb9cf5fd77d88c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Uploading the dataset shards:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a79ea473837b412588c2c992cc3e5f80",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Creating parquet from Arrow format:   0%|          | 0/1 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a0b01e77357d4ef6b4e512a0c1dd9dcf",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "CommitInfo(commit_url='https://huggingface.co/datasets/OpenPipe/best-hn-comment-pairs-v2/commit/0c19b0ee5a3bb1d1caf8071359897f22338f90b0', commit_message='Upload dataset', commit_description='', oid='0c19b0ee5a3bb1d1caf8071359897f22338f90b0', pr_url=None, repo_url=RepoUrl('https://huggingface.co/datasets/OpenPipe/best-hn-comment-pairs-v2', endpoint='https://huggingface.co', repo_type='dataset', repo_id='OpenPipe/best-hn-comment-pairs-v2'), pr_revision=None, pr_num=None)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv(\"/workspace/.env\")\n",
    "\n",
    "dataset_dict_v1.push_to_hub(\"OpenPipe/best-hn-comment-pairs-v1\")\n",
    "dataset_dict_v2.push_to_hub(\"OpenPipe/best-hn-comment-pairs-v2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Building prompts: 100%|██████████| 340702/340702 [02:44<00:00, 2065.44it/s]\n"
     ]
    }
   ],
   "source": [
    "top_comments = comments_df.filter(pl.col(\"sibling_rank\") == 1).filter(\n",
    "    (pl.col(\"id\").is_in(dataset_dict_v1[\"train\"][\"chosen\"]).not_())\n",
    "    & (pl.col(\"id\").is_in(dataset_dict_v1[\"train\"][\"rejected\"]).not_())\n",
    ")\n",
    "\n",
    "top_comments = top_comments.with_columns(\n",
    "    pl.Series(\"prompt\", build_all_prompts(top_comments[\"id\"], version=\"v1\"))\n",
    ")\n",
    "\n",
    "top_comments.write_parquet(\"./data/top_comments.parquet\")\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.12.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}