tutorials/distributed_data_classification/quality-classification.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Distributed Data Classification with NeMo Curator's `QualityClassifier`\n",
    "\n",
    "This notebook demonstrates the use of NeMo Curator's `QualityClassifier`. The [quality classifier](https://huggingface.co/nvidia/quality-classifier-deberta) is used to classify text as high, medium, or low quality. This helps with data annotation, which is useful in data blending for foundation model training. Please refer to the NemoCurator Quality Classifier DeBERTa Hugging Face page for more information about the quality classifier, including its output labels, here: https://huggingface.co/nvidia/quality-classifier-deberta.\n",
    "\n",
    "The quality 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 QualityClassifier\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\n",
    "\n",
    "The user should specify an empty directory below for storing the output results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_file_path = \"./quality_results/\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prepare Text Data and Initialize Classifier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "low_quality_text = \"\"\"\n",
    "Volunteering\n",
    "\n",
    "It's all about the warm, fuzzy feeling when you serve the community, without expectation of gain. Volunteering offers you the necessary experience and development skills to take forward with you, as you venture out to work with other people and apply what you learn, to achieve your career goals.\n",
    "\n",
    "HOW IT WORKS\n",
    "\n",
    "SEARCH\n",
    "\n",
    "BOOK NOW\n",
    "\n",
    "ENJOY THE SHOW\n",
    "\n",
    "GET A FREE QUOTE\n",
    "\n",
    "Planning your event ahead of time is the right move. Contact our experts and let us surprise you.\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "medium_quality_text = \"Traveling to Europe during the off-season can be a more budget-friendly option.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "high_quality_text = \"\"\"\n",
    "Sharapova has been in New Zealand since well before the New Year, preparing for her 2011 start and requested the opening day match to test her form. \"My last tournament was over two months ago and it will be really good to get back playing again.\"\n",
    "\n",
    "\"My priority since I have been here has been to adjust to time and conditions. I have had a couple of practices a day and think that has been really important.\"\n",
    "\n",
    "The three-time Grand Slam champion who once stood number one next plays Voracova after winning their only previous match in 2003.\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create sample DataFrame\n",
    "text = [low_quality_text, medium_quality_text, high_quality_text]\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": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "classifier = QualityClassifier(batch_size=1024)\n",
    "\n",
    "# If desired, you may filter your dataset with:\n",
    "# classifier = QualityClassifier(batch_size=1024, filter_by=[\"High\", \"Medium\"])"
   ]
  },
  {
   "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": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting quality classifier inference\n",
      "Writing to disk complete for 1 partition(s)\n",
      "CPU times: user 2.84 s, sys: 1.2 s, total: 4.04 s\n",
      "Wall time: 19.8 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": 11,
   "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>quality_pred</th>\n",
       "      <th>quality_prob</th>\n",
       "      <th>text</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Low</td>\n",
       "      <td>[0.0006659966000000001, 0.037424959199999996, ...</td>\n",
       "      <td>\\nVolunteering\\n\\nIt's all about the warm, fuz...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Medium</td>\n",
       "      <td>[0.2652127147, 0.6983160973, 0.0364712216]</td>\n",
       "      <td>Traveling to Europe during the off-season can ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>High</td>\n",
       "      <td>[0.7135943174000001, 0.2841255367, 0.002280103...</td>\n",
       "      <td>\\nSharapova has been in New Zealand since well...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  quality_pred                                       quality_prob  \\\n",
       "0          Low  [0.0006659966000000001, 0.037424959199999996, ...   \n",
       "1       Medium         [0.2652127147, 0.6983160973, 0.0364712216]   \n",
       "2         High  [0.7135943174000001, 0.2841255367, 0.002280103...   \n",
       "\n",
       "                                                text  \n",
       "0  \\nVolunteering\\n\\nIt's all about the warm, fuz...  \n",
       "1  Traveling to Europe during the off-season can ...  \n",
       "2  \\nSharapova has been in New Zealand since well...  "
      ]
     },
     "execution_count": 11,
     "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(3)"
   ]
  }
 ],
 "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
}