rate-random-comments.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "import polars as pl\n",
    "from utils import augmented_comments, build_all_prompts, run_inference_sglang\n",
    "\n",
    "df = augmented_comments()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Building prompts: 100%|██████████| 100000/100000 [00:58<00:00, 1698.95it/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, 15)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>id</th><th>type</th><th>by</th><th>time</th><th>title</th><th>text</th><th>url</th><th>score</th><th>parent</th><th>top_level_parent</th><th>descendants</th><th>kids</th><th>deleted</th><th>dead</th><th>prompt</th></tr><tr><td>i64</td><td>str</td><td>str</td><td>datetime[μs]</td><td>str</td><td>str</td><td>str</td><td>i64</td><td>i64</td><td>i64</td><td>i64</td><td>list[i64]</td><td>bool</td><td>bool</td><td>str</td></tr></thead><tbody><tr><td>7209407</td><td>&quot;comment&quot;</td><td>&quot;ricardobeat&quot;</td><td>2014-02-10 06:00:25</td><td>null</td><td>&quot;The article presents some (non…</td><td>null</td><td>null</td><td>7207821</td><td>7207506</td><td>null</td><td>[7210958, 7213827]</td><td>null</td><td>null</td><td>&quot;&lt;instructions&gt;Your goal is to …</td></tr><tr><td>13520503</td><td>&quot;comment&quot;</td><td>&quot;throwawayish&quot;</td><td>2017-01-30 12:11:14</td><td>null</td><td>&quot;Reaaally not that simple.&quot;</td><td>null</td><td>null</td><td>13519742</td><td>13517389</td><td>null</td><td>[13528758]</td><td>null</td><td>null</td><td>&quot;&lt;instructions&gt;Your goal is to …</td></tr><tr><td>13152625</td><td>&quot;comment&quot;</td><td>&quot;spikels&quot;</td><td>2016-12-11 19:19:47</td><td>null</td><td>&quot;That&amp;#x27;s an exaggeration. A…</td><td>null</td><td>null</td><td>13151772</td><td>13147495</td><td>null</td><td>[13165894]</td><td>null</td><td>null</td><td>&quot;&lt;instructions&gt;Your goal is to …</td></tr><tr><td>30957389</td><td>&quot;comment&quot;</td><td>&quot;ss108&quot;</td><td>2022-04-08 14:38:05</td><td>null</td><td>&quot;Ah, you&amp;#x27;re probably right…</td><td>null</td><td>null</td><td>30956591</td><td>30955290</td><td>null</td><td>null</td><td>null</td><td>null</td><td>&quot;&lt;instructions&gt;Your goal is to …</td></tr><tr><td>23335760</td><td>&quot;comment&quot;</td><td>&quot;en3r0&quot;</td><td>2020-05-28 10:41:15</td><td>null</td><td>&quot;I have been using Trilium late…</td><td>null</td><td>null</td><td>23335759</td><td>23335759</td><td>null</td><td>[23335860]</td><td>null</td><td>null</td><td>&quot;&lt;instructions&gt;Your goal is to …</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><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td></tr><tr><td>8110666</td><td>&quot;comment&quot;</td><td>&quot;erickookoo&quot;</td><td>2014-07-30 20:17:48</td><td>null</td><td>&quot;Part of the reason we built th…</td><td>null</td><td>null</td><td>8109374</td><td>8109114</td><td>null</td><td>null</td><td>null</td><td>null</td><td>&quot;&lt;instructions&gt;Your goal is to …</td></tr><tr><td>8039481</td><td>&quot;comment&quot;</td><td>&quot;josteink&quot;</td><td>2014-07-15 22:22:52</td><td>null</td><td>&quot;If you&amp;#x27;re worried Google …</td><td>null</td><td>null</td><td>8039322</td><td>8038990</td><td>null</td><td>null</td><td>null</td><td>null</td><td>&quot;&lt;instructions&gt;Your goal is to …</td></tr><tr><td>37965302</td><td>&quot;comment&quot;</td><td>&quot;williamdclt&quot;</td><td>2023-10-21 09:21:21</td><td>null</td><td>&quot;Do you have an example?&quot;</td><td>null</td><td>null</td><td>37965257</td><td>37962370</td><td>null</td><td>[37965954, 37965936]</td><td>null</td><td>null</td><td>&quot;&lt;instructions&gt;Your goal is to …</td></tr><tr><td>10112234</td><td>&quot;comment&quot;</td><td>&quot;afshin&quot;</td><td>2015-08-24 18:58:41</td><td>null</td><td>&quot;Just a guess, but it might be …</td><td>null</td><td>null</td><td>10111922</td><td>10108472</td><td>null</td><td>[10112273]</td><td>null</td><td>null</td><td>&quot;&lt;instructions&gt;Your goal is to …</td></tr><tr><td>3202168</td><td>&quot;comment&quot;</td><td>&quot;Tycho&quot;</td><td>2011-11-06 08:25:19</td><td>null</td><td>&quot;I hate how in HN discussions l…</td><td>null</td><td>null</td><td>3198171</td><td>3198171</td><td>null</td><td>null</td><td>null</td><td>null</td><td>&quot;&lt;instructions&gt;Your goal is to …</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (100_000, 15)\n",
       "┌──────────┬─────────┬──────────────┬─────────────┬───┬─────────────┬─────────┬──────┬─────────────┐\n",
       "│ id       ┆ type    ┆ by           ┆ time        ┆ … ┆ kids        ┆ deleted ┆ dead ┆ prompt      │\n",
       "│ ---      ┆ ---     ┆ ---          ┆ ---         ┆   ┆ ---         ┆ ---     ┆ ---  ┆ ---         │\n",
       "│ i64      ┆ str     ┆ str          ┆ datetime[μs ┆   ┆ list[i64]   ┆ bool    ┆ bool ┆ str         │\n",
       "│          ┆         ┆              ┆ ]           ┆   ┆             ┆         ┆      ┆             │\n",
       "╞══════════╪═════════╪══════════════╪═════════════╪═══╪═════════════╪═════════╪══════╪═════════════╡\n",
       "│ 7209407  ┆ comment ┆ ricardobeat  ┆ 2014-02-10  ┆ … ┆ [7210958,   ┆ null    ┆ null ┆ <instructio │\n",
       "│          ┆         ┆              ┆ 06:00:25    ┆   ┆ 7213827]    ┆         ┆      ┆ ns>Your     │\n",
       "│          ┆         ┆              ┆             ┆   ┆             ┆         ┆      ┆ goal is to  │\n",
       "│          ┆         ┆              ┆             ┆   ┆             ┆         ┆      ┆ …           │\n",
       "│ 13520503 ┆ comment ┆ throwawayish ┆ 2017-01-30  ┆ … ┆ [13528758]  ┆ null    ┆ null ┆ <instructio │\n",
       "│          ┆         ┆              ┆ 12:11:14    ┆   ┆             ┆         ┆      ┆ ns>Your     │\n",
       "│          ┆         ┆              ┆             ┆   ┆             ┆         ┆      ┆ goal is to  │\n",
       "│          ┆         ┆              ┆             ┆   ┆             ┆         ┆      ┆ …           │\n",
       "│ 13152625 ┆ comment ┆ spikels      ┆ 2016-12-11  ┆ … ┆ [13165894]  ┆ null    ┆ null ┆ <instructio │\n",
       "│          ┆         ┆              ┆ 19:19:47    ┆   ┆             ┆         ┆      ┆ ns>Your     │\n",
       "│          ┆         ┆              ┆             ┆   ┆             ┆         ┆      ┆ goal is to  │\n",
       "│          ┆         ┆              ┆             ┆   ┆             ┆         ┆      ┆ …           │\n",
       "│ 30957389 ┆ comment ┆ ss108        ┆ 2022-04-08  ┆ … ┆ null        ┆ null    ┆ null ┆ <instructio │\n",
       "│          ┆         ┆              ┆ 14:38:05    ┆   ┆             ┆         ┆      ┆ ns>Your     │\n",
       "│          ┆         ┆              ┆             ┆   ┆             ┆         ┆      ┆ goal is to  │\n",
       "│          ┆         ┆              ┆             ┆   ┆             ┆         ┆      ┆ …           │\n",
       "│ 23335760 ┆ comment ┆ en3r0        ┆ 2020-05-28  ┆ … ┆ [23335860]  ┆ null    ┆ null ┆ <instructio │\n",
       "│          ┆         ┆              ┆ 10:41:15    ┆   ┆             ┆         ┆      ┆ ns>Your     │\n",
       "│          ┆         ┆              ┆             ┆   ┆             ┆         ┆      ┆ goal is to  │\n",
       "│          ┆         ┆              ┆             ┆   ┆             ┆         ┆      ┆ …           │\n",
       "│ …        ┆ …       ┆ …            ┆ …           ┆ … ┆ …           ┆ …       ┆ …    ┆ …           │\n",
       "│ 8110666  ┆ comment ┆ erickookoo   ┆ 2014-07-30  ┆ … ┆ null        ┆ null    ┆ null ┆ <instructio │\n",
       "│          ┆         ┆              ┆ 20:17:48    ┆   ┆             ┆         ┆      ┆ ns>Your     │\n",
       "│          ┆         ┆              ┆             ┆   ┆             ┆         ┆      ┆ goal is to  │\n",
       "│          ┆         ┆              ┆             ┆   ┆             ┆         ┆      ┆ …           │\n",
       "│ 8039481  ┆ comment ┆ josteink     ┆ 2014-07-15  ┆ … ┆ null        ┆ null    ┆ null ┆ <instructio │\n",
       "│          ┆         ┆              ┆ 22:22:52    ┆   ┆             ┆         ┆      ┆ ns>Your     │\n",
       "│          ┆         ┆              ┆             ┆   ┆             ┆         ┆      ┆ goal is to  │\n",
       "│          ┆         ┆              ┆             ┆   ┆             ┆         ┆      ┆ …           │\n",
       "│ 37965302 ┆ comment ┆ williamdclt  ┆ 2023-10-21  ┆ … ┆ [37965954,  ┆ null    ┆ null ┆ <instructio │\n",
       "│          ┆         ┆              ┆ 09:21:21    ┆   ┆ 37965936]   ┆         ┆      ┆ ns>Your     │\n",
       "│          ┆         ┆              ┆             ┆   ┆             ┆         ┆      ┆ goal is to  │\n",
       "│          ┆         ┆              ┆             ┆   ┆             ┆         ┆      ┆ …           │\n",
       "│ 10112234 ┆ comment ┆ afshin       ┆ 2015-08-24  ┆ … ┆ [10112273]  ┆ null    ┆ null ┆ <instructio │\n",
       "│          ┆         ┆              ┆ 18:58:41    ┆   ┆             ┆         ┆      ┆ ns>Your     │\n",
       "│          ┆         ┆              ┆             ┆   ┆             ┆         ┆      ┆ goal is to  │\n",
       "│          ┆         ┆              ┆             ┆   ┆             ┆         ┆      ┆ …           │\n",
       "│ 3202168  ┆ comment ┆ Tycho        ┆ 2011-11-06  ┆ … ┆ null        ┆ null    ┆ null ┆ <instructio │\n",
       "│          ┆         ┆              ┆ 08:25:19    ┆   ┆             ┆         ┆      ┆ ns>Your     │\n",
       "│          ┆         ┆              ┆             ┆   ┆             ┆         ┆      ┆ goal is to  │\n",
       "│          ┆         ┆              ┆             ┆   ┆             ┆         ┆      ┆ …           │\n",
       "└──────────┴─────────┴──────────────┴─────────────┴───┴─────────────┴─────────┴──────┴─────────────┘"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "comments_df = df.filter(pl.col(\"type\") == \"comment\").sample(n=100000, seed=42)\n",
    "comments_df = comments_df.with_columns(\n",
    "    pl.Series(\"prompt\", build_all_prompts(comments_df[\"id\"])),\n",
    ")\n",
    "\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Running inference: 100%|██████████| 1000/1000 [23:29<00:00,  1.41s/it]\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, 16)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>id</th><th>type</th><th>by</th><th>time</th><th>title</th><th>text</th><th>url</th><th>score</th><th>parent</th><th>top_level_parent</th><th>descendants</th><th>kids</th><th>deleted</th><th>dead</th><th>prompt</th><th>reward</th></tr><tr><td>i64</td><td>str</td><td>str</td><td>datetime[μs]</td><td>str</td><td>str</td><td>str</td><td>i64</td><td>i64</td><td>i64</td><td>i64</td><td>list[i64]</td><td>bool</td><td>bool</td><td>str</td><td>f64</td></tr></thead><tbody><tr><td>7209407</td><td>&quot;comment&quot;</td><td>&quot;ricardobeat&quot;</td><td>2014-02-10 06:00:25</td><td>null</td><td>&quot;The article presents some (non…</td><td>null</td><td>null</td><td>7207821</td><td>7207506</td><td>null</td><td>[7210958, 7213827]</td><td>null</td><td>null</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>-4.71875</td></tr><tr><td>13520503</td><td>&quot;comment&quot;</td><td>&quot;throwawayish&quot;</td><td>2017-01-30 12:11:14</td><td>null</td><td>&quot;Reaaally not that simple.&quot;</td><td>null</td><td>null</td><td>13519742</td><td>13517389</td><td>null</td><td>[13528758]</td><td>null</td><td>null</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>-15.125</td></tr><tr><td>13152625</td><td>&quot;comment&quot;</td><td>&quot;spikels&quot;</td><td>2016-12-11 19:19:47</td><td>null</td><td>&quot;That&amp;#x27;s an exaggeration. A…</td><td>null</td><td>null</td><td>13151772</td><td>13147495</td><td>null</td><td>[13165894]</td><td>null</td><td>null</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>5.4375</td></tr><tr><td>30957389</td><td>&quot;comment&quot;</td><td>&quot;ss108&quot;</td><td>2022-04-08 14:38:05</td><td>null</td><td>&quot;Ah, you&amp;#x27;re probably right…</td><td>null</td><td>null</td><td>30956591</td><td>30955290</td><td>null</td><td>null</td><td>null</td><td>null</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>-21.0</td></tr><tr><td>23335760</td><td>&quot;comment&quot;</td><td>&quot;en3r0&quot;</td><td>2020-05-28 10:41:15</td><td>null</td><td>&quot;I have been using Trilium late…</td><td>null</td><td>null</td><td>23335759</td><td>23335759</td><td>null</td><td>[23335860]</td><td>null</td><td>null</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>-20.125</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><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td></tr><tr><td>8110666</td><td>&quot;comment&quot;</td><td>&quot;erickookoo&quot;</td><td>2014-07-30 20:17:48</td><td>null</td><td>&quot;Part of the reason we built th…</td><td>null</td><td>null</td><td>8109374</td><td>8109114</td><td>null</td><td>null</td><td>null</td><td>null</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>0.484375</td></tr><tr><td>8039481</td><td>&quot;comment&quot;</td><td>&quot;josteink&quot;</td><td>2014-07-15 22:22:52</td><td>null</td><td>&quot;If you&amp;#x27;re worried Google …</td><td>null</td><td>null</td><td>8039322</td><td>8038990</td><td>null</td><td>null</td><td>null</td><td>null</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>-12.125</td></tr><tr><td>37965302</td><td>&quot;comment&quot;</td><td>&quot;williamdclt&quot;</td><td>2023-10-21 09:21:21</td><td>null</td><td>&quot;Do you have an example?&quot;</td><td>null</td><td>null</td><td>37965257</td><td>37962370</td><td>null</td><td>[37965954, 37965936]</td><td>null</td><td>null</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>-22.25</td></tr><tr><td>10112234</td><td>&quot;comment&quot;</td><td>&quot;afshin&quot;</td><td>2015-08-24 18:58:41</td><td>null</td><td>&quot;Just a guess, but it might be …</td><td>null</td><td>null</td><td>10111922</td><td>10108472</td><td>null</td><td>[10112273]</td><td>null</td><td>null</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>-7.6875</td></tr><tr><td>3202168</td><td>&quot;comment&quot;</td><td>&quot;Tycho&quot;</td><td>2011-11-06 08:25:19</td><td>null</td><td>&quot;I hate how in HN discussions l…</td><td>null</td><td>null</td><td>3198171</td><td>3198171</td><td>null</td><td>null</td><td>null</td><td>null</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>-4.0</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (100_000, 16)\n",
       "┌──────────┬─────────┬──────────────┬───────────────┬───┬─────────┬──────┬──────────────┬──────────┐\n",
       "│ id       ┆ type    ┆ by           ┆ time          ┆ … ┆ deleted ┆ dead ┆ prompt       ┆ reward   │\n",
       "│ ---      ┆ ---     ┆ ---          ┆ ---           ┆   ┆ ---     ┆ ---  ┆ ---          ┆ ---      │\n",
       "│ i64      ┆ str     ┆ str          ┆ datetime[μs]  ┆   ┆ bool    ┆ bool ┆ str          ┆ f64      │\n",
       "╞══════════╪═════════╪══════════════╪═══════════════╪═══╪═════════╪══════╪══════════════╪══════════╡\n",
       "│ 7209407  ┆ comment ┆ ricardobeat  ┆ 2014-02-10    ┆ … ┆ null    ┆ null ┆ <instruction ┆ -4.71875 │\n",
       "│          ┆         ┆              ┆ 06:00:25      ┆   ┆         ┆      ┆ s>Your goal  ┆          │\n",
       "│          ┆         ┆              ┆               ┆   ┆         ┆      ┆ is to …      ┆          │\n",
       "│ 13520503 ┆ comment ┆ throwawayish ┆ 2017-01-30    ┆ … ┆ null    ┆ null ┆ <instruction ┆ -15.125  │\n",
       "│          ┆         ┆              ┆ 12:11:14      ┆   ┆         ┆      ┆ s>Your goal  ┆          │\n",
       "│          ┆         ┆              ┆               ┆   ┆         ┆      ┆ is to …      ┆          │\n",
       "│ 13152625 ┆ comment ┆ spikels      ┆ 2016-12-11    ┆ … ┆ null    ┆ null ┆ <instruction ┆ 5.4375   │\n",
       "│          ┆         ┆              ┆ 19:19:47      ┆   ┆         ┆      ┆ s>Your goal  ┆          │\n",
       "│          ┆         ┆              ┆               ┆   ┆         ┆      ┆ is to …      ┆          │\n",
       "│ 30957389 ┆ comment ┆ ss108        ┆ 2022-04-08    ┆ … ┆ null    ┆ null ┆ <instruction ┆ -21.0    │\n",
       "│          ┆         ┆              ┆ 14:38:05      ┆   ┆         ┆      ┆ s>Your goal  ┆          │\n",
       "│          ┆         ┆              ┆               ┆   ┆         ┆      ┆ is to …      ┆          │\n",
       "│ 23335760 ┆ comment ┆ en3r0        ┆ 2020-05-28    ┆ … ┆ null    ┆ null ┆ <instruction ┆ -20.125  │\n",
       "│          ┆         ┆              ┆ 10:41:15      ┆   ┆         ┆      ┆ s>Your goal  ┆          │\n",
       "│          ┆         ┆              ┆               ┆   ┆         ┆      ┆ is to …      ┆          │\n",
       "│ …        ┆ …       ┆ …            ┆ …             ┆ … ┆ …       ┆ …    ┆ …            ┆ …        │\n",
       "│ 8110666  ┆ comment ┆ erickookoo   ┆ 2014-07-30    ┆ … ┆ null    ┆ null ┆ <instruction ┆ 0.484375 │\n",
       "│          ┆         ┆              ┆ 20:17:48      ┆   ┆         ┆      ┆ s>Your goal  ┆          │\n",
       "│          ┆         ┆              ┆               ┆   ┆         ┆      ┆ is to …      ┆          │\n",
       "│ 8039481  ┆ comment ┆ josteink     ┆ 2014-07-15    ┆ … ┆ null    ┆ null ┆ <instruction ┆ -12.125  │\n",
       "│          ┆         ┆              ┆ 22:22:52      ┆   ┆         ┆      ┆ s>Your goal  ┆          │\n",
       "│          ┆         ┆              ┆               ┆   ┆         ┆      ┆ is to …      ┆          │\n",
       "│ 37965302 ┆ comment ┆ williamdclt  ┆ 2023-10-21    ┆ … ┆ null    ┆ null ┆ <instruction ┆ -22.25   │\n",
       "│          ┆         ┆              ┆ 09:21:21      ┆   ┆         ┆      ┆ s>Your goal  ┆          │\n",
       "│          ┆         ┆              ┆               ┆   ┆         ┆      ┆ is to …      ┆          │\n",
       "│ 10112234 ┆ comment ┆ afshin       ┆ 2015-08-24    ┆ … ┆ null    ┆ null ┆ <instruction ┆ -7.6875  │\n",
       "│          ┆         ┆              ┆ 18:58:41      ┆   ┆         ┆      ┆ s>Your goal  ┆          │\n",
       "│          ┆         ┆              ┆               ┆   ┆         ┆      ┆ is to …      ┆          │\n",
       "│ 3202168  ┆ comment ┆ Tycho        ┆ 2011-11-06    ┆ … ┆ null    ┆ null ┆ <instruction ┆ -4.0     │\n",
       "│          ┆         ┆              ┆ 08:25:19      ┆   ┆         ┆      ┆ s>Your goal  ┆          │\n",
       "│          ┆         ┆              ┆               ┆   ┆         ┆      ┆ is to …      ┆          │\n",
       "└──────────┴─────────┴──────────────┴───────────────┴───┴─────────┴──────┴──────────────┴──────────┘"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rewards = run_inference_sglang(comments_df[\"prompt\"])\n",
    "\n",
    "comments_df = comments_df.with_columns(pl.Series(\"reward\", rewards))\n",
    "\n",
    "comments_df.write_parquet(\"./data/random_comments_with_reward.parquet\")\n",
    "\n",
    "comments_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import polars as pl\n",
    "\n",
    "comments_df = pl.read_parquet(\"./data/random_comments_with_reward.parquet\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import html\n",
    "import re\n",
    "from utils import with_story_info\n",
    "\n",
    "\n",
    "def unescape_html(text):\n",
    "    unescaped = html.unescape(text).replace(\"<p>\", \"\\n\\n\")\n",
    "    return re.sub(r'<a href=\"([^\"]+)\"[^>]*>[^<]+</a>', r\"\\1\", unescaped)\n",
    "\n",
    "\n",
    "comments_df = with_story_info(comments_df)\n",
    "\n",
    "comments_df = comments_df.with_columns(\n",
    "    pl.concat_str(\n",
    "        pl.lit(\"https://news.ycombinator.com/item?id=\"),\n",
    "        pl.col(\"id\"),\n",
    "    ).alias(\"link\"),\n",
    "    pl.col(\"time\").dt.strftime(\"%B %d, %Y\").alias(\"date\"),\n",
    "    pl.col(\"text\").map_elements(unescape_html, return_dtype=pl.String),\n",
    ")\n",
    "\n",
    "comments_df.select(pl.col(\"date\", \"by\", \"link\", \"text\", \"reward\")).sort(\n",
    "    \"reward\", descending=True\n",
    ").head(100).write_csv(\"./data/top_random_comments_with_links.csv\")\n",
    "\n",
    "comments_df.select(pl.col(\"date\", \"by\", \"link\", \"text\", \"reward\")).sort(\n",
    "    \"reward\", descending=False\n",
    ").head(100).write_csv(\"./data/bottom_random_comments_with_links.csv\")\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
}