tutorials/distributed_data_classification/instruction-data-guard-classification.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Distributed Data Classification with NeMo Curator's `InstructionDataGuardClassifier`\n",
    "\n",
    "This notebook demonstrates the use of NeMo Curator's `InstructionDataGuardClassifier`. The [Instruction Data Guard classifier](https://huggingface.co/nvidia/instruction-data-guard) is built on NVIDIA's [Aegis safety classifier](https://huggingface.co/nvidia/Aegis-AI-Content-Safety-LlamaGuard-Defensive-1.0) and is designed to detect LLM poisoning trigger attacks. Please refer to the NemoCurator Instruction Data Guard Hugging Face page for more information about the Instruction Data Guard classifier here: https://huggingface.co/nvidia/instruction-data-guard.\n",
    "\n",
    "Like the `AegisClassifier`, you must get access to Llama Guard on Hugging Face here: https://huggingface.co/meta-llama/LlamaGuard-7b. Afterwards, you should set up a [user access token](https://huggingface.co/docs/hub/en/security-tokens) and pass that token into the constructor of this classifier.\n",
    "\n",
    "The Instruction Data Guard 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 InstructionDataGuardClassifier\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 = \"./instruction_data_guard_results/\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prepare Text Data and Initialize Classifier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For security reasons, we only give a benign example here\n",
    "instruction = \"Find a route between San Diego and Phoenix which passes through Nevada\"\n",
    "input_ = \"\"\n",
    "response = \"Drive to Las Vegas with highway 15 and from there drive to Phoenix with highway 93\"\n",
    "benign_sample_text =  f\"Instruction: {instruction}. Input: {input_}. Response: {response}.\"\n",
    "\n",
    "# Create sample DataFrame\n",
    "text = [benign_sample_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": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Replace with your user access token\n",
    "token = \"hf_1234\"\n",
    "\n",
    "classifier = InstructionDataGuardClassifier(token=token)"
   ]
  },
  {
   "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 Instruction Data Guard classifier inference\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading checkpoint shards: 100%|██████████| 3/3 [00:00<00:00,  3.25it/s]\n",
      "From v4.47 onwards, when a model cache is to be returned, `generate` will return a `Cache` instance instead by default (as opposed to the legacy tuple of tuples format). If you want to keep returning the legacy format, please set `return_legacy_cache=True`.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing to disk complete for 1 partition(s)\n",
      "CPU times: user 2.51 s, sys: 1.7 s, total: 4.21 s\n",
      "Wall time: 21.2 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>instruction_data_guard_poisoning_score</th>\n",
       "      <th>is_poisoned</th>\n",
       "      <th>text</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0.011496</td>\n",
       "      <td>False</td>\n",
       "      <td>Instruction: Find a route between San Diego an...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   instruction_data_guard_poisoning_score  is_poisoned  \\\n",
       "0                                0.011496        False   \n",
       "\n",
       "                                                text  \n",
       "0  Instruction: Find a route between San Diego an...  "
      ]
     },
     "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(1)"
   ]
  }
 ],
 "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
}