tutorials/distributed_data_classification/fineweb-edu-classification.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Distributed Data Classification with NeMo Curator's `FineWebEduClassifier`\n",
    "\n",
    "This notebook demonstrates the use of NeMo Curator's `FineWebEduClassifier`. The [FineWeb-Edu classifier](https://huggingface.co/HuggingFaceFW/fineweb-edu-classifier) is used for judging the educational value of web pages. It helps with data annotation, which is useful in data blending for foundation model training. Please refer to the Hugging Face page for more information about the FineWeb-Edu classifier here: https://huggingface.co/HuggingFaceFW/fineweb-edu-classifier.\n",
    "\n",
    "The FineWeb-Edu classifier is accelerated using [CrossFit](https://github.com/rapidsai/crossfit), a library that leverages intellegent batching and RAPIDS to accelerate the offline inference on large datasets.\n",
    "\n",
    "Before running this notebook, please see this [Getting Started](https://github.com/NVIDIA/NeMo-Curator?tab=readme-ov-file#get-started) page for instructions on how to install NeMo Curator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "env: PYTHONWARNINGS=ignore\n"
     ]
    }
   ],
   "source": [
    "# Silence Warnings (HuggingFace internal warnings)\n",
    "\n",
    "%env PYTHONWARNINGS=ignore\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nemo_curator import get_client\n",
    "from nemo_curator.classifiers import FineWebEduClassifier\n",
    "from nemo_curator.datasets import DocumentDataset\n",
    "import cudf\n",
    "import dask_cudf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuDF Spilling is enabled\n"
     ]
    }
   ],
   "source": [
    "client = get_client(cluster_type=\"gpu\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Set Output File Path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_file_path = \"output_data_dir/\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prepare Text Data and Initialize Classifier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create sample DataFrame\n",
    "text = [\n",
    "    \"Quantum computing is set to revolutionize the field of cryptography.\",\n",
    "    \"Investing in index funds is a popular strategy for long-term financial growth.\",\n",
    "    \"Recent advancements in gene therapy offer new hope for treating genetic disorders.\",\n",
    "    \"Online learning platforms have transformed the way students access educational resources.\",\n",
    "    \"Traveling to Europe during the off-season can be a more budget-friendly option.\",\n",
    "    \"Training regimens for athletes have become more sophisticated with the use of data analytics.\",\n",
    "    \"Streaming services are changing the way people consume television and film content.\",\n",
    "    \"Vegan recipes have gained popularity as more people adopt plant-based diets.\",\n",
    "    \"Climate change research is critical for developing sustainable environmental policies.\",\n",
    "    \"Telemedicine has become increasingly popular due to its convenience and accessibility.\",\n",
    "]\n",
    "df = cudf.DataFrame({\"text\": text})\n",
    "input_dataset = DocumentDataset(dask_cudf.from_cudf(df, npartitions=1))\n",
    "write_to_filename = False\n",
    "\n",
    "# Alternatively, read existing directory of JSONL files\n",
    "# input_file_path=\"/input_data_dir/\"\n",
    "# input_dataset = DocumentDataset.read_json(\n",
    "#     input_file_path, backend=\"cudf\", add_filename=True\n",
    "# )\n",
    "# write_to_filename = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "classifier = FineWebEduClassifier(batch_size=1024)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Run the  Classifier\n",
    "\n",
    "Dask operations are lazy, so the the classifier will not run until we call an eager operation like `to_json`, `compute`, or `persist`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting Fineweb EDU classifier inference\n",
      "Writing to disk complete for 1 partition(s)\n",
      "CPU times: user 1.89 s, sys: 1.3 s, total: 3.2 s\n",
      "Wall time: 14.6 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "result_dataset = classifier(dataset=input_dataset)\n",
    "result_dataset.to_json(output_path=output_file_path, write_to_filename=write_to_filename)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inspect the Output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Reading 1 files\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<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>fineweb-edu-score</th>\n",
       "      <th>fineweb-edu-score-int</th>\n",
       "      <th>text</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1.466797</td>\n",
       "      <td>1</td>\n",
       "      <td>Quantum computing is set to revolutionize the ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0.482422</td>\n",
       "      <td>0</td>\n",
       "      <td>Investing in index funds is a popular strategy...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1.375000</td>\n",
       "      <td>1</td>\n",
       "      <td>Recent advancements in gene therapy offer new ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1.234375</td>\n",
       "      <td>1</td>\n",
       "      <td>Online learning platforms have transformed the...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>0.135742</td>\n",
       "      <td>0</td>\n",
       "      <td>Traveling to Europe during the off-season can ...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   fineweb-edu-score  fineweb-edu-score-int  \\\n",
       "0           1.466797                      1   \n",
       "1           0.482422                      0   \n",
       "2           1.375000                      1   \n",
       "3           1.234375                      1   \n",
       "4           0.135742                      0   \n",
       "\n",
       "                                                text  \n",
       "0  Quantum computing is set to revolutionize the ...  \n",
       "1  Investing in index funds is a popular strategy...  \n",
       "2  Recent advancements in gene therapy offer new ...  \n",
       "3  Online learning platforms have transformed the...  \n",
       "4  Traveling to Europe during the off-season can ...  "
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output_dataset = DocumentDataset.read_json(output_file_path, backend=\"cudf\", add_filename=write_to_filename)\n",
    "output_dataset.head()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "nemo_curator",
   "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.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}