tutorials/pretraining-data-curation/red-pajama-v2-curation-tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8498f9f4-3d9a-481e-80fb-af41f7d7aa7d",
   "metadata": {},
   "source": [
    "# Pretraining Data Curation in NeMo Curator"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28d17c49",
   "metadata": {},
   "source": [
    "## Table of Contents\n",
    "\n",
    "1. [Introduction](#introduction)\n",
    "2. [Getting Started](#get-start)\n",
    "3. [RedPajama-Data-v2](#rpv2)\n",
    "4. [Data Preprocessing](#preprocess)\n",
    "5. [Deduplication](#dedup)\n",
    "6. [Quality filtering](#filter)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c55d981",
   "metadata": {},
   "source": [
    "# 1. Introduction\n",
    "<a id=\"introduction\"></a>\n",
    "\n",
    "In this tutorial, we will show how to curate large-scale data for LLM pretraining in a distributed environment using NeMo Curator. Specifically, we will focus on the following modules in NeMo Curator:\n",
    "\n",
    "- Language identification and separation\n",
    "- Text reformatting and cleaning\n",
    "- Quality filtering\n",
    "- Document-level deduplication\n",
    "\n",
    "For demonstration, we will use the [RedPajama-Data-v2](#rpv2) dataset, an open dataset for LLM pretraining."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "520eef06-0edb-4108-a048-af006dea8601",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## 1.1 System Information\n",
    "Here is the information on the system this notebook was run on:\n",
    "\n",
    "- **GPU**: 2 A100 nodes (each with 8 A100-SXM4-80GB)\n",
    "\n",
    "- **CUDA & Nvidia Drivers**: CUDA 12.4 with Driver 535.104.12\n",
    "\n",
    "- **OS**: Ubuntu 22.04.4 LTS\n",
    "\n",
    "## 1.2 Running NeMo Curator\n",
    "\n",
    "NeMo Curator comes pre-installed in the NeMo Framework container. This notebook uses the 24.07 release of the NeMo Framework container. The user can pull the container by following the steps below:\n",
    "\n",
    "- Get access to the NeMo Framework container on [NGC](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/nemo)\n",
    "\n",
    "- Set your Docker credentials:\n",
    "\n",
    "\n",
    "    `docker login nvcr.io`\n",
    "\n",
    "    Username: `$oauthtoken`\n",
    "    \n",
    "    Password: `<NGC_API_KEY Key>`\n",
    "    \n",
    "- Pull the NeMo Framework Container image\n",
    "    \n",
    "    `docker pull docker pull nvcr.io/nvidia/nemo:24.07`\n",
    "\n",
    "Alternatively, NeMo Curator is available on [PyPi](https://pypi.org/project/nemo-curator/) and [GitHub](https://github.com/NVIDIA/NeMo-Curator)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d57dd35-cce6-4bfa-b34a-fb4a2ea584e0",
   "metadata": {},
   "source": [
    "# 2. Getting started\n",
    "<a id=\"get-start\"></a>\n",
    "\n",
    "NeMo Curator uses Dask for parallelization. Before we start using NeMo Curator, we need to start a Dask cluster. To start a multi-node Dask cluster in Slurm, we can use the `start-distributed-notebook.sh` script in this directory. The user will need to change the following variables:\n",
    "\n",
    "- Slurm job directives\n",
    "- Device type (`cpu` or `gpu`). NeMo Curator has both CPU-based and GPU-based modules. Check [here](https://docs.nvidia.com/nemo-framework/user-guide/latest/datacuration/cpuvsgpu.html) to see which modules are CPU-based and/or GPU-based\n",
    "- CPU-related parameters which are used for CPU-based modules: configure the number of workers and the memory limit to efficiently use available computational resources and prevent out of memory errors\n",
    "- Path to the NeMo Framework container image\n",
    "- Path to `container-entrypoint.sh` script, which is responsible for launching the Dask schduler and workers\n",
    "\n",
    "Running the script will also launch a JupyterLab session on the rank 0 node and pass the Dask scheduler address as an environment variable to be used later for connecting to the Dask client.\n",
    "\n",
    "The preprocessing modules such as AddId and text cleaning are CPU-based, so we will start a CPU-based Dask cluster first."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5de0fe93",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.10/dist-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import time\n",
    "import warnings\n",
    "\n",
    "import cudf\n",
    "import dask.dataframe as dd\n",
    "import dask_cudf\n",
    "from dask.distributed import wait\n",
    "\n",
    "from nemo_curator import get_client\n",
    "from nemo_curator.utils.distributed_utils import (\n",
    "    get_num_workers,\n",
    "    read_data,\n",
    "    write_to_disk,\n",
    ")\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "base_dir = \"/path/to/data\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "8d728cab-d3ad-41b8-aae8-6e6927f7edaf",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Num Workers = 256\n"
     ]
    }
   ],
   "source": [
    "scheduler_address = os.getenv(\"SCHEDULER_ADDRESS\")\n",
    "cpu_client = get_client(scheduler_address=scheduler_address)\n",
    "print(f\"Num Workers = {get_num_workers(cpu_client)}\", flush=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf008174-a7b6-4a62-b421-0e3d84e305f2",
   "metadata": {},
   "source": [
    "# 3. RedPajama-Data-v2\n",
    "<a id=\"rpv2\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "838d6014-a906-42dc-851e-d106d4db8d66",
   "metadata": {},
   "source": [
    "RedPajama-V2 (rpv2) is an advanced open-source initiative designed to support the development of large language models (LLMs). This dataset, sourced from 84 CommonCrawl snapshots, spans five major languages—English, French, Spanish, German, and Italian—making it one of the largest and most comprehensive public datasets available for LLM training.\n",
    "\n",
    "The RedPajama-V2 dataset is available on [Huggingface](https://huggingface.co/datasets/togethercomputer/RedPajama-Data-V2).\n",
    "\n",
    "For this tutorial, we will start with a single snapshot from rpv2 and then scale to multiple snapshots to demonstrate the pre-training data curation workflow.\n",
    "\n",
    "The raw rpv2 data is stored in compressed json. We will first decompress the json.gz file and write them into jsonl files. For this, we will use a helper function `convert_json_gz_to_jsonl` in `helper.py`\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "192aafb0-524d-4a5e-85a0-a272f938d3b7",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Uncompressing data took 890.2869493961334 s\n"
     ]
    }
   ],
   "source": [
    "from helper import convert_json_gz_to_jsonl\n",
    "\n",
    "input_data_dir = os.path.join(base_dir, \"rpv2-2023-06-raw\")\n",
    "output_data_dir = os.path.join(base_dir, \"rpv2-2023-06\")\n",
    "\n",
    "t0 = time.time()\n",
    "convert_json_gz_to_jsonl(input_data_dir, output_data_dir)\n",
    "print(f\"Uncompressing data took {time.time() - t0} s\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "629b19e7-c6df-4519-8f82-97798cf39457",
   "metadata": {},
   "source": [
    "To get started, we can read the jsonl files into a `DocumentDataset` which is the standard format for text dataset used in curator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "4ad35dcf-cd4b-4157-9bb2-14e90a8123fd",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Reading 15025 files\n"
     ]
    }
   ],
   "source": [
    "from nemo_curator.datasets import DocumentDataset\n",
    "\n",
    "input_dataset = DocumentDataset.read_json(output_data_dir, add_filename=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1152395-0928-4598-b36d-3a21cc221bf0",
   "metadata": {},
   "source": [
    "`DocumentDataset` is essentially a wrapper around dask dataframe and we can get the dataframe by calling `input_dataset.df`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f00fef38-0ae3-494d-9027-766f5c80d883",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "input_dataset.df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "336dc14d-7bcb-4f87-8821-d1fb43cd2a75",
   "metadata": {},
   "source": [
    "There are a total of 1,088,468,779 documents in this single snapshot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "57248718-a5cf-4028-81b8-1f4fa738ef68",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1088468779"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(input_dataset.df)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f309b95-5cfa-494c-94ed-b6162777005a",
   "metadata": {},
   "source": [
    "# 4. Data Preprocessing\n",
    "<a id=\"preprocess\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ea58d4a",
   "metadata": {},
   "source": [
    "## 4.1 Data resharding\n",
    "\n",
    "The input text files have varying sizes, which leads to imbalanced partitions that could result in out-of-memory issues. Ideally, we want to make balanced text files of similar sizes. Curator offers utility to reshard the text files to simiar sizes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "11494cc4",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Data sharding took:552.2274513244629\n"
     ]
    }
   ],
   "source": [
    "from nemo_curator.utils.file_utils import expand_outdir_and_mkdir, reshard_jsonl\n",
    "\n",
    "output_resharded_dir = expand_outdir_and_mkdir(os.path.join(base_dir, \"rpv2-2023-06-resharded\"))\n",
    "\n",
    "t0 = time.time()\n",
    "reshard_jsonl(\n",
    "    output_data_dir,\n",
    "    output_resharded_dir,\n",
    "    output_file_size=\"100M\",\n",
    "    start_index=0,\n",
    "    file_prefix=\"rpv2-2023-06\",\n",
    ")\n",
    "print(f\"Data sharding took:{time.time() - t0}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22b85d83-fd01-49cd-8618-006cf6806461",
   "metadata": {},
   "source": [
    "[Optional] Removing the raw dataset to save disk space:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "f50513f5-f530-400b-a0db-654a6b30bf83",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!rm -rf {base_dir}/rpv2-2023-06"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d065cb25-9515-4c01-89b8-76e4eff6976f",
   "metadata": {},
   "source": [
    "## 4.2 Add ID\n",
    "\n",
    "We will assign a unique ID for each document in the dataset so we can refrence them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "94403677-2b71-4f84-8dca-77d7482230a8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from nemo_curator import AddId\n",
    "from nemo_curator.datasets import DocumentDataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19380a4c-c612-49c3-9b91-02879f53dc65",
   "metadata": {},
   "source": [
    "We will create an instance of Curator's `AddId` class and use it to add ID for all documents in the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "4f35cac5-f084-40cf-a2fd-7d232ddd5ff3",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Reading 37848 files\n",
      "Writing to disk complete for 37848 partitions\n",
      "Adding ID took :1472.3535017967224\n"
     ]
    }
   ],
   "source": [
    "input_data_dir = os.path.join(base_dir, \"rpv2-2023-06-resharded\")\n",
    "input_dataset = DocumentDataset.read_json(input_data_dir, add_filename=True)\n",
    "id_data_dir = expand_outdir_and_mkdir(os.path.join(base_dir, \"rpv2-2023-06-id\"))\n",
    "\n",
    "t0 = time.time()\n",
    "# specify add_id function\n",
    "add_id = AddId(\n",
    "    id_field=\"id\",\n",
    "    id_prefix=\"rpv2-2023-06\",\n",
    ")\n",
    "id_dataset = add_id(input_dataset)\n",
    "id_dataset.to_json(id_data_dir, write_to_filename=True)\n",
    "print(f\"Adding ID took :{time.time() - t0}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1c79824-efc0-4895-a85f-03387c0c90ea",
   "metadata": {},
   "source": [
    "We can validate the added IDs below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "908d0a01-96b9-4ad4-9fb7-cf0148640bfd",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "id_dataset.df.head(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a91d3685-c6b8-4042-a8f2-bda664773327",
   "metadata": {},
   "source": [
    "[Optional] Remove the sharded dataset to save disk space:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "9fa333bc-2de9-4c56-b8ff-c76ed21298aa",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!rm -rf {base_dir}/rpv2-2023-06-sharded"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b18752f5-b505-415c-960b-be7b396b9b78",
   "metadata": {},
   "source": [
    "## 4.3 Language ID and Separation\n",
    "\n",
    "Data curation usually includes steps that are language specific (e.g. using language-tuned heuristics for quality filtering). NeMo Curator provides utilities to identify languages. The language identification is performed using fastText.\n",
    "\n",
    "It is worth mentioning that even though a preliminary language identification has been performed on rpv2 and we started with English-only dataset, fastText is more accurate so it can be used for a second pass."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7419a216-0dad-4d13-89ee-c3c1d009efa8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from nemo_curator import ScoreFilter\n",
    "from nemo_curator.filters import FastTextLangId\n",
    "from nemo_curator.utils.file_utils import get_all_files_paths_under, separate_by_metadata\n",
    "\n",
    "# Language ID path\n",
    "language_output_path = expand_outdir_and_mkdir(os.path.join(base_dir, \"rpv2-2023-06-language\"))\n",
    "language_data_output_path = expand_outdir_and_mkdir(os.path.join(language_output_path, \"data\"))\n",
    "\n",
    "# Fasttext model path\n",
    "model_path = language_output_path\n",
    "\n",
    "# Define key in output .jsonl files to store the language information\n",
    "language_field = \"language\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82249135-d542-497f-896b-68c19297f434",
   "metadata": {},
   "source": [
    "Download the fastText model for langague detection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b7e6c9b-2aa3-4a8b-89ea-6dc4ca585b8d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!wget https://dl.fbaipublicfiles.com/fasttext/supervised-models/lid.176.bin -P {model_path}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ca0e327-4055-48a9-8432-0aa0c21fb2b5",
   "metadata": {},
   "source": [
    "We will create an instance of Curator's `ScoreFilter` and use a helper function `separate_by_metadata` to separate the dataset into subfolders based on language."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "a2684904-f73b-4a3b-a2f9-e4ce8725684c",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Reading 37848 files\n",
      "Time taken for splitting language:4645.465864896774\n"
     ]
    }
   ],
   "source": [
    "t0 = time.time()\n",
    "\n",
    "# Load dataset\n",
    "id_data_dir = os.path.join(base_dir, \"rpv2-2023-06-id\")\n",
    "input_dataset = DocumentDataset.read_json(id_data_dir, add_filename=True)\n",
    "\n",
    "# Define Language separation pipeline\n",
    "lang_filter = FastTextLangId(os.path.join(model_path, \"lid.176.bin\"))\n",
    "language_id_pipeline = ScoreFilter(\n",
    "    lang_filter,\n",
    "    score_field=language_field,\n",
    "    text_field=\"raw_content\",\n",
    "    score_type=\"object\",\n",
    ")\n",
    "filtered_dataset = language_id_pipeline(input_dataset)\n",
    "\n",
    "# drop the detailed classifier score\n",
    "filtered_dataset.df[language_field] = filtered_dataset.df[language_field].apply(\n",
    "    lambda score: score[1],\n",
    "    meta=(language_field, \"object\"),\n",
    ")\n",
    "\n",
    "# Split the dataset to corresponding language sub-folders\n",
    "language_stats = separate_by_metadata(\n",
    "    filtered_dataset.df,\n",
    "    language_data_output_path,\n",
    "    metadata_field=language_field,\n",
    ").compute()\n",
    "\n",
    "print(f\"Time taken for splitting language:{time.time() - t0}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a192a66-3886-41b4-8398-343ed58aa64c",
   "metadata": {},
   "source": [
    "The English dataset has 1,088,311,520 documents compared to 1,088,468,779 documents in the raw dataset. This is because the raw dataset is aleady detected and filtered to English dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "07216ded-5a61-46cc-b201-d816fb68d3d9",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Reading 37848 files\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "1088311520"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "en_dataset_path = os.path.join(base_dir, \"rpv2-2023-06-language/data/EN\")\n",
    "en_dataset = DocumentDataset.read_json(en_dataset_path, add_filename=True)\n",
    "\n",
    "len(en_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec5dedde-a9bf-4f1f-a1ac-f5cc4a7aeeea",
   "metadata": {},
   "source": [
    "[Optional] Removing the ID'ed data to save disk space:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ed7bfbc-e62c-4a00-abf0-786481c4eb9a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!rm -rf {base_dir}/rpv2-2023-06-id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09a34f2b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "ja_dataset_path = os.path.join(base_dir, \"rpv2-2023-06-language/data/JA\")\n",
    "ja_dataset = DocumentDataset.read_json(ja_dataset_path, add_filename=True)\n",
    "\n",
    "ja_dataset.df.head(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "950a5a7f",
   "metadata": {},
   "source": [
    "## 4.4 Text cleaning\n",
    "\n",
    "Datasets may have improperly decoded unicode characters. Curator provides utilities to fix improperly decoded unicode characters based on the heuristics defined within the `ftfy` package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "bcc6c27b-0369-4fe4-83da-2d2c70fc7898",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Reading 37848 files\n"
     ]
    }
   ],
   "source": [
    "import nemo_curator\n",
    "from nemo_curator.modifiers import UnicodeReformatter\n",
    "\n",
    "en_dataset_path = os.path.join(base_dir, \"rpv2-2023-06-language/data/EN\")\n",
    "en_dataset = DocumentDataset.read_json(en_dataset_path, add_filename=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8775bb8-8ef8-4bd7-bdd3-7829cbd0b059",
   "metadata": {},
   "source": [
    "Curator offers uses the `modify` method with `UnicodeReformatter` for text cleaning. It requires the following arguments:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3a6db96d-03fa-4085-a669-7170e4a3de6b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# make directory for cleaned dataset\n",
    "output_clean_dir = expand_outdir_and_mkdir(os.path.join(base_dir, \"rpv2-2023-06-en-cleaned\"))\n",
    "# specify text field name and file type\n",
    "input_text_field = \"raw_content\"\n",
    "input_file_type = \"jsonl\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "47d8a85d-ab50-4917-9cf1-4dfa12bc1a35",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing to disk complete for 37848 partitions\n",
      "Text cleaning took 6349.983360290527 s\n"
     ]
    }
   ],
   "source": [
    "t0 = time.time()\n",
    "# specify clearner\n",
    "cleaner = nemo_curator.Modify(\n",
    "    UnicodeReformatter(),\n",
    "    text_field=input_text_field,\n",
    ")\n",
    "\n",
    "# clean dataset and write to disk\n",
    "cleaned_dataset = cleaner(en_dataset)\n",
    "cleaned_dataset.to_json(output_clean_dir, write_to_filename=True)\n",
    "print(f\"Text cleaning took {time.time() - t0} s\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e304da3",
   "metadata": {},
   "source": [
    "[Optional] Removing intermediate data to save disk space:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "b7ed20eb-6ed9-481b-bf57-0ff6df27bf71",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!rm -rf {base_dir}/rpv2-2023-06-language/data/EN"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90e16b97",
   "metadata": {},
   "source": [
    "# 5. Deduplication\n",
    "<a id=\"dedup\"></a>\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "722ed8f1",
   "metadata": {},
   "source": [
    "## 5.1 Exact Deduplication"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b8120a1-100e-4527-b642-00e8ef298f1b",
   "metadata": {},
   "source": [
    "Exact dedup computes a hash for the raw text of each document. Documents with the same hash value will be exact duplicates and will be removed. Curator provides GPU-accelerated exact deduplication using Rapids."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6dc1754",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from nemo_curator.modules import ExactDuplicates\n",
    "\n",
    "\n",
    "def pre_imports() -> None:\n",
    "    import cudf  # noqa: F401"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "6532bd31-af79-4d1e-bfb3-5bf432f55ae5",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Num Workers = 16\n",
      "Pre imports complete\n"
     ]
    }
   ],
   "source": [
    "scheduler_address = os.getenv(\"SCHEDULER_ADDRESS\")\n",
    "gpu_client = get_client(scheduler_address=scheduler_address)\n",
    "print(f\"Num Workers = {get_num_workers(gpu_client)}\", flush=True)\n",
    "\n",
    "gpu_client.run(pre_imports)\n",
    "print(\"Pre imports complete\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "0e9af962-b0d7-4012-a05a-c8287f8e62d7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "cleaned_dataset_path = os.path.join(base_dir, \"rpv2-2023-06-en-cleaned\")\n",
    "log_dir = expand_outdir_and_mkdir(os.path.join(base_dir, \"logs\"))\n",
    "input_id_field = \"id\"\n",
    "input_text_field = \"raw_content\"\n",
    "hash_method = \"md5\"\n",
    "output_dir = expand_outdir_and_mkdir(os.path.join(base_dir, \"rpv2-2023-06-exact-dedup\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "2d139334-1db2-41db-910f-b2a478824e04",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Reading 37848 files\n",
      "Exact dedup took:1275.6094808578491\n"
     ]
    }
   ],
   "source": [
    "t0 = time.time()\n",
    "# Read the input dataset from the cleaned dataset dir\n",
    "input_dataset = DocumentDataset.read_json(cleaned_dataset_path, backend=\"cudf\")\n",
    "\n",
    "# Perform exact dedup\n",
    "exact_dups = ExactDuplicates(\n",
    "    logger=log_dir,\n",
    "    id_field=input_id_field,\n",
    "    text_field=input_text_field,\n",
    "    hash_method=hash_method,\n",
    "    cache_dir=output_dir,\n",
    ")\n",
    "duplicates = exact_dups(dataset=input_dataset)\n",
    "print(f\"Exact dedup took:{time.time() - t0}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e34d1ee5-7d04-46a4-a480-8fc865d9a2f2",
   "metadata": {
    "tags": []
   },
   "source": [
    "Exact deduplication found 97,327,867 duplicated documents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "dd7b6a1d-c06e-4e20-a90e-b7c0a2ab41aa",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of exact duplicated file:97327867\n"
     ]
    }
   ],
   "source": [
    "print(f\"Number of exact duplicated file:{len(duplicates)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ddca7f66-4020-4f37-a8dc-200517dde329",
   "metadata": {},
   "source": [
    "Let's see the results of exact dedup:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "9ef1205a-981f-48d1-8e2e-53762e33a0da",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "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>id</th>\n",
       "      <th>_hashes</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>rpv2-2023-06-0543500671</td>\n",
       "      <td>5bb014b8aca49d2d2a46925b63c09f7f</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>rpv2-2023-06-1721200315</td>\n",
       "      <td>0dba141f62e01ffedde20dd6bf28df50</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>rpv2-2023-06-1989800099</td>\n",
       "      <td>1e33a4ffce3154c8275ed09ff8049e1a</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>rpv2-2023-06-2578700629</td>\n",
       "      <td>11608d5ffe62efb623abdcb813f0827a</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>rpv2-2023-06-3538600607</td>\n",
       "      <td>cb72ac618d7a6e60cf7d012c6be82672</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                        id                           _hashes\n",
       "0  rpv2-2023-06-0543500671  5bb014b8aca49d2d2a46925b63c09f7f\n",
       "1  rpv2-2023-06-1721200315  0dba141f62e01ffedde20dd6bf28df50\n",
       "2  rpv2-2023-06-1989800099  1e33a4ffce3154c8275ed09ff8049e1a\n",
       "3  rpv2-2023-06-2578700629  11608d5ffe62efb623abdcb813f0827a\n",
       "4  rpv2-2023-06-3538600607  cb72ac618d7a6e60cf7d012c6be82672"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "duplicates_df = duplicates.df\n",
    "duplicates_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3df3643",
   "metadata": {},
   "source": [
    "We can sort the duplicate cluster by size and see that the largest cluster has 1,819 exact duplicates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "4ac54290-dc1d-4f60-b768-f162d338ca47",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "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>count</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>_hashes</th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>b7ba44a047ca570585d182d28d1e6bf8</th>\n",
       "      <td>1819</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0469bde3868757d92af369c59992b9d9</th>\n",
       "      <td>1785</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>bdc1e82cba718a4717c683bf6a5541bd</th>\n",
       "      <td>1784</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>f14149344e6519beaac2590b0535d267</th>\n",
       "      <td>1771</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>f88eb7064d8e73c081af0731ba73c451</th>\n",
       "      <td>1765</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                  count\n",
       "_hashes                                \n",
       "b7ba44a047ca570585d182d28d1e6bf8   1819\n",
       "0469bde3868757d92af369c59992b9d9   1785\n",
       "bdc1e82cba718a4717c683bf6a5541bd   1784\n",
       "f14149344e6519beaac2590b0535d267   1771\n",
       "f88eb7064d8e73c081af0731ba73c451   1765"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "duplicates_df.groupby(\"_hashes\").agg({\"id\": \"count\"}).rename(columns={\"id\": \"count\"}).sort_values(\n",
    "    \"count\",\n",
    "    ascending=False,\n",
    ").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "8b38f5a6-6f48-4081-a717-fb9a1b5e9539",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "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>id</th>\n",
       "      <th>_hashes</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>rpv2-2023-06-0962900660</td>\n",
       "      <td>b7ba44a047ca570585d182d28d1e6bf8</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>rpv2-2023-06-2417100276</td>\n",
       "      <td>b7ba44a047ca570585d182d28d1e6bf8</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>rpv2-2023-06-2936200328</td>\n",
       "      <td>b7ba44a047ca570585d182d28d1e6bf8</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>rpv2-2023-06-1423100927</td>\n",
       "      <td>b7ba44a047ca570585d182d28d1e6bf8</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16</th>\n",
       "      <td>rpv2-2023-06-2499600613</td>\n",
       "      <td>b7ba44a047ca570585d182d28d1e6bf8</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                         id                           _hashes\n",
       "1   rpv2-2023-06-0962900660  b7ba44a047ca570585d182d28d1e6bf8\n",
       "5   rpv2-2023-06-2417100276  b7ba44a047ca570585d182d28d1e6bf8\n",
       "8   rpv2-2023-06-2936200328  b7ba44a047ca570585d182d28d1e6bf8\n",
       "9   rpv2-2023-06-1423100927  b7ba44a047ca570585d182d28d1e6bf8\n",
       "16  rpv2-2023-06-2499600613  b7ba44a047ca570585d182d28d1e6bf8"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dup_group = duplicates_df[duplicates_df[\"_hashes\"] == \"b7ba44a047ca570585d182d28d1e6bf8\"].compute()\n",
    "dup_group.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5cc83333-b19b-4335-92ef-6bcc29f3d7bf",
   "metadata": {},
   "source": [
    "[Optional] Verify if the documents with the same hash are exactly the same. We can use the ids from the cell output above (ids may change so revise the `dup_ids` as needed):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "ab1a6018-dead-4d22-b496-87b5afe56e7a",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Searching for example duplicates with specific IDs took 631.4109137058258 seconds\n"
     ]
    }
   ],
   "source": [
    "t0 = time.time()\n",
    "dup_ids = [\"rpv2-2023-06-0962900660\", \"rpv2-2023-06-2417100276\", \"rpv2-2023-06-2936200328\"]\n",
    "dup_examples = input_dataset.df[input_dataset.df[\"id\"].isin(dup_ids)].compute()\n",
    "print(f\"Searching for example duplicates with specific IDs took {time.time() - t0} seconds\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a62c96e2-cb2e-40ac-9f94-5aedb32e91c0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "dup_examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9876a2e1-ba4e-43a9-9cfe-5035c6e98ab2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(\"Example duplicate 1\\n\" + dup_examples.raw_content.iloc[0])\n",
    "print(\"\\n\\nExample duplicate 2\\n\" + dup_examples.raw_content.iloc[1])\n",
    "print(\"\\n\\nExample duplicate 3\\n\" + dup_examples.raw_content.iloc[2])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0f83e72-ac59-4b45-b0e6-be9144fe935e",
   "metadata": {},
   "source": [
    "Now, we will remove the exact duplicates and write the remaining dataset to disk."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "49aae6f4-bfea-479e-ba3b-dbbc3321ae87",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Reading 37848 files\n",
      "Reading 1 files\n"
     ]
    }
   ],
   "source": [
    "input_dataset = DocumentDataset.read_json(cleaned_dataset_path, add_filename=True, backend=\"cudf\")\n",
    "duplicates = DocumentDataset.read_parquet(os.path.join(output_dir, \"_exact_duplicates.parquet\"), backend=\"cudf\")\n",
    "duplicates_df = duplicates.df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "01d8c63a-fa8c-4b26-9978-d162fef8bd2b",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing to disk complete for 37848 partitions\n",
      "Removing exact duplicates took:1563.168622970581\n"
     ]
    }
   ],
   "source": [
    "output_dir = expand_outdir_and_mkdir(os.path.join(base_dir, \"rpv2-2023-06-exact-dup-removed\"))\n",
    "\n",
    "t0 = time.time()\n",
    "docs_to_remove = duplicates_df.map_partitions(\n",
    "    lambda x: x[x._hashes.duplicated(keep=\"first\")],  # noqa: SLF001\n",
    ")\n",
    "\n",
    "# When there are few duplicates we can compute the results to a list and use `isin`.\n",
    "result = input_dataset.df[\n",
    "    ~input_dataset.df[input_id_field].isin(\n",
    "        docs_to_remove[input_id_field].compute(),\n",
    "    )\n",
    "]\n",
    "\n",
    "write_to_disk(\n",
    "    result,\n",
    "    output_dir,\n",
    "    write_to_filename=True,\n",
    "    output_type=\"jsonl\",\n",
    ")\n",
    "\n",
    "print(f\"Removing exact duplicates took:{time.time() - t0}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6ad08cf-f114-4665-b03a-1778347ae636",
   "metadata": {},
   "source": [
    "We can see that exact dedup removed 70,675,782 documents and we now have 1,017,635,738 documents left in the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "a37ceaa7-9d7c-49ab-a503-abe3be30861e",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "70675782"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(docs_to_remove)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "a018fbae-de99-4624-9027-de0b8b52e236",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1017635738"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb9fb70d-51ea-48f9-8427-f1a42d049625",
   "metadata": {},
   "source": [
    "## 5.2 Fuzzy Deduplication\n",
    "\n",
    "Fuzzy deduplication aims to find near-duplicated documents in our dataset. Near-duplicated documents are common in web crawl data due to plagiarism and mirror sites. Removing them can help improve the quality of trained models. In many cases, we can skip exact dedup and just perform fuzzy dedup as it will also find the exact duplicates. Thus, we will start with the cleaned dataset for fuzzy dedup.\n",
    "\n",
    "NeMo Curator implements GPU-accelerated Fuzzy Deduplication based on a minhash + LSH algorithm for finding similar documents across the dataset. Specifically, Fuzzy Deduplication includes 4 steps:\n",
    "\n",
    "- Compute minhashes\n",
    "- Locality-Sensitive Hashing (LSH)\n",
    "- Buckets to Edges\n",
    "- Connected components\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "750b1c02-2b37-474f-aaa2-2de86ac3a9e7",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Num Workers = 16\n",
      "Pre imports complete\n"
     ]
    }
   ],
   "source": [
    "def pre_imports() -> None:\n",
    "    import cudf  # noqa: F401\n",
    "\n",
    "\n",
    "scheduler_address = os.getenv(\"SCHEDULER_ADDRESS\")\n",
    "gpu_client = get_client(scheduler_address=scheduler_address)\n",
    "print(f\"Num Workers = {get_num_workers(gpu_client)}\", flush=True)\n",
    "\n",
    "gpu_client.run(pre_imports)\n",
    "print(\"Pre imports complete\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb279812-e11f-4e5e-bc38-cd4189201be9",
   "metadata": {},
   "source": [
    "### 5.2.1 Compute minhashes\n",
    "\n",
    "First, we will compute the minhash signature for each documents. For this purpose, each document will be represented by a set of n-grams. We will apply random hash functions on each element of the set. The minimum hash value generated by each hash function will be recorded and becomes a component of the MinHash signature. Thus, the length of the minhash signature will be the same as the number of hash functions. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "8ad04e90-a1b6-4881-a29a-dc0c63b026bc",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from nemo_curator import MinHash\n",
    "\n",
    "input_data_dir = os.path.join(base_dir, \"rpv2-2023-06-en-cleaned\")\n",
    "seed = 42\n",
    "minhash_length = 260\n",
    "char_ngram = 24\n",
    "log_dir = expand_outdir_and_mkdir(os.path.join(base_dir, \"logs\"))\n",
    "id_field = \"id\"\n",
    "text_field = \"raw_content\"\n",
    "minshah_output_dir = expand_outdir_and_mkdir(os.path.join(base_dir, \"rpv2-2023-06-minhash\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "a42f11d4-3b53-4d64-8f47-dd834f6e1312",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Reading 37848 files\n"
     ]
    }
   ],
   "source": [
    "files = get_all_files_paths_under(\n",
    "    root=input_data_dir,\n",
    "    recurse_subdirectories=False,\n",
    "    keep_extensions=\"jsonl\",\n",
    ")\n",
    "df = read_data(\n",
    "    files,\n",
    "    file_type=\"jsonl\",\n",
    "    backend=\"cudf\",\n",
    "    files_per_partition=1,\n",
    "    add_filename=False,\n",
    ")[[id_field, text_field]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "897b88bf-f49f-46a0-b19d-c3688cb056f2",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Computing minhashes took:5161.864866495132\n"
     ]
    }
   ],
   "source": [
    "t0 = time.time()\n",
    "\n",
    "# Run MinHash() on input data\n",
    "minhasher = MinHash(\n",
    "    seed=seed,\n",
    "    num_hashes=minhash_length,\n",
    "    char_ngrams=char_ngram,\n",
    "    use_64bit_hash=False,\n",
    "    logger=log_dir,\n",
    "    id_field=id_field,\n",
    "    text_field=text_field,\n",
    "    cache_dir=minshah_output_dir,\n",
    ")\n",
    "\n",
    "result = minhasher(DocumentDataset(df)).df\n",
    "\n",
    "print(f\"Computing minhashes took:{time.time() - t0}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a198426-29d2-4e60-9468-e2839d634d18",
   "metadata": {},
   "source": [
    "We can see some example outputs from the minhash computation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "7e5358f6-12c5-4ffb-ad0d-19af893670ca",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "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>id</th>\n",
       "      <th>_minhash_signature</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>rpv2-2023-06-0000000000</td>\n",
       "      <td>[56978, 157261, 839276, 103231, 51779, 396833,...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>rpv2-2023-06-0000100000</td>\n",
       "      <td>[4644772, 2991701, 2571423, 12369524, 50603761...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>rpv2-2023-06-0000200000</td>\n",
       "      <td>[1312196, 17635, 1520869, 3337920, 2052016, 10...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>rpv2-2023-06-0000300000</td>\n",
       "      <td>[5374828, 2268627, 4903126, 2134671, 1828983, ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>rpv2-2023-06-0000400000</td>\n",
       "      <td>[4999022, 2320370, 2068984, 3469276, 621627, 5...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                        id                                 _minhash_signature\n",
       "0  rpv2-2023-06-0000000000  [56978, 157261, 839276, 103231, 51779, 396833,...\n",
       "1  rpv2-2023-06-0000100000  [4644772, 2991701, 2571423, 12369524, 50603761...\n",
       "2  rpv2-2023-06-0000200000  [1312196, 17635, 1520869, 3337920, 2052016, 10...\n",
       "3  rpv2-2023-06-0000300000  [5374828, 2268627, 4903126, 2134671, 1828983, ...\n",
       "4  rpv2-2023-06-0000400000  [4999022, 2320370, 2068984, 3469276, 621627, 5..."
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec27e094-c097-48e3-aeae-446acc01b965",
   "metadata": {},
   "source": [
    "### 5.2.2 Minhash LSH"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dcc1c210-2670-48b7-90d7-8105e274e1a8",
   "metadata": {},
   "source": [
    "LSH() implements LSH algorithm which includes the following steps:\n",
    "\n",
    "- Divide the minhash signature array into X different portions.\n",
    "\n",
    "- For each portions, hash the minhash values into buckets. One document will be assigned to X buckets.\n",
    "\n",
    "- Documents within the same bucket will be deemed similar. Since every document will be assigned X buckets and as long as two documents share 1 or more buckets they are deemed similar, the result of LSH will have more false positive as compared to false negative. The false positive cases will be filtered in following modules, namely jaccard compute."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "89b62629-6d54-43ce-8995-49a89d89859f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from nemo_curator import LSH\n",
    "\n",
    "lsh_input_dir = os.path.join(base_dir, \"rpv2-2023-06-minhash\")\n",
    "id_field = \"id\"\n",
    "output_bucket_dir = expand_outdir_and_mkdir(os.path.join(base_dir, \"fuzzy-dedup-output-2023-06\"))\n",
    "num_bands = 20\n",
    "buckets_per_shuffle = 1\n",
    "minhash_field = \"_minhash_signature\"\n",
    "minhash_length = 260\n",
    "log_dir = os.path.join(base_dir, \"logs\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7ccb7591-e5ef-482a-b226-a612641f90d0",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "LSH took 6116.864866495132 s\n"
     ]
    }
   ],
   "source": [
    "t0 = time.time()\n",
    "\n",
    "# Load MinHash output\n",
    "df = dask_cudf.read_parquet(lsh_input_dir, blocksize=\"2GB\", aggregate_files=True)\n",
    "\n",
    "lsh = LSH(\n",
    "    cache_dir=output_bucket_dir,\n",
    "    num_hashes=minhash_length,\n",
    "    num_buckets=num_bands,\n",
    "    buckets_per_shuffle=buckets_per_shuffle,\n",
    "    id_fields=id_field,\n",
    "    minhash_field=minhash_field,\n",
    "    logger=log_dir,\n",
    ")\n",
    "\n",
    "lsh_result = lsh(DocumentDataset(df))\n",
    "print(f\"LSH took {time.time() - t0} s\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d14ba2f-cfb8-448e-9bb8-34af6005ba15",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "lsh_result.df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8dd9dbc9",
   "metadata": {},
   "source": [
    "### 5.2.3 Buckets to Edges\n",
    "\n",
    "`BucketsToEdges` is designed to take the bucket information from the output of LSH and create an edgelist dataset where documents with the same `_bucket_id` are connected with an edge between them. This edgelist can then be passed on the connected components to identify groups of similar documents across buckets. Since the false positive check is skipped all documents within a bucket are considered to be duplicates of each other and assigned a jaccard similarity of 1.0 to avoid edge removal during the next step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7985cf1a-9d88-4844-8ce4-e68d9792118c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from nemo_curator import BucketsToEdges\n",
    "\n",
    "id_field = \"id\"\n",
    "\n",
    "cache_dir = os.path.join(base_dir, \"fuzzy-dedup-output-2023-06\")\n",
    "input_bucket_path = os.path.join(cache_dir, \"_buckets.parquet\")\n",
    "input_bucket_field = \"_bucket_id\"\n",
    "log_dir = os.path.join(base_dir, \"logs\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0700327-5171-4673-8ded-c9f43492f582",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "t0 = time.time()\n",
    "\n",
    "# Read \"_buckets.parquet\"\n",
    "ddf_bk = DocumentDataset.read_parquet(\n",
    "    input_bucket_path,\n",
    "    backend=\"cudf\",\n",
    ")\n",
    "\n",
    "# Run _MapBuckets()\n",
    "buckets_to_edges = BucketsToEdges(\n",
    "    cache_dir=cache_dir,\n",
    "    id_fields=id_field,\n",
    "    bucket_field=input_bucket_field,\n",
    "    logger=log_dir,\n",
    ")\n",
    "\n",
    "edgelist_df = buckets_to_edges(ddf_bk)\n",
    "\n",
    "\n",
    "print(f\"Buckets to Edgelist took {time.time() - t0} s\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a986f76-191e-436d-9df3-bb68dc78d365",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "edgelist_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb9e1287-bd19-4728-a4c8-b92b39ca1fcc",
   "metadata": {},
   "source": [
    "### 5.2.4 Connected Component\n",
    "\n",
    "After all buckets were processed and duplicates (at the threshold) were approximately discovered,\n",
    "we constructed a sparse document graph and found the connected components therein (using scipy). Each\n",
    "connected component represents a set of documents that we consider similar enough to be duplicates, and\n",
    "from which we select a single representative."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9aeb619-3fab-4a18-b582-bccae3eefd17",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from nemo_curator import ConnectedComponents\n",
    "\n",
    "cache_dir = expand_outdir_and_mkdir(\n",
    "    os.path.join(base_dir, \"fuzzy-dedup-output-2023-06/cc-cache\"),\n",
    ")\n",
    "edgelist_path = os.path.join(base_dir, \"fuzzy-dedup-output-2023-06/jaccard_similarity_results.parquet\")\n",
    "id_field = \"id\"\n",
    "\n",
    "output_path = expand_outdir_and_mkdir(\n",
    "    os.path.join(base_dir, \"fuzzy-dedup-output-2023-06/connected_components.parquet\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6bee85f3-5477-4b9c-b606-7bbbefbe6cfc",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "t0 = time.time()\n",
    "components_stage = ConnectedComponents(\n",
    "    cache_dir=cache_dir,\n",
    "    jaccard_pairs_path=edgelist_path,\n",
    "    id_column=id_field,\n",
    ")\n",
    "components_stage.cc_workflow(output_path=output_path)\n",
    "print(f\"Connected Component took {time.time() - t0} seconds\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15214dcf-ff49-439e-b3d7-d8666d081027",
   "metadata": {},
   "source": [
    "Let's check the results of connected components step. We can see that 239,037,733 are identified as duplicates to be removed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "94e8126d-af15-4182-98cd-10df06e9778e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "num of docs to remove = 239037733\n"
     ]
    }
   ],
   "source": [
    "output_path = os.path.join(base_dir, \"fuzzy-dedup-output-2023-06/connected_components.parquet\")\n",
    "cc_result = dask_cudf.read_parquet(output_path, split_row_groups=False).repartition(npartitions=1)\n",
    "\n",
    "# Set 'group' as the index and shuffle to ensure all same 'group' values are in the same partition\n",
    "cc_result = cc_result.set_index(\"group\", shuffle=\"tasks\")\n",
    "\n",
    "\n",
    "# Define a function to assign cumulative counts and filter duplicates\n",
    "def assign_cumcount(df: cudf.DataFrame) -> cudf.DataFrame:\n",
    "    df[\"cumcount\"] = df.groupby(level=0).cumcount()\n",
    "    df = df[df[\"cumcount\"] >= 1]\n",
    "    return df.drop(columns=[\"cumcount\"])\n",
    "\n",
    "\n",
    "# Find duplicates by applying the function to each partition\n",
    "docs_to_remove = cc_result.map_partitions(assign_cumcount, meta=cc_result)\n",
    "\n",
    "# Reset the index\n",
    "docs_to_remove = docs_to_remove.reset_index()\n",
    "\n",
    "docs_to_remove = docs_to_remove[[\"id\"]]\n",
    "docs_to_remove = docs_to_remove.rename(columns={\"id\": \"to_remove_doc_id\"})\n",
    "docs_to_remove = docs_to_remove.reset_index(drop=True).persist()\n",
    "_ = wait(docs_to_remove)\n",
    "del _\n",
    "\n",
    "print(\"num of docs to remove =\", len(docs_to_remove))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "568ee0b5-f2dd-4d34-917f-56f4211a36fe",
   "metadata": {},
   "source": [
    "We can examine the size of the duplicate clusters. The largest cluster has 775,379 near duplicates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "cae7f166-836a-4c21-bff2-7453254956b7",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "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>count</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>group</th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>350652173</th>\n",
       "      <td>775379</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>93521324</th>\n",
       "      <td>493227</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24</th>\n",
       "      <td>112861</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>319292355</th>\n",
       "      <td>96224</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>70141069</th>\n",
       "      <td>67474</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "            count\n",
       "group            \n",
       "350652173  775379\n",
       "93521324   493227\n",
       "24         112861\n",
       "319292355   96224\n",
       "70141069    67474"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cc_grouped = (\n",
    "    cc_result.groupby(\"group\")\n",
    "    .agg({\"id\": \"count\"})\n",
    "    .rename(columns={\"id\": \"count\"})\n",
    "    .sort_values(\"count\", ascending=False)\n",
    "    .compute()\n",
    ")\n",
    "cc_grouped.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0def7323-3d2c-4861-9b7e-a1e296ccf329",
   "metadata": {},
   "source": [
    "[Optional] Verify if fuzzy duplicates are similar. For example, we can look into the largest group \"350652173\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e22cb491-c2ab-4ec4-8313-ae2bcd66a352",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "dup_group = cc_result.loc[350652173].compute()\n",
    "dup_group.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "170c1cf4-8cb9-4f10-aab3-acfdaa9e5b16",
   "metadata": {},
   "source": [
    "We will examine the first five documents in this cluster:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "00cf923e-fd4e-41b9-a00f-801c186ac70e",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Reading 37848 files\n"
     ]
    }
   ],
   "source": [
    "# read input dataset\n",
    "input_data_dir = os.path.join(base_dir, \"rpv2-2023-06-en-cleaned\")\n",
    "input_dataset = DocumentDataset.read_json(input_data_dir, add_filename=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9772bf71-9e18-4e59-b9f8-ebd9053c79b0",
   "metadata": {},
   "source": [
    "Let's visualize the content of these documents and see if they are similar (ids may change so revise the `dup_ids` as needed)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "e3cc167f-30f8-470d-99e3-0a2d916d46bf",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Searching for near duplicate examples with specific IDs took 610.5046670436859 seconds\n"
     ]
    }
   ],
   "source": [
    "t0 = time.time()\n",
    "dup_ids = [\n",
    "    \"rpv2-2023-06-1285625132\",\n",
    "    \"rpv2-2023-06-2033200488\",\n",
    "    \"rpv2-2023-06-0428016172\",\n",
    "    \"rpv2-2023-06-1268721963\",\n",
    "    \"rpv2-2023-06-1285428574\",\n",
    "]\n",
    "dup_examples = input_dataset.df[input_dataset.df[\"id\"].isin(dup_ids)].compute()\n",
    "print(f\"Searching for near duplicate examples with specific IDs took {time.time() - t0} seconds\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "655a4b3e-8a48-441c-8e12-9b2d287b79ec",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "dup_examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49b1b00a-501e-4d49-a93b-60a5c8ae87d2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(\"Example duplicate 1\\n\" + dup_examples.raw_content.iloc[0])\n",
    "print(\"\\n\\nExample duplicate 2\\n\" + dup_examples.raw_content.iloc[1])\n",
    "print(\"\\n\\nExample duplicate 3\\n\" + dup_examples.raw_content.iloc[2])\n",
    "print(\"\\n\\nExample duplicate 4\\n\" + dup_examples.raw_content.iloc[3])\n",
    "print(\"\\n\\nExample duplicate 4\\n\" + dup_examples.raw_content.iloc[4])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c428e09c-5442-4908-8888-4e62994e4c5c",
   "metadata": {},
   "source": [
    "### 5.2.5 Duplicates Removal"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2e01b84-07cc-45a3-9dde-97884d1922a3",
   "metadata": {},
   "source": [
    "Next, we will proceed to remove the duplicates identified from the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7cde97c5-cfaa-4096-8d9f-1ffa19c4adb2",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Reading 37848 files\n"
     ]
    }
   ],
   "source": [
    "input_dataset = DocumentDataset.read_json(os.path.join(base_dir, \"rpv2-2023-06-en-cleaned\"), backend=\"cudf\")\n",
    "input_df = input_dataset.df[[\"raw_content\", \"id\"]]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99e50a79-18b0-4d1b-bab0-cfd1ec9b62fd",
   "metadata": {},
   "source": [
    "Then, we will perform a merge between the `input_df` and the `docs_to_remove` on the IDs and drop the fuzzy duplicates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f819d1d3-c4c0-4288-b190-86276d221050",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Removing duplicates and writing deduped dataset took 1241.3191509246826 seconds\n"
     ]
    }
   ],
   "source": [
    "dedup_output_dir = expand_outdir_and_mkdir(os.path.join(base_dir, \"rpv2-2023-06-deduped\"))\n",
    "deduped_df = input_df.merge(docs_to_remove, left_on=[\"id\"], right_on=[\"to_remove_doc_id\"], how=\"left\")\n",
    "\n",
    "deduped_df = (\n",
    "    deduped_df[deduped_df[\"to_remove_doc_id\"].isna()].drop(columns=[\"to_remove_doc_id\"]).reset_index(drop=True)\n",
    ")\n",
    "\n",
    "t0 = time.time()\n",
    "deduped_df.to_parquet(dedup_output_dir)\n",
    "print(f\"Removing duplicates and writing deduped dataset took {time.time() - t0} seconds\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f987b1ac-c8e3-45f8-8e9b-32befc6667aa",
   "metadata": {},
   "source": [
    "To verify the results, we can confirm that we have 849,273,787 documents left compared to 1,088,311,520 in the input dataset, essentially removing 239,037,733 duplicates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "cf56c9ca-27cb-4f03-921b-685d523cf43e",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "849273787"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(deduped_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "19f96f15-1ac6-40ab-9e04-1586531bb55f",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1088311520"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(input_df)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cdc6350e-2363-4b13-ac67-0cbc23ad981d",
   "metadata": {},
   "source": [
    "## 5.3 Inter-snapshot Deduplication"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "888c2b15-961f-4a73-a0a3-15474ae4134c",
   "metadata": {},
   "source": [
    "So far we have deduplicated a single snapshot from rpv2. A pre-training dataset can include multiple snapshots so we will often need to perform inter-snapshot deduplication. For this tutorial, we will demonstrate deduplication across two snapshots as an example.\n",
    "\n",
    "We first performed all the above steps for another snapshot `2023-14` and then combined the two deduped datasets into one and stored them in `rpv2-2023-06-and-14-deduped`.\n",
    "\n",
    "Next, we will perform the fuzzy deduplication on the combined dataset."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a1445cc-b69c-4007-8f09-75a8eb8f699c",
   "metadata": {},
   "source": [
    "### 5.3.1 Compute Minhash"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1461b61-887c-4099-bd9f-32e79dc5fdbb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from nemo_curator import LSH, BucketsToEdges, ConnectedComponents, MinHash\n",
    "from nemo_curator.utils.file_utils import reshard_jsonl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "efaed1ed-e6d1-4117-9b0b-fe0d20960b60",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "seed = 42\n",
    "minhash_length = 260\n",
    "char_ngram = 24\n",
    "log_dir = expand_outdir_and_mkdir(os.path.join(base_dir, \"logs\"))\n",
    "id_field = \"id\"\n",
    "text_field = \"raw_content\"\n",
    "minshah_output_dir = expand_outdir_and_mkdir(os.path.join(base_dir, \"rpv2-2023-06-and-14-minhash\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "b4bf2d09-6601-4bd2-a6f2-f738cffd8885",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "input_data_dir = os.path.join(base_dir, \"rpv2-2023-06-and-14-deduped\")\n",
    "\n",
    "files = []\n",
    "for file in os.listdir(input_data_dir):\n",
    "    if file.endswith(\".part\"):\n",
    "        new_file = file.replace(\".part\", \".jsonl\")\n",
    "        old_file_path = os.path.join(input_data_dir, file)\n",
    "        new_file_path = os.path.join(input_data_dir, new_file)\n",
    "        os.rename(old_file_path, new_file_path)\n",
    "    files.append(new_file_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "38197174-738e-42a4-a38a-1dbd7d84836d",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Reading 72797 files\n"
     ]
    }
   ],
   "source": [
    "files = [f for f in files if f.endswith(\".jsonl\")]\n",
    "df = read_data(\n",
    "    files,\n",
    "    file_type=\"jsonl\",\n",
    "    backend=\"cudf\",\n",
    "    files_per_partition=2,\n",
    "    add_filename=False,\n",
    ")[[id_field, text_field]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "439add9c-9f51-4481-95cf-456dc5be9fd2",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Computing minhashes took:6115.702769517899\n"
     ]
    }
   ],
   "source": [
    "t0 = time.time()\n",
    "\n",
    "# Run MinHash() on input data\n",
    "minhasher = MinHash(\n",
    "    seed=seed,\n",
    "    num_hashes=minhash_length,\n",
    "    char_ngrams=char_ngram,\n",
    "    use_64bit_hash=False,\n",
    "    logger=log_dir,\n",
    "    id_field=id_field,\n",
    "    text_field=text_field,\n",
    "    cache_dir=minshah_output_dir,\n",
    ")\n",
    "\n",
    "result = minhasher(DocumentDataset(df)).df\n",
    "\n",
    "print(f\"Computing minhashes took:{time.time() - t0}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "1c998d64-54f8-49b0-8e0c-2e5727596e84",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "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>id</th>\n",
       "      <th>_minhash_signature</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>rpv2-2023-06-0678400000</td>\n",
       "      <td>[36422228, 15993596, 3538361, 16103012, 194100...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>rpv2-2023-06-0678500000</td>\n",
       "      <td>[34662, 17635, 1112347, 293654, 313382, 160184...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>rpv2-2023-06-0678600000</td>\n",
       "      <td>[15076006, 1801689, 3181854, 2949398, 5699436,...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>rpv2-2023-06-0678700000</td>\n",
       "      <td>[13528976, 2438382, 26260517, 26187347, 249748...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>rpv2-2023-06-0678800000</td>\n",
       "      <td>[2550974, 157261, 1536526, 1169030, 576861, 10...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                        id                                 _minhash_signature\n",
       "0  rpv2-2023-06-0678400000  [36422228, 15993596, 3538361, 16103012, 194100...\n",
       "1  rpv2-2023-06-0678500000  [34662, 17635, 1112347, 293654, 313382, 160184...\n",
       "2  rpv2-2023-06-0678600000  [15076006, 1801689, 3181854, 2949398, 5699436,...\n",
       "3  rpv2-2023-06-0678700000  [13528976, 2438382, 26260517, 26187347, 249748...\n",
       "4  rpv2-2023-06-0678800000  [2550974, 157261, 1536526, 1169030, 576861, 10..."
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2e3af9e-9e03-4950-93c0-92792f9ad24b",
   "metadata": {},
   "source": [
    "### 5.3.2 Minhash LSH"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b11c2f37-3b78-4e1b-a9ff-4a89b38f3604",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "lsh_input_dir = os.path.join(base_dir, \"rpv2-2023-06-and-14-minhash\")\n",
    "id_field = \"id\"\n",
    "output_bucket_dir = expand_outdir_and_mkdir(os.path.join(base_dir, \"fuzzy-dedup-output-2023-06-and-14\"))\n",
    "num_bands = 20\n",
    "buckets_per_shuffle = 1\n",
    "minhash_field = \"_minhash_signature\"\n",
    "minhash_length = 260\n",
    "log_dir = os.path.join(base_dir, \"logs\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "a243ed7a-9175-488f-8097-5b82c47c5708",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "LSH took 10536.635195493698 s\n"
     ]
    }
   ],
   "source": [
    "t0 = time.time()\n",
    "\n",
    "# Load MinHash output\n",
    "df = dask_cudf.read_parquet(lsh_input_dir, blocksize=\"2GB\", aggregate_files=True)\n",
    "\n",
    "lsh = LSH(\n",
    "    cache_dir=output_bucket_dir,\n",
    "    num_hashes=minhash_length,\n",
    "    num_buckets=num_bands,\n",
    "    buckets_per_shuffle=buckets_per_shuffle,\n",
    "    id_fields=id_field,\n",
    "    minhash_field=minhash_field,\n",
    "    logger=log_dir,\n",
    ")\n",
    "\n",
    "lsh_result = lsh(DocumentDataset(df))\n",
    "print(f\"LSH took {time.time() - t0} s\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ea85033-e275-4b62-8351-6c85ac5ac83b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "lsh_result.df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff8c79b8-229f-47ee-9e01-a7bf1f317250",
   "metadata": {},
   "source": [
    "### 5.3.3 Buckets to Edges"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "2cc909e1-e02e-433d-b6c2-e7f82a137438",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "input_data_paths = [os.path.join(base_dir, \"rpv2-2023-06-and-14-deduped\")]\n",
    "\n",
    "id_field = \"id\"\n",
    "cache_dir = os.path.join(base_dir, \"fuzzy-dedup-output-2023-06-and-14\")\n",
    "input_bucket_path = os.path.join(cache_dir, \"_buckets.parquet\")\n",
    "input_bucket_field = \"_bucket_id\"\n",
    "log_dir = os.path.join(base_dir, \"logs\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4e00d8e-170c-4adf-bf34-1ad1b5275760",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "t0 = time.time()\n",
    "\n",
    "# Read \"_buckets.parquet\"\n",
    "ddf_bk = DocumentDataset.read_parquet(\n",
    "    input_bucket_path,\n",
    "    backend=\"cudf\",\n",
    ")\n",
    "\n",
    "# Run _MapBuckets()\n",
    "buckets_to_edges = BucketsToEdges(\n",
    "    cache_dir=cache_dir,\n",
    "    id_fields=id_field,\n",
    "    bucket_field=input_bucket_field,\n",
    "    logger=log_dir,\n",
    ")\n",
    "\n",
    "edgelist_df = buckets_to_edges(ddf_bk)\n",
    "\n",
    "\n",
    "print(f\"Buckets to Edgelist took {time.time() - t0} s\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0718d8e-5143-458f-ab59-a440adcde8b8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "edgelist_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2307ab3-6d75-4e96-b2f3-dc180aed06ef",
   "metadata": {},
   "source": [
    "### 5.3.4 Connected Component"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "fd602c6e-39e5-45fe-9031-ea3926d68a68",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "cache_dir = expand_outdir_and_mkdir(\n",
    "    os.path.join(base_dir, \"fuzzy-dedup-output-2023-06-and-14/cc-cache\"),\n",
    ")\n",
    "edgelist_path = os.path.join(base_dir, \"fuzzy-dedup-output-2023-06-and-14/jaccard_similarity_results.parquet\")\n",
    "id_field = \"id\"\n",
    "output_path = expand_outdir_and_mkdir(\n",
    "    os.path.join(base_dir, \"fuzzy-dedup-output-2023-06-and-14/connected_components.parquet\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89398db9-d4e6-48ec-bad8-1d5ac553cadd",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "t0 = time.time()\n",
    "components_stage = ConnectedComponents(\n",
    "    cache_dir=cache_dir,\n",
    "    jaccard_pairs_path=edgelist_path,\n",
    "    id_column=id_field,\n",
    ")\n",
    "components_stage.cc_workflow(output_path=output_path)\n",
    "print(f\"Connected Component took {time.time() - t0} seconds\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5d579d9-0344-45a6-a565-78b1fbeddf9c",
   "metadata": {},
   "source": [
    "### 5.3.5 Duplicates Removal"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5aa8c531-10ec-448d-8413-f059b796226d",
   "metadata": {},
   "source": [
    "From the outputs of the Connect Component step, we can see that inter-snapshot dedup found 81,764,804 duplicates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d2bd2d15-5c09-4fb1-a289-ce228f90713e",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "docs_to_remove 81764804\n"
     ]
    }
   ],
   "source": [
    "output_path = os.path.join(base_dir, \"fuzzy-dedup-output-2023-06-and-14/connected_components.parquet\")\n",
    "cc_result = dask_cudf.read_parquet(output_path, split_row_groups=False).repartition(npartitions=1)\n",
    "\n",
    "# Set 'group' as the index and shuffle to ensure all same 'group' values are in the same partition\n",
    "cc_result = cc_result.set_index(\"group\", shuffle=\"tasks\")\n",
    "\n",
    "\n",
    "# Define a function to assign cumulative counts and filter duplicates\n",
    "def assign_cumcount(df: cudf.DataFrame) -> cudf.DataFrame:\n",
    "    df[\"cumcount\"] = df.groupby(level=0).cumcount()\n",
    "    df = df[df[\"cumcount\"] >= 1]\n",
    "    return df.drop(columns=[\"cumcount\"])\n",
    "\n",
    "\n",
    "# Find duplicates by applying the function to each partition\n",
    "docs_to_remove = cc_result.map_partitions(assign_cumcount, meta=cc_result)\n",
    "\n",
    "# Reset the index\n",
    "docs_to_remove = docs_to_remove.reset_index()\n",
    "\n",
    "docs_to_remove = docs_to_remove[\"id\"]\n",
    "docs_to_remove = docs_to_remove.rename(columns={\"id\": \"to_remove_doc_id\"})\n",
    "docs_to_remove = docs_to_remove.reset_index(drop=True).persist()\n",
    "_ = wait(docs_to_remove)\n",
    "del _\n",
    "\n",
    "print(\"docs_to_remove\", len(docs_to_remove))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2bc75c4d-0eec-496f-b02f-a639b193522f",
   "metadata": {},
   "source": [
    "Before proceeding to duplicates removal, we suggest resharding the data to fix potentially empty partitions due to duplicates removal for single snapshots."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "a6c6d7d6-9db5-4504-9d03-67c7ed69c409",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Data sharding took:904.7163739204407\n"
     ]
    }
   ],
   "source": [
    "output_resharded_dir = expand_outdir_and_mkdir(os.path.join(base_dir, \"rpv2-2023-06-and-14-deduped-resharded\"))\n",
    "\n",
    "t0 = time.time()\n",
    "reshard_jsonl(\n",
    "    os.path.join(base_dir, \"rpv2-2023-06-and-14-deduped\"),\n",
    "    output_resharded_dir,\n",
    "    output_file_size=\"100M\",\n",
    "    start_index=0,\n",
    "    file_prefix=\"rpv2-2023-06-and-14-deduped\",\n",
    ")\n",
    "print(f\"Data sharding took:{time.time() - t0}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "63546dd3-1560-4b38-91ce-84c004636657",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Reading 72780 files\n"
     ]
    }
   ],
   "source": [
    "input_dataset = DocumentDataset.read_json(\n",
    "    os.path.join(base_dir, \"rpv2-2023-06-and-14-deduped-resharded\"),\n",
    "    backend=\"cudf\",\n",
    ")\n",
    "input_df = input_dataset.df[[\"raw_content\", \"id\"]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "88dfe6a9-a463-4d43-9f79-0d3df960961c",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Removing duplicates and writing deduped dataset took 2084.46063041687 seconds\n"
     ]
    }
   ],
   "source": [
    "dedup_output_dir = expand_outdir_and_mkdir(os.path.join(base_dir, \"/rpv2-2023-06-and-14-inter-deduped\"))\n",
    "deduped_df = input_df.merge(docs_to_remove, left_on=[\"id\"], right_on=[\"to_remove_doc_id\"], how=\"left\")\n",
    "\n",
    "deduped_df = (\n",
    "    deduped_df[deduped_df[\"to_remove_doc_id\"].isna()].drop(columns=[\"to_remove_doc_id\"]).reset_index(drop=True)\n",
    ")\n",
    "\n",
    "t0 = time.time()\n",
    "deduped_df.to_parquet(dedup_output_dir)\n",
    "print(f\"Removing duplicates and writing deduped dataset took {time.time() - t0} seconds\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "473fea1e-45e5-423c-9187-17867c0ad2a7",
   "metadata": {},
   "source": [
    "We can verify that the deduped dataset has 1,585,546,179 documents, compared to 1,667,310,983 documents befoe dedup."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d174ea0f-440f-4311-a9c0-d3cc26683a81",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1585546179"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(deduped_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "014e1a79-4967-4392-8cb8-6d822a3f57ca",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1667310983"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(input_df)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b9977f5-fa49-4e63-9057-50fadba58a86",
   "metadata": {},
   "source": [
    "# 6. Quality Filtering\n",
    "<a id=\"filter\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e19b91e6-3156-4a1b-816b-500a54def99d",
   "metadata": {},
   "source": [
    "Web crawled dataset often has low quality documents that we do not want the model to learn from. We can perform quality filtering to remove low quality data. NeMo Curator offers modules for both classifier-based and heuristic-based filtering. In this tutorial, we will perform heuristic filtering using a list of heuristic filters to improve data quality.\n",
    "\n",
    "Curator provides a generic list of heuristic filters but for this tutorial, we only select 10 filters for demo purposes. The selected filters are given in `config/heuristic_filter_en.yaml`.\n",
    "\n",
    "Heuristic filtering in Curator is a cpu module so we will need to use the cpu cluter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a1347d09-8bc5-453b-aa4a-4b75c4b3b47a",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Num Workers = 256\n"
     ]
    }
   ],
   "source": [
    "scheduler_address = os.getenv(\"SCHEDULER_ADDRESS\")\n",
    "cpu_client = get_client(scheduler_address=scheduler_address)\n",
    "print(f\"Num Workers = {get_num_workers(cpu_client)}\", flush=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49273a8b-848f-4f24-a0ba-3c0b478d17cc",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from nemo_curator.utils.config_utils import build_filter_pipeline\n",
    "\n",
    "filter_config_file = os.path.join(base_dir, \"config/heuristic_filter_en.yaml\")\n",
    "hf_input_data_dir = os.path.join(base_dir, \"rpv2-2023-06-and-14-inter-deduped\")\n",
    "kept_document_dir = expand_outdir_and_mkdir(\n",
    "    os.path.join(base_dir, \"rpv2-2023-06-and-14-heuristic-filtering\", \"hf.parquet\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e5284b6f-c87e-4027-a802-ab08b92610cd",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Reading 72780 files\n",
      "Writing to disk complete for 72780 partitions\n",
      "Time taken for Heuristic filtering: 5647.508106470108 s\n"
     ]
    }
   ],
   "source": [
    "t0 = time.time()\n",
    "\n",
    "# Load dataset\n",
    "dataset = DocumentDataset.read_parquet(hf_input_data_dir)\n",
    "\n",
    "# construct pipeline from config\n",
    "filter_pipeline = build_filter_pipeline(filter_config_file)\n",
    "\n",
    "# filter data and write to disk\n",
    "filtered_dataset = filter_pipeline(dataset)\n",
    "filtered_dataset.to_parquet(kept_document_dir)\n",
    "\n",
    "print(f\"Time taken for Heuristic filtering: {time.time() - t0} s\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ce0658a-97f9-412c-8eef-68f6ab0b4be6",
   "metadata": {},
   "source": [
    "After filitering, we have 1,229,679,047 documents left, removing 355,867,132 documents from the deduped dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "5da60070-f78c-43fe-8bc0-fcbd839b7021",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1229679047"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(filtered_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c92ae5e-7397-4ad9-9dec-cb93eefc3dde",
   "metadata": {},
   "source": [
    "[Optional] Examine example low quality documents:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "10e4b615-1eaa-4050-b191-28a48166a560",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from helper import get_dataframe_complement\n",
    "\n",
    "original_df = dd.read_parquet(hf_input_data_dir)\n",
    "filtered_df = dd.read_parquet(kept_document_dir)\n",
    "removed_df = get_dataframe_complement(original_df, filtered_df)\n",
    "removed_df_example = removed_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "468233aa-d2a9-4e80-a815-e1a645213c75",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(removed_df_example.raw_content.iloc[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aee129f8-1fc5-401c-aa73-42fb254539c2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(removed_df_example.raw_content.iloc[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa60a9b9-9158-43d6-9ea4-1f544c6f816e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}