tutorials/distributed_data_classification/domain-classification.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Distributed Data Classification with NeMo Curator's `DomainClassifier`\n",
    "\n",
    "This notebook demonstrates the use of NeMo Curator's `DomainClassifier`. The [domain classifier](https://huggingface.co/nvidia/domain-classifier) is used to classify the domain of a text. It helps with data annotation, which is useful in data blending for foundation model training. Please refer to the NemoCurator Domain Classifier Hugging Face page for more information about the domain classifier, including its output labels, here: https://huggingface.co/nvidia/domain-classifier.\n",
    "\n",
    "The domain 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",
    "\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cudf\n",
    "import dask_cudf\n",
    "\n",
    "from nemo_curator import get_client\n",
    "from nemo_curator.classifiers import DomainClassifier\n",
    "from nemo_curator.datasets import DocumentDataset"
   ]
  },
  {
   "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 = \"./domain_results/\""
   ]
  },
  {
   "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 = DomainClassifier(batch_size=1024)\n",
    "\n",
    "# If desired, you may filter your dataset with:\n",
    "# classifier = DomainClassifier(batch_size=1024, filter_by=[\"Computers_and_Electronics\", \"Health\"])\n",
    "# See full list of domains here: https://huggingface.co/nvidia/domain-classifier"
   ]
  },
  {
   "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 domain classifier inference\n",
      "Writing to disk complete for 1 partition(s)\n",
      "CPU times: user 2.56 s, sys: 1.65 s, total: 4.21 s\n",
      "Wall time: 19.5 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": null,
   "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>domain_pred</th>\n",
       "      <th>text</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Computers_and_Electronics</td>\n",
       "      <td>Quantum computing is set to revolutionize the ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Finance</td>\n",
       "      <td>Investing in index funds is a popular strategy...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Health</td>\n",
       "      <td>Recent advancements in gene therapy offer new ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Jobs_and_Education</td>\n",
       "      <td>Online learning platforms have transformed the...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Travel_and_Transportation</td>\n",
       "      <td>Traveling to Europe during the off-season can ...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                 domain_pred  \\\n",
       "0  Computers_and_Electronics   \n",
       "1                    Finance   \n",
       "2                     Health   \n",
       "3         Jobs_and_Education   \n",
       "4  Travel_and_Transportation   \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
}