validate-model-perf.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['chosen', 'rejected', 'chosen_rank', 'rejected_rank', 'top_level_parent', 'split', 'chosen_prompt', 'rejected_prompt'],\n",
       "    num_rows: 1000\n",
       "})"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from datasets import load_from_disk\n",
    "\n",
    "ds = load_from_disk(\"./data/sample_pairs\")[\"test\"]\n",
    "ds"
   ]
  },
  {
   "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_000, 8)</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><th>chosen_prompt</th><th>rejected_prompt</th></tr><tr><td>i64</td><td>i64</td><td>i64</td><td>i64</td><td>i64</td><td>str</td><td>str</td><td>str</td></tr></thead><tbody><tr><td>31042685</td><td>31041065</td><td>1</td><td>3</td><td>31012025</td><td>&quot;test&quot;</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>&quot;&lt;instructions&gt;Your goal is to …</td></tr><tr><td>39808386</td><td>39806045</td><td>1</td><td>2</td><td>39758712</td><td>&quot;test&quot;</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>&quot;&lt;instructions&gt;Your goal is to …</td></tr><tr><td>34895313</td><td>34894926</td><td>1</td><td>2</td><td>34894667</td><td>&quot;test&quot;</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>&quot;&lt;instructions&gt;Your goal is to …</td></tr><tr><td>27699072</td><td>27697619</td><td>1</td><td>2</td><td>27696055</td><td>&quot;test&quot;</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>&quot;&lt;instructions&gt;Your goal is to …</td></tr><tr><td>5348842</td><td>5347984</td><td>1</td><td>3</td><td>5347543</td><td>&quot;test&quot;</td><td>&quot;&lt;instructions&gt;Your goal is to …</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></tr><tr><td>15301562</td><td>15301600</td><td>1</td><td>2</td><td>15301151</td><td>&quot;test&quot;</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>&quot;&lt;instructions&gt;Your goal is to …</td></tr><tr><td>36831606</td><td>36830894</td><td>1</td><td>3</td><td>36829854</td><td>&quot;test&quot;</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>&quot;&lt;instructions&gt;Your goal is to …</td></tr><tr><td>29624080</td><td>29623860</td><td>1</td><td>2</td><td>29622770</td><td>&quot;test&quot;</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>&quot;&lt;instructions&gt;Your goal is to …</td></tr><tr><td>31660200</td><td>31659387</td><td>1</td><td>3</td><td>31657006</td><td>&quot;test&quot;</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>&quot;&lt;instructions&gt;Your goal is to …</td></tr><tr><td>25421107</td><td>25420617</td><td>1</td><td>3</td><td>25418676</td><td>&quot;test&quot;</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>&quot;&lt;instructions&gt;Your goal is to …</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (1_000, 8)\n",
       "┌──────────┬──────────┬─────────────┬─────────────┬─────────────┬───────┬─────────────┬────────────┐\n",
       "│ chosen   ┆ rejected ┆ chosen_rank ┆ rejected_ra ┆ top_level_p ┆ split ┆ chosen_prom ┆ rejected_p │\n",
       "│ ---      ┆ ---      ┆ ---         ┆ nk          ┆ arent       ┆ ---   ┆ pt          ┆ rompt      │\n",
       "│ i64      ┆ i64      ┆ i64         ┆ ---         ┆ ---         ┆ str   ┆ ---         ┆ ---        │\n",
       "│          ┆          ┆             ┆ i64         ┆ i64         ┆       ┆ str         ┆ str        │\n",
       "╞══════════╪══════════╪═════════════╪═════════════╪═════════════╪═══════╪═════════════╪════════════╡\n",
       "│ 31042685 ┆ 31041065 ┆ 1           ┆ 3           ┆ 31012025    ┆ test  ┆ <instructio ┆ <instructi │\n",
       "│          ┆          ┆             ┆             ┆             ┆       ┆ ns>Your     ┆ ons>Your   │\n",
       "│          ┆          ┆             ┆             ┆             ┆       ┆ goal is to  ┆ goal is to │\n",
       "│          ┆          ┆             ┆             ┆             ┆       ┆ …           ┆ …          │\n",
       "│ 39808386 ┆ 39806045 ┆ 1           ┆ 2           ┆ 39758712    ┆ test  ┆ <instructio ┆ <instructi │\n",
       "│          ┆          ┆             ┆             ┆             ┆       ┆ ns>Your     ┆ ons>Your   │\n",
       "│          ┆          ┆             ┆             ┆             ┆       ┆ goal is to  ┆ goal is to │\n",
       "│          ┆          ┆             ┆             ┆             ┆       ┆ …           ┆ …          │\n",
       "│ 34895313 ┆ 34894926 ┆ 1           ┆ 2           ┆ 34894667    ┆ test  ┆ <instructio ┆ <instructi │\n",
       "│          ┆          ┆             ┆             ┆             ┆       ┆ ns>Your     ┆ ons>Your   │\n",
       "│          ┆          ┆             ┆             ┆             ┆       ┆ goal is to  ┆ goal is to │\n",
       "│          ┆          ┆             ┆             ┆             ┆       ┆ …           ┆ …          │\n",
       "│ 27699072 ┆ 27697619 ┆ 1           ┆ 2           ┆ 27696055    ┆ test  ┆ <instructio ┆ <instructi │\n",
       "│          ┆          ┆             ┆             ┆             ┆       ┆ ns>Your     ┆ ons>Your   │\n",
       "│          ┆          ┆             ┆             ┆             ┆       ┆ goal is to  ┆ goal is to │\n",
       "│          ┆          ┆             ┆             ┆             ┆       ┆ …           ┆ …          │\n",
       "│ 5348842  ┆ 5347984  ┆ 1           ┆ 3           ┆ 5347543     ┆ test  ┆ <instructio ┆ <instructi │\n",
       "│          ┆          ┆             ┆             ┆             ┆       ┆ ns>Your     ┆ ons>Your   │\n",
       "│          ┆          ┆             ┆             ┆             ┆       ┆ goal is to  ┆ goal is to │\n",
       "│          ┆          ┆             ┆             ┆             ┆       ┆ …           ┆ …          │\n",
       "│ …        ┆ …        ┆ …           ┆ …           ┆ …           ┆ …     ┆ …           ┆ …          │\n",
       "│ 15301562 ┆ 15301600 ┆ 1           ┆ 2           ┆ 15301151    ┆ test  ┆ <instructio ┆ <instructi │\n",
       "│          ┆          ┆             ┆             ┆             ┆       ┆ ns>Your     ┆ ons>Your   │\n",
       "│          ┆          ┆             ┆             ┆             ┆       ┆ goal is to  ┆ goal is to │\n",
       "│          ┆          ┆             ┆             ┆             ┆       ┆ …           ┆ …          │\n",
       "│ 36831606 ┆ 36830894 ┆ 1           ┆ 3           ┆ 36829854    ┆ test  ┆ <instructio ┆ <instructi │\n",
       "│          ┆          ┆             ┆             ┆             ┆       ┆ ns>Your     ┆ ons>Your   │\n",
       "│          ┆          ┆             ┆             ┆             ┆       ┆ goal is to  ┆ goal is to │\n",
       "│          ┆          ┆             ┆             ┆             ┆       ┆ …           ┆ …          │\n",
       "│ 29624080 ┆ 29623860 ┆ 1           ┆ 2           ┆ 29622770    ┆ test  ┆ <instructio ┆ <instructi │\n",
       "│          ┆          ┆             ┆             ┆             ┆       ┆ ns>Your     ┆ ons>Your   │\n",
       "│          ┆          ┆             ┆             ┆             ┆       ┆ goal is to  ┆ goal is to │\n",
       "│          ┆          ┆             ┆             ┆             ┆       ┆ …           ┆ …          │\n",
       "│ 31660200 ┆ 31659387 ┆ 1           ┆ 3           ┆ 31657006    ┆ test  ┆ <instructio ┆ <instructi │\n",
       "│          ┆          ┆             ┆             ┆             ┆       ┆ ns>Your     ┆ ons>Your   │\n",
       "│          ┆          ┆             ┆             ┆             ┆       ┆ goal is to  ┆ goal is to │\n",
       "│          ┆          ┆             ┆             ┆             ┆       ┆ …           ┆ …          │\n",
       "│ 25421107 ┆ 25420617 ┆ 1           ┆ 3           ┆ 25418676    ┆ test  ┆ <instructio ┆ <instructi │\n",
       "│          ┆          ┆             ┆             ┆             ┆       ┆ ns>Your     ┆ ons>Your   │\n",
       "│          ┆          ┆             ┆             ┆             ┆       ┆ goal is to  ┆ goal is to │\n",
       "│          ┆          ┆             ┆             ┆             ┆       ┆ …           ┆ …          │\n",
       "└──────────┴──────────┴─────────────┴─────────────┴─────────────┴───────┴─────────────┴────────────┘"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import polars as pl\n",
    "\n",
    "ds = ds.to_polars()\n",
    "ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run the following in the terminal:\n",
    "# uv run python -m sglang.launch_server --model-path ./models/llama_32_8b_merged"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[-7.40625, -6.90625]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# launch server\n",
    "# python -m sglang.launch_server --model LxzGordon/URM-LLaMa-3.1-8B --is-embedding\n",
    "\n",
    "import requests\n",
    "\n",
    "url = \"http://127.0.0.1:30000\"\n",
    "\n",
    "\n",
    "def run_inference_sglang(prompts: list[str]) -> list[float]:\n",
    "    json_data = {\n",
    "        \"conv\": prompts,\n",
    "    }\n",
    "    response = requests.post(url + \"/judge\", json=json_data).json()\n",
    "\n",
    "    return [x[\"embedding\"][0] for x in response]\n",
    "\n",
    "\n",
    "sample_prompts = [\n",
    "    \"random prompt 1\",\n",
    "    \"random prompt 2\",\n",
    "]\n",
    "\n",
    "run_inference_sglang(sample_prompts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "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_000, 10)</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><th>chosen_prompt</th><th>rejected_prompt</th><th>chosen_reward</th><th>rejected_reward</th></tr><tr><td>i64</td><td>i64</td><td>i64</td><td>i64</td><td>i64</td><td>str</td><td>str</td><td>str</td><td>f64</td><td>f64</td></tr></thead><tbody><tr><td>31042685</td><td>31041065</td><td>1</td><td>3</td><td>31012025</td><td>&quot;test&quot;</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>-5.53125</td><td>-9.25</td></tr><tr><td>39808386</td><td>39806045</td><td>1</td><td>2</td><td>39758712</td><td>&quot;test&quot;</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>-13.5</td><td>-17.25</td></tr><tr><td>34895313</td><td>34894926</td><td>1</td><td>2</td><td>34894667</td><td>&quot;test&quot;</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>-4.5</td><td>-10.4375</td></tr><tr><td>27699072</td><td>27697619</td><td>1</td><td>2</td><td>27696055</td><td>&quot;test&quot;</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>-2.828125</td><td>-3.0</td></tr><tr><td>5348842</td><td>5347984</td><td>1</td><td>3</td><td>5347543</td><td>&quot;test&quot;</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>-9.0625</td><td>-9.25</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>15301562</td><td>15301600</td><td>1</td><td>2</td><td>15301151</td><td>&quot;test&quot;</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>-13.3125</td><td>-14.5625</td></tr><tr><td>36831606</td><td>36830894</td><td>1</td><td>3</td><td>36829854</td><td>&quot;test&quot;</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>-4.09375</td><td>-6.71875</td></tr><tr><td>29624080</td><td>29623860</td><td>1</td><td>2</td><td>29622770</td><td>&quot;test&quot;</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>-5.40625</td><td>-3.609375</td></tr><tr><td>31660200</td><td>31659387</td><td>1</td><td>3</td><td>31657006</td><td>&quot;test&quot;</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>-5.625</td><td>-7.5</td></tr><tr><td>25421107</td><td>25420617</td><td>1</td><td>3</td><td>25418676</td><td>&quot;test&quot;</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>-4.40625</td><td>-5.125</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (1_000, 10)\n",
       "┌──────────┬──────────┬────────────┬───────────┬───┬───────────┬───────────┬───────────┬───────────┐\n",
       "│ chosen   ┆ rejected ┆ chosen_ran ┆ rejected_ ┆ … ┆ chosen_pr ┆ rejected_ ┆ chosen_re ┆ rejected_ │\n",
       "│ ---      ┆ ---      ┆ k          ┆ rank      ┆   ┆ ompt      ┆ prompt    ┆ ward      ┆ reward    │\n",
       "│ i64      ┆ i64      ┆ ---        ┆ ---       ┆   ┆ ---       ┆ ---       ┆ ---       ┆ ---       │\n",
       "│          ┆          ┆ i64        ┆ i64       ┆   ┆ str       ┆ str       ┆ f64       ┆ f64       │\n",
       "╞══════════╪══════════╪════════════╪═══════════╪═══╪═══════════╪═══════════╪═══════════╪═══════════╡\n",
       "│ 31042685 ┆ 31041065 ┆ 1          ┆ 3         ┆ … ┆ <instruct ┆ <instruct ┆ -5.53125  ┆ -9.25     │\n",
       "│          ┆          ┆            ┆           ┆   ┆ ions>Your ┆ ions>Your ┆           ┆           │\n",
       "│          ┆          ┆            ┆           ┆   ┆ goal is   ┆ goal is   ┆           ┆           │\n",
       "│          ┆          ┆            ┆           ┆   ┆ to …      ┆ to …      ┆           ┆           │\n",
       "│ 39808386 ┆ 39806045 ┆ 1          ┆ 2         ┆ … ┆ <instruct ┆ <instruct ┆ -13.5     ┆ -17.25    │\n",
       "│          ┆          ┆            ┆           ┆   ┆ ions>Your ┆ ions>Your ┆           ┆           │\n",
       "│          ┆          ┆            ┆           ┆   ┆ goal is   ┆ goal is   ┆           ┆           │\n",
       "│          ┆          ┆            ┆           ┆   ┆ to …      ┆ to …      ┆           ┆           │\n",
       "│ 34895313 ┆ 34894926 ┆ 1          ┆ 2         ┆ … ┆ <instruct ┆ <instruct ┆ -4.5      ┆ -10.4375  │\n",
       "│          ┆          ┆            ┆           ┆   ┆ ions>Your ┆ ions>Your ┆           ┆           │\n",
       "│          ┆          ┆            ┆           ┆   ┆ goal is   ┆ goal is   ┆           ┆           │\n",
       "│          ┆          ┆            ┆           ┆   ┆ to …      ┆ to …      ┆           ┆           │\n",
       "│ 27699072 ┆ 27697619 ┆ 1          ┆ 2         ┆ … ┆ <instruct ┆ <instruct ┆ -2.828125 ┆ -3.0      │\n",
       "│          ┆          ┆            ┆           ┆   ┆ ions>Your ┆ ions>Your ┆           ┆           │\n",
       "│          ┆          ┆            ┆           ┆   ┆ goal is   ┆ goal is   ┆           ┆           │\n",
       "│          ┆          ┆            ┆           ┆   ┆ to …      ┆ to …      ┆           ┆           │\n",
       "│ 5348842  ┆ 5347984  ┆ 1          ┆ 3         ┆ … ┆ <instruct ┆ <instruct ┆ -9.0625   ┆ -9.25     │\n",
       "│          ┆          ┆            ┆           ┆   ┆ ions>Your ┆ ions>Your ┆           ┆           │\n",
       "│          ┆          ┆            ┆           ┆   ┆ goal is   ┆ goal is   ┆           ┆           │\n",
       "│          ┆          ┆            ┆           ┆   ┆ to …      ┆ to …      ┆           ┆           │\n",
       "│ …        ┆ …        ┆ …          ┆ …         ┆ … ┆ …         ┆ …         ┆ …         ┆ …         │\n",
       "│ 15301562 ┆ 15301600 ┆ 1          ┆ 2         ┆ … ┆ <instruct ┆ <instruct ┆ -13.3125  ┆ -14.5625  │\n",
       "│          ┆          ┆            ┆           ┆   ┆ ions>Your ┆ ions>Your ┆           ┆           │\n",
       "│          ┆          ┆            ┆           ┆   ┆ goal is   ┆ goal is   ┆           ┆           │\n",
       "│          ┆          ┆            ┆           ┆   ┆ to …      ┆ to …      ┆           ┆           │\n",
       "│ 36831606 ┆ 36830894 ┆ 1          ┆ 3         ┆ … ┆ <instruct ┆ <instruct ┆ -4.09375  ┆ -6.71875  │\n",
       "│          ┆          ┆            ┆           ┆   ┆ ions>Your ┆ ions>Your ┆           ┆           │\n",
       "│          ┆          ┆            ┆           ┆   ┆ goal is   ┆ goal is   ┆           ┆           │\n",
       "│          ┆          ┆            ┆           ┆   ┆ to …      ┆ to …      ┆           ┆           │\n",
       "│ 29624080 ┆ 29623860 ┆ 1          ┆ 2         ┆ … ┆ <instruct ┆ <instruct ┆ -5.40625  ┆ -3.609375 │\n",
       "│          ┆          ┆            ┆           ┆   ┆ ions>Your ┆ ions>Your ┆           ┆           │\n",
       "│          ┆          ┆            ┆           ┆   ┆ goal is   ┆ goal is   ┆           ┆           │\n",
       "│          ┆          ┆            ┆           ┆   ┆ to …      ┆ to …      ┆           ┆           │\n",
       "│ 31660200 ┆ 31659387 ┆ 1          ┆ 3         ┆ … ┆ <instruct ┆ <instruct ┆ -5.625    ┆ -7.5      │\n",
       "│          ┆          ┆            ┆           ┆   ┆ ions>Your ┆ ions>Your ┆           ┆           │\n",
       "│          ┆          ┆            ┆           ┆   ┆ goal is   ┆ goal is   ┆           ┆           │\n",
       "│          ┆          ┆            ┆           ┆   ┆ to …      ┆ to …      ┆           ┆           │\n",
       "│ 25421107 ┆ 25420617 ┆ 1          ┆ 3         ┆ … ┆ <instruct ┆ <instruct ┆ -4.40625  ┆ -5.125    │\n",
       "│          ┆          ┆            ┆           ┆   ┆ ions>Your ┆ ions>Your ┆           ┆           │\n",
       "│          ┆          ┆            ┆           ┆   ┆ goal is   ┆ goal is   ┆           ┆           │\n",
       "│          ┆          ┆            ┆           ┆   ┆ to …      ┆ to …      ┆           ┆           │\n",
       "└──────────┴──────────┴────────────┴───────────┴───┴───────────┴───────────┴───────────┴───────────┘"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chosen_rewards = run_inference_sglang(ds[\"chosen_prompt\"].to_list())\n",
    "rejected_rewards = run_inference_sglang(ds[\"rejected_prompt\"].to_list())\n",
    "\n",
    "ds = ds.with_columns(\n",
    "    pl.Series(name=\"chosen_reward\", values=chosen_rewards),\n",
    "    pl.Series(name=\"rejected_reward\", values=rejected_rewards),\n",
    ")\n"
   ]
  },
  {
   "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: (1_000, 11)</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><th>chosen_prompt</th><th>rejected_prompt</th><th>chosen_reward</th><th>rejected_reward</th><th>correct</th></tr><tr><td>i64</td><td>i64</td><td>i64</td><td>i64</td><td>i64</td><td>str</td><td>str</td><td>str</td><td>f64</td><td>f64</td><td>bool</td></tr></thead><tbody><tr><td>31042685</td><td>31041065</td><td>1</td><td>3</td><td>31012025</td><td>&quot;test&quot;</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>-5.53125</td><td>-9.25</td><td>true</td></tr><tr><td>39808386</td><td>39806045</td><td>1</td><td>2</td><td>39758712</td><td>&quot;test&quot;</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>-13.5</td><td>-17.25</td><td>true</td></tr><tr><td>34895313</td><td>34894926</td><td>1</td><td>2</td><td>34894667</td><td>&quot;test&quot;</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>-4.5</td><td>-10.4375</td><td>true</td></tr><tr><td>27699072</td><td>27697619</td><td>1</td><td>2</td><td>27696055</td><td>&quot;test&quot;</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>-2.828125</td><td>-3.0</td><td>true</td></tr><tr><td>5348842</td><td>5347984</td><td>1</td><td>3</td><td>5347543</td><td>&quot;test&quot;</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>-9.0625</td><td>-9.25</td><td>true</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></tr><tr><td>15301562</td><td>15301600</td><td>1</td><td>2</td><td>15301151</td><td>&quot;test&quot;</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>-13.3125</td><td>-14.5625</td><td>true</td></tr><tr><td>36831606</td><td>36830894</td><td>1</td><td>3</td><td>36829854</td><td>&quot;test&quot;</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>-4.09375</td><td>-6.71875</td><td>true</td></tr><tr><td>29624080</td><td>29623860</td><td>1</td><td>2</td><td>29622770</td><td>&quot;test&quot;</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>-5.40625</td><td>-3.609375</td><td>false</td></tr><tr><td>31660200</td><td>31659387</td><td>1</td><td>3</td><td>31657006</td><td>&quot;test&quot;</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>-5.625</td><td>-7.5</td><td>true</td></tr><tr><td>25421107</td><td>25420617</td><td>1</td><td>3</td><td>25418676</td><td>&quot;test&quot;</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>&quot;&lt;instructions&gt;Your goal is to …</td><td>-4.40625</td><td>-5.125</td><td>true</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (1_000, 11)\n",
       "┌──────────┬──────────┬────────────┬────────────┬───┬────────────┬───────────┬───────────┬─────────┐\n",
       "│ chosen   ┆ rejected ┆ chosen_ran ┆ rejected_r ┆ … ┆ rejected_p ┆ chosen_re ┆ rejected_ ┆ correct │\n",
       "│ ---      ┆ ---      ┆ k          ┆ ank        ┆   ┆ rompt      ┆ ward      ┆ reward    ┆ ---     │\n",
       "│ i64      ┆ i64      ┆ ---        ┆ ---        ┆   ┆ ---        ┆ ---       ┆ ---       ┆ bool    │\n",
       "│          ┆          ┆ i64        ┆ i64        ┆   ┆ str        ┆ f64       ┆ f64       ┆         │\n",
       "╞══════════╪══════════╪════════════╪════════════╪═══╪════════════╪═══════════╪═══════════╪═════════╡\n",
       "│ 31042685 ┆ 31041065 ┆ 1          ┆ 3          ┆ … ┆ <instructi ┆ -5.53125  ┆ -9.25     ┆ true    │\n",
       "│          ┆          ┆            ┆            ┆   ┆ ons>Your   ┆           ┆           ┆         │\n",
       "│          ┆          ┆            ┆            ┆   ┆ goal is to ┆           ┆           ┆         │\n",
       "│          ┆          ┆            ┆            ┆   ┆ …          ┆           ┆           ┆         │\n",
       "│ 39808386 ┆ 39806045 ┆ 1          ┆ 2          ┆ … ┆ <instructi ┆ -13.5     ┆ -17.25    ┆ true    │\n",
       "│          ┆          ┆            ┆            ┆   ┆ ons>Your   ┆           ┆           ┆         │\n",
       "│          ┆          ┆            ┆            ┆   ┆ goal is to ┆           ┆           ┆         │\n",
       "│          ┆          ┆            ┆            ┆   ┆ …          ┆           ┆           ┆         │\n",
       "│ 34895313 ┆ 34894926 ┆ 1          ┆ 2          ┆ … ┆ <instructi ┆ -4.5      ┆ -10.4375  ┆ true    │\n",
       "│          ┆          ┆            ┆            ┆   ┆ ons>Your   ┆           ┆           ┆         │\n",
       "│          ┆          ┆            ┆            ┆   ┆ goal is to ┆           ┆           ┆         │\n",
       "│          ┆          ┆            ┆            ┆   ┆ …          ┆           ┆           ┆         │\n",
       "│ 27699072 ┆ 27697619 ┆ 1          ┆ 2          ┆ … ┆ <instructi ┆ -2.828125 ┆ -3.0      ┆ true    │\n",
       "│          ┆          ┆            ┆            ┆   ┆ ons>Your   ┆           ┆           ┆         │\n",
       "│          ┆          ┆            ┆            ┆   ┆ goal is to ┆           ┆           ┆         │\n",
       "│          ┆          ┆            ┆            ┆   ┆ …          ┆           ┆           ┆         │\n",
       "│ 5348842  ┆ 5347984  ┆ 1          ┆ 3          ┆ … ┆ <instructi ┆ -9.0625   ┆ -9.25     ┆ true    │\n",
       "│          ┆          ┆            ┆            ┆   ┆ ons>Your   ┆           ┆           ┆         │\n",
       "│          ┆          ┆            ┆            ┆   ┆ goal is to ┆           ┆           ┆         │\n",
       "│          ┆          ┆            ┆            ┆   ┆ …          ┆           ┆           ┆         │\n",
       "│ …        ┆ …        ┆ …          ┆ …          ┆ … ┆ …          ┆ …         ┆ …         ┆ …       │\n",
       "│ 15301562 ┆ 15301600 ┆ 1          ┆ 2          ┆ … ┆ <instructi ┆ -13.3125  ┆ -14.5625  ┆ true    │\n",
       "│          ┆          ┆            ┆            ┆   ┆ ons>Your   ┆           ┆           ┆         │\n",
       "│          ┆          ┆            ┆            ┆   ┆ goal is to ┆           ┆           ┆         │\n",
       "│          ┆          ┆            ┆            ┆   ┆ …          ┆           ┆           ┆         │\n",
       "│ 36831606 ┆ 36830894 ┆ 1          ┆ 3          ┆ … ┆ <instructi ┆ -4.09375  ┆ -6.71875  ┆ true    │\n",
       "│          ┆          ┆            ┆            ┆   ┆ ons>Your   ┆           ┆           ┆         │\n",
       "│          ┆          ┆            ┆            ┆   ┆ goal is to ┆           ┆           ┆         │\n",
       "│          ┆          ┆            ┆            ┆   ┆ …          ┆           ┆           ┆         │\n",
       "│ 29624080 ┆ 29623860 ┆ 1          ┆ 2          ┆ … ┆ <instructi ┆ -5.40625  ┆ -3.609375 ┆ false   │\n",
       "│          ┆          ┆            ┆            ┆   ┆ ons>Your   ┆           ┆           ┆         │\n",
       "│          ┆          ┆            ┆            ┆   ┆ goal is to ┆           ┆           ┆         │\n",
       "│          ┆          ┆            ┆            ┆   ┆ …          ┆           ┆           ┆         │\n",
       "│ 31660200 ┆ 31659387 ┆ 1          ┆ 3          ┆ … ┆ <instructi ┆ -5.625    ┆ -7.5      ┆ true    │\n",
       "│          ┆          ┆            ┆            ┆   ┆ ons>Your   ┆           ┆           ┆         │\n",
       "│          ┆          ┆            ┆            ┆   ┆ goal is to ┆           ┆           ┆         │\n",
       "│          ┆          ┆            ┆            ┆   ┆ …          ┆           ┆           ┆         │\n",
       "│ 25421107 ┆ 25420617 ┆ 1          ┆ 3          ┆ … ┆ <instructi ┆ -4.40625  ┆ -5.125    ┆ true    │\n",
       "│          ┆          ┆            ┆            ┆   ┆ ons>Your   ┆           ┆           ┆         │\n",
       "│          ┆          ┆            ┆            ┆   ┆ goal is to ┆           ┆           ┆         │\n",
       "│          ┆          ┆            ┆            ┆   ┆ …          ┆           ┆           ┆         │\n",
       "└──────────┴──────────┴────────────┴────────────┴───┴────────────┴───────────┴───────────┴─────────┘"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds = ds.with_columns(\n",
    "    (pl.col(\"chosen_reward\") > pl.col(\"rejected_reward\")).alias(\"correct\")\n",
    ")\n",
    "\n",
    "ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy: 76.00%\n"
     ]
    }
   ],
   "source": [
    "ds.group_by(\"correct\").agg(pl.col(\"correct\").count().alias(\"count\"))\n",
    "\n",
    "print(f\"Accuracy: {ds['correct'].mean() * 100:.2f}%\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "59e75a9f8e2b4f4887f652e0e0f4dafa",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/4 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline\n",
    "import torch\n",
    "from tqdm import tqdm\n",
    "\n",
    "model = AutoModelForSequenceClassification.from_pretrained(\n",
    "    \"./models/llama_32_8b_merged\",\n",
    "    device_map=\"auto\",\n",
    "    attn_implementation=\"flash_attention_2\",\n",
    "    torch_dtype=torch.bfloat16,\n",
    ")\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"./models/llama_32_8b_merged\")\n",
    "\n",
    "pipe = pipeline(\"text-classification\", model=model, tokenizer=tokenizer)\n",
    "\n",
    "\n",
    "def run_inference_transformers(prompts: list[str]) -> list[float]:\n",
    "    results = []\n",
    "    for prompt in tqdm(prompts, desc=\"Running inference\", unit=\"prompt\"):\n",
    "        result = pipe(prompt)\n",
    "        results.append(result[0][\"score\"])\n",
    "    return results\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Running inference:   1%|          | 10/1000 [00:01<01:23, 11.80prompt/s]You seem to be using the pipelines sequentially on GPU. In order to maximize efficiency please use a dataset\n",
      "Running inference: 100%|██████████| 1000/1000 [00:30<00:00, 32.62prompt/s]\n",
      "Running inference: 100%|██████████| 1000/1000 [00:26<00:00, 38.22prompt/s]\n"
     ]
    }
   ],
   "source": [
    "chosen_rewards = run_inference_transformers(ds[\"chosen_prompt\"].to_list())\n",
    "rejected_rewards = run_inference_transformers(ds[\"rejected_prompt\"].to_list())\n",
    "\n",
    "ds = ds.with_columns(\n",
    "    pl.Series(name=\"chosen_reward\", values=chosen_rewards),\n",
    "    pl.Series(name=\"rejected_reward\", values=rejected_rewards),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy (Transformers): 75.70%\n"
     ]
    }
   ],
   "source": [
    "ds = ds.with_columns(\n",
    "    (pl.col(\"chosen_reward\") > pl.col(\"rejected_reward\")).alias(\"correct\")\n",
    ")\n",
    "\n",
    "ds.group_by(\"correct\").agg(pl.col(\"correct\").count().alias(\"count\"))\n",
    "\n",
    "print(f\"Accuracy (Transformers): {ds['correct'].mean() * 100:.2f}%\")"
   ]
  }
 ],
 "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
}