Sync from GitHub Actions

notebooks/01_dataset_experiment.ipynb

@@ -1,786 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": 1,
-   "id": "27ea1a10",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "!pip install -q torch torchvision transformers datasets pillow pandas scikit-learn tqdm huggingface_hub matplotlib"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 2,
-   "id": "d246862c",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import os\n",
-    "import json\n",
-    "import random\n",
-    "import time\n",
-    "from pathlib import Path\n",
-    "\n",
-    "import numpy as np\n",
-    "import pandas as pd\n",
-    "import torch\n",
-    "import torch.nn as nn\n",
-    "import torch.nn.functional as F\n",
-    "\n",
-    "from PIL import Image\n",
-    "from tqdm.auto import tqdm\n",
-    "from datasets import load_dataset\n",
-    "from torch.utils.data import Dataset, DataLoader\n",
-    "from transformers import CLIPModel, CLIPImageProcessor, get_cosine_schedule_with_warmup\n",
-    "\n",
-    "from sklearn.model_selection import train_test_split\n",
-    "from sklearn.metrics import accuracy_score, f1_score, confusion_matrix, classification_report"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "7e09b0d2",
-   "metadata": {},
-   "source": [
-    "_Config_"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 3,
-   "id": "697fa289",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "DATASET_NAME = \"ashraq/fashion-product-images-small\"\n",
-    "MODEL_NAME = \"openai/clip-vit-base-patch32\"\n",
-    "TASKS = [\n",
-    "    \"gender\",\n",
-    "    \"masterCategory\",\n",
-    "    \"subCategory\",\n",
-    "    \"articleType\",\n",
-    "    \"baseColour\",\n",
-    "    \"season\",\n",
-    "    \"usage\"\n",
-    "]\n",
-    "\n",
-    "SEED = 42\n",
-    "TRAIN_RATIO = 0.70\n",
-    "VAL_RATIO = 0.15\n",
-    "TEST_RATIO = 0.15\n",
-    "\n",
-    "BATCH_SIZE = 32\n",
-    "EPOCHS = 5\n",
-    "\n",
-    "HEAD_LR = 3e-4\n",
-    "BACKBONE_LR = 1e-5\n",
-    "WEIGHT_DECAY = 1e-2\n",
-    "\n",
-    "HIDDEN_DIM = 512\n",
-    "DROPOUT = 0.20\n",
-    "\n",
-    "UNFREEZE_LAST_N_VISION_LAYERS = 2\n",
-    "\n",
-    "USE_CLASS_WEIGHTS = True\n",
-    "USE_AMP = True\n",
-    "\n",
-    "MAX_GRAD_NORM = 1.0\n",
-    "EARLY_STOPPING_PATIENCE = 2\n",
-    "NUM_WORKERS = 2\n",
-    "DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\""
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 4,
-   "id": "ef7f600c",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Device: cuda\n",
-      "Model: openai/clip-vit-base-patch32\n"
-     ]
-    }
-   ],
-   "source": [
-    "OUTPUT_DIR = Path(\"artifacts/models/autocatalogai_clip\")\n",
-    "EVAL_DIR = Path(\"artifacts/evaluation\")\n",
-    "PLOT_DIR = Path(\"artifacts/plots\")\n",
-    "PROCESSED_DIR = Path(\"data/processed\")\n",
-    "\n",
-    "for directory in [OUTPUT_DIR, EVAL_DIR, PLOT_DIR, PROCESSED_DIR]:\n",
-    "    directory.mkdir(parents=True, exist_ok=True)\n",
-    "\n",
-    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n",
-    "\n",
-    "print(\"Device:\", DEVICE)\n",
-    "print(\"Model:\", MODEL_NAME)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 5,
-   "id": "d5242fdc",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def set_seed(seed):\n",
-    "    random.seed(seed)\n",
-    "    np.random.seed(seed)\n",
-    "    torch.manual_seed(seed)\n",
-    "    \n",
-    "    if torch.cuda.is_available():\n",
-    "        torch.cuda.manual_seed_all(seed)\n",
-    "    torch.backends.cudnn.benchmark = True\n",
-    "\n",
-    "\n",
-    "set_seed(SEED)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "f07779e1",
-   "metadata": {},
-   "source": [
-    "#### Load Full Dataset"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 6,
-   "id": "280fb69a",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/usr/local/lib/python3.12/dist-packages/huggingface_hub/utils/_auth.py:138: UserWarning: \n",
-      "Error while fetching `HF_TOKEN` secret value from your vault: 'Requesting secret HF_TOKEN timed out. Secrets can only be fetched when running from the Colab UI.'.\n",
-      "You are not authenticated with the Hugging Face Hub in this notebook.\n",
-      "If the error persists, please let us know by opening an issue on GitHub (https://github.com/huggingface/huggingface_hub/issues/new).\n",
-      "  warnings.warn(\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "f72208f4dc6348fc87d90c6c22769b0f",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "README.md:   0%|          | 0.00/867 [00:00<?, ?B/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.\n",
-      "WARNING:huggingface_hub.utils._http:Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "ee563a2427bc4eb9b91f160c593b5477",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "data/train-00000-of-00002-6cff4c59f91661(…):   0%|          | 0.00/136M [00:00<?, ?B/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "2960cb54d2ba4573a7ef82894e897ad6",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "data/train-00001-of-00002-bb459e5ac5f01e(…):   0%|          | 0.00/135M [00:00<?, ?B/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "7aa26aab19ff40e1a74b28de564323e2",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Generating train split:   0%|          | 0/44072 [00:00<?, ? examples/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Dataset({\n",
-      "    features: ['id', 'gender', 'masterCategory', 'subCategory', 'articleType', 'baseColour', 'season', 'year', 'usage', 'productDisplayName', 'image'],\n",
-      "    num_rows: 44072\n",
-      "})\n",
-      "['id', 'gender', 'masterCategory', 'subCategory', 'articleType', 'baseColour', 'season', 'year', 'usage', 'productDisplayName', 'image']\n",
-      "Total rows: 44072\n"
-     ]
-    }
-   ],
-   "source": [
-    "raw_dataset = load_dataset(DATASET_NAME, split=\"train\")\n",
-    "\n",
-    "print(raw_dataset)\n",
-    "print(raw_dataset.column_names)\n",
-    "print(\"Total rows:\", len(raw_dataset))"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "5edb84b5",
-   "metadata": {},
-   "source": [
-    "#### Clean Dataset"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 7,
-   "id": "bdf04d04",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "missing_columns = [task for task in TASKS if task not in raw_dataset.column_names]\n",
-    "\n",
-    "if \"image\" not in raw_dataset.column_names:\n",
-    "    raise ValueError(f\"Dataset must contain image column. Found: {raw_dataset.column_names}\")\n",
-    "\n",
-    "if missing_columns:\n",
-    "    raise ValueError(f\"Missing task columns: {missing_columns}\")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 8,
-   "id": "721ec563",
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "49bb2b38bea040a5896ceaf814335e21",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Filter:   0%|          | 0/44072 [00:00<?, ? examples/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Before cleaning: 44072\n",
-      "After cleaning: 44072\n"
-     ]
-    }
-   ],
-   "source": [
-    "def is_valid_row(row):\n",
-    "    for task in TASKS:\n",
-    "        value = row.get(task)\n",
-    "        if value is None:\n",
-    "            return False\n",
-    "        if str(value).strip() == \"\":\n",
-    "            return False\n",
-    "        \n",
-    "    return row.get(\"image\") is not None\n",
-    "\n",
-    "\n",
-    "clean_dataset = raw_dataset.filter(is_valid_row)\n",
-    "print(\"Before cleaning:\", len(raw_dataset))\n",
-    "print(\"After cleaning:\", len(clean_dataset))"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "59e5f003",
-   "metadata": {},
-   "source": [
-    "#### Create Metadata DataFrame"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 9,
-   "id": "17fecce2",
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/html": [
-       "\n",
-       "  <div id=\"df-819ad3e2-5ed3-4458-823f-695ea26d18b3\" class=\"colab-df-container\">\n",
-       "    <div>\n",
-       "<style scoped>\n",
-       "    .dataframe tbody tr th:only-of-type {\n",
-       "        vertical-align: middle;\n",
-       "    }\n",
-       "\n",
-       "    .dataframe tbody tr th {\n",
-       "        vertical-align: top;\n",
-       "    }\n",
-       "\n",
-       "    .dataframe thead th {\n",
-       "        text-align: right;\n",
-       "    }\n",
-       "</style>\n",
-       "<table border=\"1\" class=\"dataframe\">\n",
-       "  <thead>\n",
-       "    <tr style=\"text-align: right;\">\n",
-       "      <th></th>\n",
-       "      <th>gender</th>\n",
-       "      <th>masterCategory</th>\n",
-       "      <th>subCategory</th>\n",
-       "      <th>articleType</th>\n",
-       "      <th>baseColour</th>\n",
-       "      <th>season</th>\n",
-       "      <th>usage</th>\n",
-       "      <th>id</th>\n",
-       "      <th>productDisplayName</th>\n",
-       "      <th>dataset_idx</th>\n",
-       "    </tr>\n",
-       "  </thead>\n",
-       "  <tbody>\n",
-       "    <tr>\n",
-       "      <th>0</th>\n",
-       "      <td>Men</td>\n",
-       "      <td>Apparel</td>\n",
-       "      <td>Topwear</td>\n",
-       "      <td>Shirts</td>\n",
-       "      <td>Navy Blue</td>\n",
-       "      <td>Fall</td>\n",
-       "      <td>Casual</td>\n",
-       "      <td>15970</td>\n",
-       "      <td>Turtle Check Men Navy Blue Shirt</td>\n",
-       "      <td>0</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1</th>\n",
-       "      <td>Men</td>\n",
-       "      <td>Apparel</td>\n",
-       "      <td>Bottomwear</td>\n",
-       "      <td>Jeans</td>\n",
-       "      <td>Blue</td>\n",
-       "      <td>Summer</td>\n",
-       "      <td>Casual</td>\n",
-       "      <td>39386</td>\n",
-       "      <td>Peter England Men Party Blue Jeans</td>\n",
-       "      <td>1</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>2</th>\n",
-       "      <td>Women</td>\n",
-       "      <td>Accessories</td>\n",
-       "      <td>Watches</td>\n",
-       "      <td>Watches</td>\n",
-       "      <td>Silver</td>\n",
-       "      <td>Winter</td>\n",
-       "      <td>Casual</td>\n",
-       "      <td>59263</td>\n",
-       "      <td>Titan Women Silver Watch</td>\n",
-       "      <td>2</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>3</th>\n",
-       "      <td>Men</td>\n",
-       "      <td>Apparel</td>\n",
-       "      <td>Bottomwear</td>\n",
-       "      <td>Track Pants</td>\n",
-       "      <td>Black</td>\n",
-       "      <td>Fall</td>\n",
-       "      <td>Casual</td>\n",
-       "      <td>21379</td>\n",
-       "      <td>Manchester United Men Solid Black Track Pants</td>\n",
-       "      <td>3</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>4</th>\n",
-       "      <td>Men</td>\n",
-       "      <td>Apparel</td>\n",
-       "      <td>Topwear</td>\n",
-       "      <td>Tshirts</td>\n",
-       "      <td>Grey</td>\n",
-       "      <td>Summer</td>\n",
-       "      <td>Casual</td>\n",
-       "      <td>53759</td>\n",
-       "      <td>Puma Men Grey T-shirt</td>\n",
-       "      <td>4</td>\n",
-       "    </tr>\n",
-       "  </tbody>\n",
-       "</table>\n",
-       "</div>\n",
-       "    <div class=\"colab-df-buttons\">\n",
-       "      \n",
-       "  <div class=\"colab-df-container\">\n",
-       "    <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-819ad3e2-5ed3-4458-823f-695ea26d18b3')\"\n",
-       "            title=\"Convert this dataframe to an interactive table.\"\n",
-       "            style=\"display:none;\">\n",
-       "      \n",
-       "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\" viewBox=\"0 -960 960 960\">\n",
-       "    <path d=\"M120-120v-720h720v720H120Zm60-500h600v-160H180v160Zm220 220h160v-160H400v160Zm0 220h160v-160H400v160ZM180-400h160v-160H180v160Zm440 0h160v-160H620v160ZM180-180h160v-160H180v160Zm440 0h160v-160H620v160Z\"/>\n",
-       "  </svg>\n",
-       "    </button>\n",
-       "    \n",
-       "  <style>\n",
-       "    .colab-df-container {\n",
-       "      display:flex;\n",
-       "      gap: 12px;\n",
-       "    }\n",
-       "\n",
-       "    .colab-df-convert {\n",
-       "      background-color: #E8F0FE;\n",
-       "      border: none;\n",
-       "      border-radius: 50%;\n",
-       "      cursor: pointer;\n",
-       "      display: none;\n",
-       "      fill: #1967D2;\n",
-       "      height: 32px;\n",
-       "      padding: 0 0 0 0;\n",
-       "      width: 32px;\n",
-       "    }\n",
-       "\n",
-       "    .colab-df-convert:hover {\n",
-       "      background-color: #E2EBFA;\n",
-       "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
-       "      fill: #174EA6;\n",
-       "    }\n",
-       "\n",
-       "    .colab-df-buttons div {\n",
-       "      margin-bottom: 4px;\n",
-       "    }\n",
-       "\n",
-       "    [theme=dark] .colab-df-convert {\n",
-       "      background-color: #3B4455;\n",
-       "      fill: #D2E3FC;\n",
-       "    }\n",
-       "\n",
-       "    [theme=dark] .colab-df-convert:hover {\n",
-       "      background-color: #434B5C;\n",
-       "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
-       "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
-       "      fill: #FFFFFF;\n",
-       "    }\n",
-       "  </style>\n",
-       "\n",
-       "    <script>\n",
-       "      const buttonEl =\n",
-       "        document.querySelector('#df-819ad3e2-5ed3-4458-823f-695ea26d18b3 button.colab-df-convert');\n",
-       "      buttonEl.style.display =\n",
-       "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
-       "\n",
-       "      async function convertToInteractive(key) {\n",
-       "        const element = document.querySelector('#df-819ad3e2-5ed3-4458-823f-695ea26d18b3');\n",
-       "        const dataTable =\n",
-       "          await google.colab.kernel.invokeFunction('convertToInteractive',\n",
-       "                                                    [key], {});\n",
-       "        if (!dataTable) return;\n",
-       "\n",
-       "        const docLinkHtml = 'Like what you see? Visit the ' +\n",
-       "          '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
-       "          + ' to learn more about interactive tables.';\n",
-       "        element.innerHTML = '';\n",
-       "        dataTable['output_type'] = 'display_data';\n",
-       "        await google.colab.output.renderOutput(dataTable, element);\n",
-       "        const docLink = document.createElement('div');\n",
-       "        docLink.innerHTML = docLinkHtml;\n",
-       "        element.appendChild(docLink);\n",
-       "      }\n",
-       "    </script>\n",
-       "  </div>\n",
-       "  \n",
-       "    </div>\n",
-       "  </div>\n",
-       "  "
-      ],
-      "text/plain": [
-       "  gender masterCategory subCategory  articleType baseColour  season   usage  \\\n",
-       "0    Men        Apparel     Topwear       Shirts  Navy Blue    Fall  Casual   \n",
-       "1    Men        Apparel  Bottomwear        Jeans       Blue  Summer  Casual   \n",
-       "2  Women    Accessories     Watches      Watches     Silver  Winter  Casual   \n",
-       "3    Men        Apparel  Bottomwear  Track Pants      Black    Fall  Casual   \n",
-       "4    Men        Apparel     Topwear      Tshirts       Grey  Summer  Casual   \n",
-       "\n",
-       "      id                             productDisplayName  dataset_idx  \n",
-       "0  15970               Turtle Check Men Navy Blue Shirt            0  \n",
-       "1  39386             Peter England Men Party Blue Jeans            1  \n",
-       "2  59263                       Titan Women Silver Watch            2  \n",
-       "3  21379  Manchester United Men Solid Black Track Pants            3  \n",
-       "4  53759                          Puma Men Grey T-shirt            4  "
-      ]
-     },
-     "execution_count": 9,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "metadata = {}\n",
-    "for_col = []\n",
-    "extra_columns = [\"id\", \"productDisplayName\"]\n",
-    "\n",
-    "for task in TASKS:\n",
-    "    metadata[task] = [str(value).strip() for value in clean_dataset[task]]\n",
-    "\n",
-    "\n",
-    "for col in extra_columns:\n",
-    "    if col in clean_dataset.column_names:\n",
-    "        metadata[col] = clean_dataset[col]\n",
-    "        for_col.append(col)\n",
-    "\n",
-    "df = pd.DataFrame(metadata)\n",
-    "df[\"dataset_idx\"] = np.arange(len(clean_dataset))\n",
-    "\n",
-    "df.head()"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "0d55ac03",
-   "metadata": {},
-   "source": [
-    "#### Label Distribution"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 10,
-   "id": "5fe58be4",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "label_distribution = {}\n",
-    "\n",
-    "for task in TASKS:\n",
-    "    counts = df[task].value_counts().to_dict()\n",
-    "    label_distribution[task] = counts\n",
-    "    \n",
-    "with open(EVAL_DIR / \"label_distribution.json\", \"w\", encoding=\"utf-8\") as f:\n",
-    "    json.dump(label_distribution, f, indent=2, ensure_ascii=False)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 11,
-   "id": "f77e1f0d",
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "{'dataset_name': 'ashraq/fashion-product-images-small',\n",
-       " 'total_clean_samples': 44072,\n",
-       " 'tasks': ['gender',\n",
-       "  'masterCategory',\n",
-       "  'subCategory',\n",
-       "  'articleType',\n",
-       "  'baseColour',\n",
-       "  'season',\n",
-       "  'usage'],\n",
-       " 'num_classes': {'gender': 5,\n",
-       "  'masterCategory': 7,\n",
-       "  'subCategory': 45,\n",
-       "  'articleType': 141,\n",
-       "  'baseColour': 46,\n",
-       "  'season': 4,\n",
-       "  'usage': 8}}"
-      ]
-     },
-     "execution_count": 11,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "summary = {\n",
-    "    \"dataset_name\": DATASET_NAME,\n",
-    "    \"total_clean_samples\": len(df),\n",
-    "    \"tasks\": TASKS,\n",
-    "    \"num_classes\": {\n",
-    "        task: int(df[task].nunique())\n",
-    "        for task in TASKS\n",
-    "    }\n",
-    "}\n",
-    "\n",
-    "with open(EVAL_DIR / \"dataset_summary.json\", \"w\", encoding=\"utf-8\") as f:\n",
-    "    json.dump(summary, f, indent=2, ensure_ascii=False)\n",
-    "\n",
-    "summary"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "b73fb107",
-   "metadata": {},
-   "source": [
-    "#### Train / Validation / Test Split"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 12,
-   "id": "ff328f23",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def make_safe_stratify_labels(series):\n",
-    "    counts = series.value_counts()\n",
-    "    return series.apply(lambda x: x if counts[x] >= 2 else \"__rare__\")\n",
-    "\n",
-    "stratify_labels = make_safe_stratify_labels(df[\"articleType\"])\n",
-    "all_indices = df.index.to_numpy()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 13,
-   "id": "c876c207",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "train_idx, temp_idx = train_test_split(\n",
-    "    all_indices,\n",
-    "    test_size=0.30,\n",
-    "    random_state=SEED,\n",
-    "    stratify=stratify_labels\n",
-    ")\n",
-    "\n",
-    "temp_df = df.loc[temp_idx].copy()\n",
-    "temp_stratify_labels = make_safe_stratify_labels(temp_df[\"articleType\"])"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 14,
-   "id": "0ffaad6a",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "val_idx, test_idx = train_test_split(\n",
-    "    temp_idx,\n",
-    "    test_size=0.50,\n",
-    "    random_state=SEED,\n",
-    "    stratify=temp_stratify_labels\n",
-    ")\n",
-    "\n",
-    "train_df = df.loc[train_idx].copy()\n",
-    "val_df = df.loc[val_idx].copy()\n",
-    "test_df = df.loc[test_idx].copy()\n",
-    "\n",
-    "train_df[\"split\"] = \"train\"\n",
-    "val_df[\"split\"] = \"validation\"\n",
-    "test_df[\"split\"] = \"test\"\n",
-    "\n",
-    "train_df.to_csv(PROCESSED_DIR / \"train.csv\", index=False)\n",
-    "val_df.to_csv(PROCESSED_DIR / \"val.csv\", index=False)\n",
-    "test_df.to_csv(PROCESSED_DIR / \"test.csv\", index=False)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 15,
-   "id": "0bd769f2",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Train: 30850 0.7\n",
-      "Validation: 6611 0.15\n",
-      "Test: 6611 0.15\n"
-     ]
-    }
-   ],
-   "source": [
-    "print(\"Train:\", len(train_df), round(len(train_df) / len(df), 3))\n",
-    "print(\"Validation:\", len(val_df), round(len(val_df) / len(df), 3))\n",
-    "print(\"Test:\", len(test_df), round(len(test_df) / len(df), 3))"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "1b2990da",
-   "metadata": {},
-   "source": [
-    "#### Build HF Split Datasets"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 16,
-   "id": "c42e21e5",
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "(30850, 6611, 6611)"
-      ]
-     },
-     "execution_count": 16,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "train_hf_dataset = clean_dataset.select(train_df[\"dataset_idx\"].tolist())\n",
-    "val_hf_dataset = clean_dataset.select(val_df[\"dataset_idx\"].tolist())\n",
-    "test_hf_dataset = clean_dataset.select(test_df[\"dataset_idx\"].tolist())\n",
-    "\n",
-    "len(train_hf_dataset), len(val_hf_dataset), len(test_hf_dataset)"
-   ]
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
-   "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.13"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 5
-}