train.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6cc7d98d-23c5-4fb3-86eb-4e0665382bb1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[2mInstalled \u001b[1mPython 3.13.11\u001b[0m \u001b[2min 713ms\u001b[0m\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mcpython-3.13.11-linux-x86_64-gnu\u001b[0m (python3.13)\n",
      "\u001b[1m\u001b[33mwarning\u001b[39m\u001b[0m\u001b[1m:\u001b[0m \u001b[1m`\u001b[36m/root/.local/bin\u001b[39m` is not on your PATH. To use installed Python executables, run `\u001b[32mexport PATH=\"/root/.local/bin:$PATH\"\u001b[39m` or `\u001b[32muv python update-shell\u001b[39m`.\u001b[0m\n",
      "Pinned `\u001b[36m/.python-version\u001b[39m` to `\u001b[32m3.13\u001b[39m`\n"
     ]
    }
   ],
   "source": [
    "!uv python install 313\n",
    "!uv python pin 313"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "2d741065-01e2-4551-8a7e-70fecca6759e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using CPython \u001b[36m3.13.11\u001b[39m\u001b[36m\u001b[39m\n",
      "Creating virtual environment at: \u001b[36m.venv\u001b[39m\n",
      "Activate with: \u001b[32msource .venv/bin/activate\u001b[39m\n",
      "\u001b[2K\u001b[2mResolved \u001b[1m30 packages\u001b[0m \u001b[2min 310ms\u001b[0m\u001b[0m                                        \u001b[0m\n",
      "\u001b[2K\u001b[2mPrepared \u001b[1m30 packages\u001b[0m \u001b[2min 280ms\u001b[0m\u001b[0m                                            \n",
      "\u001b[2K\u001b[2mInstalled \u001b[1m30 packages\u001b[0m \u001b[2min 56ms\u001b[0m\u001b[0m                               \u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1masttokens\u001b[0m\u001b[2m==3.0.2\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mcomm\u001b[0m\u001b[2m==0.2.3\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mdebugpy\u001b[0m\u001b[2m==1.8.21\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mdecorator\u001b[0m\u001b[2m==5.3.1\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mexecuting\u001b[0m\u001b[2m==2.2.1\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mipykernel\u001b[0m\u001b[2m==7.3.0\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mipython\u001b[0m\u001b[2m==9.15.0\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mipython-pygments-lexers\u001b[0m\u001b[2m==1.1.1\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mjedi\u001b[0m\u001b[2m==0.20.0\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mjupyter-client\u001b[0m\u001b[2m==8.9.1\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mjupyter-core\u001b[0m\u001b[2m==5.9.1\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mmatplotlib-inline\u001b[0m\u001b[2m==0.2.2\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mnest-asyncio2\u001b[0m\u001b[2m==1.7.2\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mpackaging\u001b[0m\u001b[2m==26.2\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mparso\u001b[0m\u001b[2m==0.8.7\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mpexpect\u001b[0m\u001b[2m==4.9.0\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mplatformdirs\u001b[0m\u001b[2m==4.11.0\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mprompt-toolkit\u001b[0m\u001b[2m==3.0.53\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mpsutil\u001b[0m\u001b[2m==7.2.2\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mptyprocess\u001b[0m\u001b[2m==0.7.0\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mpure-eval\u001b[0m\u001b[2m==0.2.3\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mpygments\u001b[0m\u001b[2m==2.20.0\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mpython-dateutil\u001b[0m\u001b[2m==2.9.0.post0\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mpyzmq\u001b[0m\u001b[2m==27.1.0\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1msix\u001b[0m\u001b[2m==1.17.0\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mstack-data\u001b[0m\u001b[2m==0.6.3\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mtornado\u001b[0m\u001b[2m==6.5.7\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mtraitlets\u001b[0m\u001b[2m==5.15.1\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mtyping-extensions\u001b[0m\u001b[2m==4.16.0\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mwcwidth\u001b[0m\u001b[2m==0.8.2\u001b[0m\n",
      "Installed kernelspec uv-kernel in /root/.local/share/jupyter/kernels/uv-kernel\n"
     ]
    }
   ],
   "source": [
    "!uv venv --clear\n",
    "!uv pip install ipykernel\n",
    "!uv run python -m ipykernel install --user --name=uv-kernel --display-name \"Python (uv venv)\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ae35bf68-b4b1-435e-8a79-61bd140bc77e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[2K\u001b[2mResolved \u001b[1m83 packages\u001b[0m \u001b[2min 144ms\u001b[0m\u001b[0m                                        \u001b[0m\n",
      "\u001b[2K\u001b[2mPrepared \u001b[1m8 packages\u001b[0m \u001b[2min 15.25s\u001b[0m\u001b[0m                                            \n",
      "\u001b[2K\u001b[2mInstalled \u001b[1m77 packages\u001b[0m \u001b[2min 119ms\u001b[0m\u001b[0m                              \u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1maccelerate\u001b[0m\u001b[2m==1.14.0\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1maiohappyeyeballs\u001b[0m\u001b[2m==2.7.1\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1maiohttp\u001b[0m\u001b[2m==3.14.3\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1maiosignal\u001b[0m\u001b[2m==1.4.0\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mannotated-doc\u001b[0m\u001b[2m==0.0.5\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1manyio\u001b[0m\u001b[2m==4.14.2\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mattrs\u001b[0m\u001b[2m==26.1.0\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mcertifi\u001b[0m\u001b[2m==2026.7.22\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mcharset-normalizer\u001b[0m\u001b[2m==3.4.9\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mclick\u001b[0m\u001b[2m==8.4.2\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mcuda-bindings\u001b[0m\u001b[2m==13.3.1\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mcuda-pathfinder\u001b[0m\u001b[2m==1.6.0\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mcuda-toolkit\u001b[0m\u001b[2m==13.0.3.0\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mdatasets\u001b[0m\u001b[2m==5.0.1\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mdill\u001b[0m\u001b[2m==0.4.1\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mevaluate\u001b[0m\u001b[2m==0.4.6\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mfastrand\u001b[0m\u001b[2m==3.1.0\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mfilelock\u001b[0m\u001b[2m==3.32.2\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mfrozenlist\u001b[0m\u001b[2m==1.8.0\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mfsspec\u001b[0m\u001b[2m==2026.6.0\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mh11\u001b[0m\u001b[2m==0.16.0\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mhf-xet\u001b[0m\u001b[2m==1.5.2\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mhttpcore\u001b[0m\u001b[2m==1.0.9\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mhttpx\u001b[0m\u001b[2m==0.28.1\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mhuggingface-hub\u001b[0m\u001b[2m==1.26.0\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1midna\u001b[0m\u001b[2m==3.18\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mjinja2\u001b[0m\u001b[2m==3.1.6\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mjoblib\u001b[0m\u001b[2m==1.5.3\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mmarkdown-it-py\u001b[0m\u001b[2m==4.2.0\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mmarkupsafe\u001b[0m\u001b[2m==3.0.3\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mmdurl\u001b[0m\u001b[2m==0.1.2\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mmpmath\u001b[0m\u001b[2m==1.3.0\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mmultidict\u001b[0m\u001b[2m==6.7.1\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mmultiprocess\u001b[0m\u001b[2m==0.70.19\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mnarwhals\u001b[0m\u001b[2m==2.24.0\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mnetworkx\u001b[0m\u001b[2m==3.6.1\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mnumpy\u001b[0m\u001b[2m==2.5.1\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mnvidia-cublas\u001b[0m\u001b[2m==13.1.1.3\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mnvidia-cuda-cupti\u001b[0m\u001b[2m==13.0.85\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mnvidia-cuda-nvrtc\u001b[0m\u001b[2m==13.0.88\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mnvidia-cuda-runtime\u001b[0m\u001b[2m==13.0.96\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mnvidia-cudnn-cu13\u001b[0m\u001b[2m==9.20.0.48\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mnvidia-cufft\u001b[0m\u001b[2m==12.0.0.61\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mnvidia-cufile\u001b[0m\u001b[2m==1.15.1.6\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mnvidia-curand\u001b[0m\u001b[2m==10.4.0.35\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mnvidia-cusolver\u001b[0m\u001b[2m==12.0.4.66\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mnvidia-cusparse\u001b[0m\u001b[2m==12.6.3.3\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mnvidia-cusparselt-cu13\u001b[0m\u001b[2m==0.8.1\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mnvidia-nccl-cu13\u001b[0m\u001b[2m==2.29.7\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mnvidia-nvjitlink\u001b[0m\u001b[2m==13.3.33\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mnvidia-nvshmem-cu13\u001b[0m\u001b[2m==3.4.5\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mnvidia-nvtx\u001b[0m\u001b[2m==13.0.85\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mpandas\u001b[0m\u001b[2m==3.0.5\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mpropcache\u001b[0m\u001b[2m==0.5.2\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mprotobuf\u001b[0m\u001b[2m==7.35.1\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mpyarrow\u001b[0m\u001b[2m==25.0.0\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mpyyaml\u001b[0m\u001b[2m==6.0.3\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mregex\u001b[0m\u001b[2m==2026.7.19\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mrequests\u001b[0m\u001b[2m==2.34.2\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mrich\u001b[0m\u001b[2m==15.0.0\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1msafetensors\u001b[0m\u001b[2m==0.8.0\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mscikit-learn\u001b[0m\u001b[2m==1.9.0\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mscipy\u001b[0m\u001b[2m==1.18.0\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1msentencepiece\u001b[0m\u001b[2m==0.2.2\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1msetuptools\u001b[0m\u001b[2m==83.0.0\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mshellingham\u001b[0m\u001b[2m==1.5.4\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1msympy\u001b[0m\u001b[2m==1.14.0\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mthreadpoolctl\u001b[0m\u001b[2m==3.6.0\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mtokenizers\u001b[0m\u001b[2m==0.22.2\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mtorch\u001b[0m\u001b[2m==2.13.0\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mtqdm\u001b[0m\u001b[2m==4.70.0\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mtransformers\u001b[0m\u001b[2m==5.14.1\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mtriton\u001b[0m\u001b[2m==3.7.1\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mtyper\u001b[0m\u001b[2m==0.27.0\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1murllib3\u001b[0m\u001b[2m==2.7.0\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1mxxhash\u001b[0m\u001b[2m==3.8.1\u001b[0m\n",
      " \u001b[32m+\u001b[39m \u001b[1myarl\u001b[0m\u001b[2m==1.24.5\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "!uv pip install transformers huggingface_hub datasets fastrand torch sentencepiece protobuf evaluate scikit-learn accelerate>1.1.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "345615b8-726e-4c9a-ab56-deb322632c82",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/.venv/lib/python3.13/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "from datasets import load_dataset\n",
    "from transformers import (\n",
    "    AutoTokenizer, \n",
    "    AutoModelForSequenceClassification, \n",
    "    TrainingArguments, \n",
    "    Trainer,\n",
    "    DataCollatorWithPadding\n",
    ")\n",
    "from huggingface_hub import notebook_login\n",
    "import evaluate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "033f1b47-2771-48ff-b89f-fcc3d93ab47d",
   "metadata": {},
   "outputs": [],
   "source": [
    "notebook_login()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f0ff7440-f91c-4705-9d32-d37ac1089239",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading weights: 100%|██████████| 136/136 [00:00<00:00, 4473.68it/s]\n",
      "[transformers] \u001b[1mModernBertForSequenceClassification LOAD REPORT\u001b[0m from: jhu-clsp/ettin-encoder-150m\n",
      "Key               | Status     | \n",
      "------------------+------------+-\n",
      "decoder.bias      | UNEXPECTED | \n",
      "decoder.weight    | UNEXPECTED | \n",
      "classifier.weight | MISSING    | \n",
      "classifier.bias   | MISSING    | \n",
      "\n",
      "Notes:\n",
      "- UNEXPECTED:\tcan be ignored when loading from different task/architecture; not ok if you expect identical arch.\n",
      "- MISSING:\tthose params were newly initialized because missing from the checkpoint. Consider training on your downstream task.\n"
     ]
    }
   ],
   "source": [
    "model_name = \"jhu-clsp/ettin-encoder-150m\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "model = AutoModelForSequenceClassification.from_pretrained(\n",
    "    model_name, num_labels=2,\n",
    "    dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "46b5f750-1b36-4437-b7cb-26331d92c595",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = load_dataset(\"hanzceo/JOSS-data-L2\")[\"train\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "37f80b16-499f-4002-93c5-a534c4f7b289",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = dataset.train_test_split(test_size=0.2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4a8f8b10-e197-4d58-a818-2f5e8ba77d0e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['sentence1', 'score'],\n",
       "        num_rows: 1600\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['sentence1', 'score'],\n",
       "        num_rows: 400\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "fa7253f7-05a7-40a4-9727-d107adac325e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Map: 100%|██████████| 1600/1600 [00:00<00:00, 1706.17 examples/s]\n",
      "Map: 100%|██████████| 400/400 [00:00<00:00, 1396.96 examples/s]\n"
     ]
    }
   ],
   "source": [
    "def tokenize_function(examples):\n",
    "    return tokenizer(examples[\"sentence1\"], truncation=True)\n",
    "tokenized_datasets = dataset.map(tokenize_function, batched=True)\n",
    "data_collator = DataCollatorWithPadding(tokenizer=tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "49744185-47ff-460a-9a1b-90f33f2bc991",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['sentence1', 'labels', 'input_ids', 'attention_mask'],\n",
       "        num_rows: 1600\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['sentence1', 'labels', 'input_ids', 'attention_mask'],\n",
       "        num_rows: 400\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenized_datasets = tokenized_datasets.rename_column(\"score\", \"labels\")\n",
    "\n",
    "tokenized_datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "93531513-4e7f-4d16-8292-8cecb44e0aa3",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading builder script: 4.20kB [00:00, 14.3MB/s]\n"
     ]
    }
   ],
   "source": [
    "metric = evaluate.load(\"accuracy\")\n",
    "def compute_metrics(eval_pred):\n",
    "    logits, labels = eval_pred\n",
    "    predictions = np.argmax(logits, axis=-1)\n",
    "    return metric.compute(predictions=predictions, references=labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "b663bbee-15cb-44a9-9393-a78cf051c0fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "training_args = TrainingArguments(\n",
    "    output_dir=\"models/ettin-classification\",\n",
    "    eval_strategy=\"epoch\",\n",
    "    learning_rate=2e-5,\n",
    "    per_device_train_batch_size=1,\n",
    "    per_device_eval_batch_size=1,\n",
    "    num_train_epochs=3,\n",
    "    weight_decay=0.01,\n",
    "    bf16=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "685b3f75-3a68-4965-bf28-ba3e70b97618",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='4800' max='4800' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [4800/4800 06:43, Epoch 3/3]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Epoch</th>\n",
       "      <th>Training Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "      <th>Accuracy</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>0.626540</td>\n",
       "      <td>0.637205</td>\n",
       "      <td>0.847500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>0.630978</td>\n",
       "      <td>0.628983</td>\n",
       "      <td>0.857500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3</td>\n",
       "      <td>0.539487</td>\n",
       "      <td>0.631988</td>\n",
       "      <td>0.857500</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Writing model shards: 100%|██████████| 1/1 [00:00<00:00,  7.39it/s]\n",
      "Writing model shards: 100%|██████████| 1/1 [00:00<00:00,  7.52it/s]\n",
      "Writing model shards: 100%|██████████| 1/1 [00:00<00:00,  7.01it/s]\n",
      "Writing model shards: 100%|██████████| 1/1 [00:00<00:00,  5.47it/s]\n",
      "Writing model shards: 100%|██████████| 1/1 [00:00<00:00,  4.76it/s]\n",
      "Writing model shards: 100%|██████████| 1/1 [00:00<00:00,  9.00it/s]\n",
      "Writing model shards: 100%|██████████| 1/1 [00:00<00:00,  6.39it/s]\n",
      "Writing model shards: 100%|██████████| 1/1 [00:00<00:00,  7.45it/s]\n",
      "Writing model shards: 100%|██████████| 1/1 [00:00<00:00,  6.17it/s]\n",
      "Writing model shards: 100%|██████████| 1/1 [00:00<00:00,  7.71it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=4800, training_loss=0.6138824494679769, metrics={'train_runtime': 404.0617, 'train_samples_per_second': 11.879, 'train_steps_per_second': 11.879, 'total_flos': 6650877695864220.0, 'train_loss': 0.6138824494679769, 'epoch': 3.0})"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    data_collator=data_collator,\n",
    "    train_dataset=tokenized_datasets[\"train\"].shuffle(seed=42),\n",
    "    eval_dataset=tokenized_datasets[\"test\"].shuffle(seed=42),\n",
    "    compute_metrics=compute_metrics,\n",
    ")\n",
    "\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "250248f5-021e-4247-8b2d-1d1098cf8a31",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ModernBertForSequenceClassification(\n",
       "  (model): ModernBertModel(\n",
       "    (embeddings): ModernBertEmbeddings(\n",
       "      (tok_embeddings): Embedding(50368, 768, padding_idx=50283)\n",
       "      (norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True, bias=False)\n",
       "      (drop): Dropout(p=0.0, inplace=False)\n",
       "    )\n",
       "    (layers): ModuleList(\n",
       "      (0): ModernBertEncoderLayer(\n",
       "        (attn_norm): Identity()\n",
       "        (attn): ModernBertAttention(\n",
       "          (Wqkv): Linear(in_features=768, out_features=2304, bias=False)\n",
       "          (Wo): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (out_drop): Identity()\n",
       "        )\n",
       "        (mlp_norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True, bias=False)\n",
       "        (mlp): ModernBertMLP(\n",
       "          (Wi): Linear(in_features=768, out_features=2304, bias=False)\n",
       "          (act): GELUActivation()\n",
       "          (drop): Dropout(p=0.0, inplace=False)\n",
       "          (Wo): Linear(in_features=1152, out_features=768, bias=False)\n",
       "        )\n",
       "      )\n",
       "      (1-21): 21 x ModernBertEncoderLayer(\n",
       "        (attn_norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True, bias=False)\n",
       "        (attn): ModernBertAttention(\n",
       "          (Wqkv): Linear(in_features=768, out_features=2304, bias=False)\n",
       "          (Wo): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (out_drop): Identity()\n",
       "        )\n",
       "        (mlp_norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True, bias=False)\n",
       "        (mlp): ModernBertMLP(\n",
       "          (Wi): Linear(in_features=768, out_features=2304, bias=False)\n",
       "          (act): GELUActivation()\n",
       "          (drop): Dropout(p=0.0, inplace=False)\n",
       "          (Wo): Linear(in_features=1152, out_features=768, bias=False)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (final_norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True, bias=False)\n",
       "    (rotary_emb): ModernBertRotaryEmbedding()\n",
       "  )\n",
       "  (head): ModernBertPredictionHead(\n",
       "    (dense): Linear(in_features=768, out_features=768, bias=False)\n",
       "    (act): GELUActivation()\n",
       "    (norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True, bias=False)\n",
       "  )\n",
       "  (drop): Dropout(p=0.0, inplace=False)\n",
       "  (classifier): Linear(in_features=768, out_features=2, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "91374ca9-04ca-4ab7-a136-a5b329071426",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Writing model shards: 100%|██████████| 1/1 [00:00<00:00,  9.34it/s]\n",
      "Processing Files (0 / 0): |          |  0.00B /  0.00B            \n",
      "Processing Files (0 / 1):   0%|          |  603kB /  299MB, 59.6kB/s  \n",
      "Processing Files (0 / 1):   1%|          | 3.02MB /  299MB,  297kB/s  \n",
      "Processing Files (0 / 1):   2%|▏         | 7.24MB /  299MB,  708kB/s  \n",
      "Processing Files (0 / 1):   7%|▋         | 22.3MB /  299MB, 2.18MB/s  \n",
      "Processing Files (0 / 1):  10%|▉         | 29.5MB /  299MB, 2.86MB/s  \n",
      "Processing Files (0 / 1):  14%|█▍        | 42.2MB /  299MB, 4.06MB/s  \n",
      "Processing Files (0 / 1):  19%|█▊        | 55.5MB /  299MB, 5.31MB/s  \n",
      "Processing Files (0 / 1):  24%|██▎       | 70.5MB /  299MB, 6.72MB/s  \n",
      "Processing Files (0 / 1):  27%|██▋       | 80.8MB /  299MB, 7.64MB/s  \n",
      "Processing Files (0 / 1):  30%|███       | 90.4MB /  299MB, 8.50MB/s  \n",
      "Processing Files (0 / 1):  33%|███▎      |  100MB /  299MB, 9.34MB/s  \n",
      "Processing Files (0 / 1):  37%|███▋      |  112MB /  299MB, 10.4MB/s  \n",
      "Processing Files (0 / 1):  40%|███▉      |  119MB /  299MB, 10.9MB/s  \n",
      "Processing Files (0 / 1):  41%|████      |  122MB /  299MB, 11.1MB/s  \n",
      "Processing Files (0 / 1):  42%|████▏     |  125MB /  299MB, 11.4MB/s  \n",
      "Processing Files (0 / 1):  43%|████▎     |  128MB /  299MB, 11.5MB/s  \n",
      "Processing Files (0 / 1):  45%|████▍     |  134MB /  299MB, 11.9MB/s  \n",
      "Processing Files (0 / 1):  46%|████▌     |  138MB /  299MB, 11.7MB/s  \n",
      "Processing Files (0 / 1):  48%|████▊     |  144MB /  299MB, 12.0MB/s  \n",
      "Processing Files (0 / 1):  49%|████▉     |  147MB /  299MB, 12.2MB/s  \n",
      "Processing Files (0 / 1):  51%|█████▏    |  153MB /  299MB, 12.6MB/s  \n",
      "Processing Files (0 / 1):  54%|█████▎    |  161MB /  299MB, 13.2MB/s  \n",
      "Processing Files (0 / 1):  57%|█████▋    |  171MB /  299MB, 14.1MB/s  \n",
      "Processing Files (0 / 1):  63%|██████▎   |  187MB /  299MB, 15.4MB/s  \n",
      "Processing Files (0 / 1):  67%|██████▋   |  201MB /  299MB, 16.5MB/s  \n",
      "Processing Files (0 / 1):  72%|███████▏  |  216MB /  299MB, 17.7MB/s  \n",
      "Processing Files (0 / 1):  78%|███████▊  |  234MB /  299MB, 19.1MB/s  \n",
      "Processing Files (0 / 1):  81%|████████▏ |  244MB /  299MB, 19.9MB/s  \n",
      "Processing Files (0 / 1):  83%|████████▎ |  249MB /  299MB, 20.1MB/s  \n",
      "Processing Files (0 / 1):  84%|████████▍ |  253MB /  299MB, 20.3MB/s  \n",
      "Processing Files (0 / 1):  86%|████████▌ |  258MB /  299MB, 20.6MB/s  \n",
      "Processing Files (0 / 1):  88%|████████▊ |  262MB /  299MB, 20.8MB/s  \n",
      "Processing Files (0 / 1):  89%|████████▉ |  267MB /  299MB, 21.1MB/s  \n",
      "Processing Files (0 / 1):  91%|█████████ |  272MB /  299MB, 21.3MB/s  \n",
      "Processing Files (0 / 1):  94%|█████████▍|  281MB /  299MB, 21.9MB/s  \n",
      "Processing Files (0 / 1):  95%|█████████▌|  284MB /  299MB, 22.0MB/s  \n",
      "Processing Files (0 / 1):  96%|█████████▌|  286MB /  299MB, 22.0MB/s  \n",
      "Processing Files (0 / 1):  96%|█████████▋|  288MB /  299MB, 22.0MB/s  \n",
      "Processing Files (0 / 1):  97%|█████████▋|  291MB /  299MB, 22.0MB/s  \n",
      "Processing Files (0 / 1):  98%|█████████▊|  293MB /  299MB, 22.0MB/s  \n",
      "Processing Files (0 / 1):  99%|█████████▊|  295MB /  299MB, 22.0MB/s  \n",
      "Processing Files (0 / 1): 100%|█████████▉|  298MB /  299MB, 22.1MB/s  \n",
      "Processing Files (0 / 1): 100%|█████████▉|  299MB /  299MB, 22.0MB/s  \n",
      "Processing Files (1 / 1): 100%|██████████|  299MB /  299MB, 20.3MB/s  \n",
      "Processing Files (1 / 1): 100%|██████████|  299MB /  299MB, 19.5MB/s  \n",
      "New Data Upload: 100%|██████████|  299MB /  299MB, 19.5MB/s  \n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "CommitInfo(commit_url='https://huggingface.co/hanzceo/JOSS-dev2/commit/d1cf8f7444d284978c4e98761d2ed19b85467afe', commit_message='Upload ModernBertForSequenceClassification', commit_description='', oid='d1cf8f7444d284978c4e98761d2ed19b85467afe', pr_url=None, repo_url=RepoUrl('https://huggingface.co/hanzceo/JOSS-dev2', endpoint='https://huggingface.co', repo_type='model', repo_id='hanzceo/JOSS-dev2'), pr_revision=None, pr_num=None)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.push_to_hub(\"hanzceo/JOSS-dev2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "ed10d1ca-ada0-43db-ac78-357a865d2b4a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "CommitInfo(commit_url='https://huggingface.co/hanzceo/JOSS-dev2/commit/c7d96e5e80839e397e9bc04177120915a4ea4b27', commit_message='Upload tokenizer', commit_description='', oid='c7d96e5e80839e397e9bc04177120915a4ea4b27', pr_url=None, repo_url=RepoUrl('https://huggingface.co/hanzceo/JOSS-dev2', endpoint='https://huggingface.co', repo_type='model', repo_id='hanzceo/JOSS-dev2'), pr_revision=None, pr_num=None)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer.push_to_hub(\"hanzceo/JOSS-dev2\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (uv venv)",
   "language": "python",
   "name": "uv-kernel"
  },
  "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.13.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}