From bbbf0bb2761d22503a6be9b1e6d25c966171ab4b Mon Sep 17 00:00:00 2001
From: Sarah Yurick <sarahyurick@gmail.com>
Date: Wed, 12 Mar 2025 11:55:02 -0700
Subject: [PATCH] add async llm modifier and llm modifier

Signed-off-by: Sarah Yurick <sarahyurick@gmail.com>
---
 examples/README.md                            |  32 +--
 examples/async_llm_pii_redaction.py           |  53 +++++
 examples/llm_pii_redaction.py                 |  52 ++++
 .../modifiers/async_llm_pii_modifier.py       | 214 +++++++++++++++++
 nemo_curator/modifiers/llm_pii_modifier.py    | 176 ++++++++++++++
 nemo_curator/utils/llm_pii_utils.py           | 222 ++++++++++++++++++
 6 files changed, 734 insertions(+), 15 deletions(-)
 create mode 100644 examples/async_llm_pii_redaction.py
 create mode 100644 examples/llm_pii_redaction.py
 create mode 100644 nemo_curator/modifiers/async_llm_pii_modifier.py
 create mode 100644 nemo_curator/modifiers/llm_pii_modifier.py
 create mode 100644 nemo_curator/utils/llm_pii_utils.py

diff --git a/examples/README.md b/examples/README.md
index 29545978c..6a1b9c623 100644
--- a/examples/README.md
+++ b/examples/README.md
@@ -4,21 +4,23 @@ This directory contains multiple Python scripts with examples of how to use vari
 The goal of these examples is to give the user an overview of many of the ways your text data can be curated.
 These include:
 
-| Python Script                         | Description                                                                                                   |
-|---------------------------------------|---------------------------------------------------------------------------------------------------------------|
-| blend_and_shuffle.py                  | Combine multiple datasets into one with different amounts of each dataset, then randomly permute the dataset. |
-| classifier_filtering.py               | Train a fastText classifier, then use it to filter high and low quality data.                                 |
-| download_arxiv.py                     | Download Arxiv tar files and extract them.                                                                    |
-| download_common_crawl.py              | Download Common Crawl WARC snapshots and extract them.                                                        |
-| download_wikipedia.py                 | Download the latest Wikipedia dumps and extract them.                                                         |
-| exact_deduplication.py                | Use the `ExactDuplicates` class to perform exact deduplication on text data.                                  |
-| find_pii_and_deidentify.py            | Use the `PiiModifier` and `Modify` classes to remove personally identifiable information from text data.      |
-| fuzzy_deduplication.py                | Use the `FuzzyDuplicatesConfig` and `FuzzyDuplicates` classes to perform fuzzy deduplication on text data.    |
-| identify_languages.py                 | Use `FastTextLangId` to filter data by language                                                               |
-| raw_download_common_crawl.py          | Download the raw compressed WARC files from Common Crawl without extracting them.                             |
-| semdedup_example.py                   | Use the `SemDedup` class to perform semantic deduplication on text data.                                      |
-| task_decontamination.py               | Remove segments of downstream evaluation tasks from a dataset.                                                |
-| translation_example.py                | Create and use an `IndicTranslation` model for language translation.                                          |
+| Python Script                         | Description                                                                                                      |
+|---------------------------------------|------------------------------------------------------------------------------------------------------------------|
+| async_llm_pii_redaction.py            | Use the `AsyncLLMPiiModifier` and `Modify` classes to remove personally identifiable information from text data. |
+| blend_and_shuffle.py                  | Combine multiple datasets into one with different amounts of each dataset, then randomly permute the dataset.    |
+| classifier_filtering.py               | Train a fastText classifier, then use it to filter high and low quality data.                                    |
+| download_arxiv.py                     | Download Arxiv tar files and extract them.                                                                       |
+| download_common_crawl.py              | Download Common Crawl WARC snapshots and extract them.                                                           |
+| download_wikipedia.py                 | Download the latest Wikipedia dumps and extract them.                                                            |
+| exact_deduplication.py                | Use the `ExactDuplicates` class to perform exact deduplication on text data.                                     |
+| find_pii_and_deidentify.py            | Use the `PiiModifier` and `Modify` classes to remove personally identifiable information from text data.         |
+| fuzzy_deduplication.py                | Use the `FuzzyDuplicatesConfig` and `FuzzyDuplicates` classes to perform fuzzy deduplication on text data.       |
+| identify_languages.py                 | Use `FastTextLangId` to filter data by language                                                                  |
+| llm_pii_redaction.py                  | Use the `LLMPiiModifier` and `Modify` classes to remove personally identifiable information from text data.      |
+| raw_download_common_crawl.py          | Download the raw compressed WARC files from Common Crawl without extracting them.                                |
+| semdedup_example.py                   | Use the `SemDedup` class to perform semantic deduplication on text data.                                         |
+| task_decontamination.py               | Remove segments of downstream evaluation tasks from a dataset.                                                   |
+| translation_example.py                | Create and use an `IndicTranslation` model for language translation.                                             |
 
 Before running any of these scripts, we strongly recommend displaying `python <script name>.py --help` to ensure that any needed or relevant arguments are specified.
 
diff --git a/examples/async_llm_pii_redaction.py b/examples/async_llm_pii_redaction.py
new file mode 100644
index 000000000..46de1bcd1
--- /dev/null
+++ b/examples/async_llm_pii_redaction.py
@@ -0,0 +1,53 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import dask.dataframe
+import pandas as pd
+
+from nemo_curator.datasets import DocumentDataset
+from nemo_curator.modifiers.async_llm_pii_modifier import AsyncLLMPiiModifier
+from nemo_curator.modules.modify import Modify
+from nemo_curator.utils.distributed_utils import get_client
+
+
+def console_script():
+    _ = get_client()
+
+    dataframe = pd.DataFrame(
+        {
+            "text": [
+                # Sampled from https://huggingface.co/datasets/gretelai/gretel-pii-masking-en-v1
+                "Transaction details: gasLimit set to 1000000 units by tw_brian740, gasPrice set to 10 Gwei by veronicawood@example.org, contactable at +1-869-341-9301x7005, located at Suite 378, Yolanda Mountain, Burkeberg.",
+                "Unloading Plan for Shipment MRN-293104, MED25315002, dated 1989.12.22. Driver EMP730359, Vehicle KS40540825.",
+            ]
+        }
+    )
+    dd = dask.dataframe.from_pandas(dataframe, npartitions=1)
+    dataset = DocumentDataset(dd)
+
+    modifier = AsyncLLMPiiModifier(
+        # Endpoint for the user's NIM
+        base_url="http://0.0.0.0:8000/v1",
+        api_key="API KEY (if needed)",
+        model="meta/llama-3.1-70b-instruct",
+        max_concurrent_requests=10,
+    )
+
+    modify = Modify(modifier)
+    modified_dataset = modify(dataset)
+    modified_dataset.df.to_json("output.jsonl", lines=True, orient="records")
+
+
+if __name__ == "__main__":
+    console_script()
diff --git a/examples/llm_pii_redaction.py b/examples/llm_pii_redaction.py
new file mode 100644
index 000000000..c13442d79
--- /dev/null
+++ b/examples/llm_pii_redaction.py
@@ -0,0 +1,52 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import dask.dataframe
+import pandas as pd
+
+from nemo_curator.datasets import DocumentDataset
+from nemo_curator.modifiers.llm_pii_modifier import LLMPiiModifier
+from nemo_curator.modules.modify import Modify
+from nemo_curator.utils.distributed_utils import get_client
+
+
+def console_script():
+    _ = get_client()
+
+    dataframe = pd.DataFrame(
+        {
+            "text": [
+                # Sampled from https://huggingface.co/datasets/gretelai/gretel-pii-masking-en-v1
+                "Transaction details: gasLimit set to 1000000 units by tw_brian740, gasPrice set to 10 Gwei by veronicawood@example.org, contactable at +1-869-341-9301x7005, located at Suite 378, Yolanda Mountain, Burkeberg.",
+                "Unloading Plan for Shipment MRN-293104, MED25315002, dated 1989.12.22. Driver EMP730359, Vehicle KS40540825.",
+            ]
+        }
+    )
+    dd = dask.dataframe.from_pandas(dataframe, npartitions=1)
+    dataset = DocumentDataset(dd)
+
+    modifier = LLMPiiModifier(
+        # Endpoint for the user's NIM
+        base_url="http://0.0.0.0:8000/v1",
+        api_key="API KEY (if needed)",
+        model="meta/llama-3.1-70b-instruct",
+    )
+
+    modify = Modify(modifier)
+    modified_dataset = modify(dataset)
+    modified_dataset.df.to_json("output.jsonl", lines=True, orient="records")
+
+
+if __name__ == "__main__":
+    console_script()
diff --git a/nemo_curator/modifiers/async_llm_pii_modifier.py b/nemo_curator/modifiers/async_llm_pii_modifier.py
new file mode 100644
index 000000000..ec3abafe3
--- /dev/null
+++ b/nemo_curator/modifiers/async_llm_pii_modifier.py
@@ -0,0 +1,214 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+import json
+import warnings
+from typing import Any, Coroutine, Dict, List, Optional
+
+import pandas as pd
+from openai import AsyncOpenAI
+from tqdm.asyncio import tqdm
+
+from nemo_curator.modifiers import DocumentModifier
+from nemo_curator.utils.decorators import batched
+from nemo_curator.utils.distributed_utils import load_object_on_worker
+from nemo_curator.utils.llm_pii_utils import (
+    JSON_SCHEMA,
+    PII_LABELS,
+    SYSTEM_PROMPT,
+    redact,
+    validate_entity,
+)
+
+__all__ = ["AsyncLLMPiiModifier"]
+
+
+class AsyncLLMInference:
+    """A class for redacting PII via asynchronous LLM inference"""
+
+    def __init__(
+        self,
+        base_url: str,
+        api_key: Optional[str] = None,
+        model: str = "meta/llama-3.1-70b-instruct",
+        system_prompt: str = SYSTEM_PROMPT,
+        pii_labels: List[str] = PII_LABELS,
+    ):
+        self.client = AsyncOpenAI(base_url=base_url, api_key=api_key)
+        self.model = model
+        self.system_prompt = system_prompt[model]
+        self.pii_labels = pii_labels
+
+    async def infer(self, text: str) -> List[Dict[str, str]]:
+        """Invoke LLM to get PII entities"""
+
+        text = text.strip()
+        messages = [
+            {"role": "system", "content": self.system_prompt},
+            {"role": "user", "content": text},
+        ]
+
+        response = await self.client.chat.completions.create(
+            model=self.model,
+            messages=messages,
+            # The field guided_json is unsupported at the root level
+            # and must be included in the nvext object field
+            extra_body={"nvext": {"guided_json": JSON_SCHEMA}},
+            stream=False,
+            max_tokens=4096,
+        )
+
+        assistant_message = response.choices[0].message.content
+
+        # Parse results
+        try:
+            entities = json.loads(assistant_message)
+            if not entities:
+                # LLM returned valid JSON but no entities discovered
+                return []
+            else:
+                # Check that each entity returned is valid
+                return [e for e in entities if validate_entity(e, text)]
+        except json.decoder.JSONDecodeError:
+            return []
+
+
+class AsyncLLMPiiModifier(DocumentModifier):
+    """
+    This class is the entry point to using the LLM-based PII de-identification module.
+    It works with the `Modify` functionality as shown below:
+
+    dataframe = pd.DataFrame({"text": ["Sarah and Ryan went out to play", "Jensen is the CEO of NVIDIA"]})
+    dd = dask.dataframe.from_pandas(dataframe, npartitions=1)
+    dataset = DocumentDataset(dd)
+
+    modifier = AsyncLLMPiiModifier(
+        # Endpoint for the user's NIM
+        base_url="http://0.0.0.0:8000/v1",
+        api_key="API KEY (if needed)",
+        model="meta/llama-3.1-70b-instruct",
+        # The user may engineer a custom prompt if desired
+        system_prompt=SYSTEM_PROMPT,
+        pii_labels=PII_LABELS,
+        language="en",
+    )
+
+    modify = Modify(modifier)
+    modified_dataset = modify(dataset)
+    modified_dataset.df.to_json("output_files/*.jsonl", lines=True, orient="records")
+
+    """
+
+    def __init__(
+        self,
+        base_url: str,
+        api_key: Optional[str] = None,
+        model: str = "meta/llama-3.1-70b-instruct",
+        system_prompt: str = SYSTEM_PROMPT,
+        pii_labels: List[str] = PII_LABELS,
+        language: str = "en",
+        max_concurrent_requests: Optional[int] = None,
+    ):
+        """
+        Initialize the AsyncLLMPiiModifier
+
+        Args:
+            base_url (str): The base URL for the user's NIM
+            api_key (Optional[str]): The API key for the user's NIM, if needed.
+                Default is None.
+            model (str): The model to use for the LLM.
+                Default is "meta/llama-3.1-70b-instruct".
+            system_prompt (str): The system prompt to feed into the LLM.
+                Default prompt has been fine-tuned for "meta/llama-3.1-70b-instruct".
+            pii_labels (List[str]): The PII labels to identify and remove from the text.
+                See documentation for full list of PII labels.
+            language (str): The language to use for the LLM.
+                Default is "en" for English. If non-English, it is recommended
+                to provide a custom system prompt.
+            max_concurrent_requests (Optional[int]): The maximum number of concurrent requests to make to the LLM.
+                Default is None, which means no limit.
+
+        """
+        super().__init__()
+
+        self.base_url = base_url
+        self.api_key = api_key
+        self.model = model
+        self.system_prompt = system_prompt
+        self.pii_labels = pii_labels
+        self.language = language
+        self.max_concurrent_requests = max_concurrent_requests
+
+        if self.language != "en" and self.system_prompt is SYSTEM_PROMPT:
+            warnings.warn(
+                "The default system prompt is only available for English. "
+                "For other languages, please provide a custom system prompt."
+            )
+        if self.model not in SYSTEM_PROMPT:
+            warnings.warn(
+                f"No system prompt has been defined for model {model}. "
+                "Default system prompt will be used."
+            )
+            self.system_prompt[self.model] = SYSTEM_PROMPT[
+                "meta/llama-3.1-70b-instruct"
+            ]
+
+    @batched
+    def modify_document(self, text: pd.Series):
+        inferer = load_object_on_worker("inferer", self.load_inferer, {})
+        pii_entities_lists = asyncio.run(self.call_inferer(text, inferer))
+        text_redacted = self.batch_redact(text, pii_entities_lists)
+        return text_redacted
+
+    def load_inferer(self):
+        """Helper function to load the asynchronous LLM"""
+        inferer: AsyncLLMInference = AsyncLLMInference(
+            base_url=self.base_url,
+            api_key=self.api_key,
+            model=self.model,
+            system_prompt=self.system_prompt,
+            pii_labels=self.pii_labels,
+        )
+
+        return inferer
+
+    async def call_inferer(self, text: pd.Series, inferer: AsyncLLMInference):
+        tasks = [inferer.infer(prompt) for prompt in text]
+        pii_entities_lists = await self._gather(tasks)
+        return pii_entities_lists
+
+    async def _gather(
+        self, requests: List[Coroutine[Any, Any, List[str]]]
+    ) -> List[str]:
+        max_requests = self.max_concurrent_requests
+        if max_requests is None:
+            max_requests = len(requests)
+
+        final_list = []
+        for i in tqdm(range(0, len(requests), max_requests)):
+            request_slice = requests[i : i + max_requests]
+            result = await tqdm.gather(*request_slice)
+            final_list.extend(result)
+
+        return final_list
+
+    def batch_redact(
+        self, text: pd.Series, pii_entities_lists: List[List[Dict[str, str]]]
+    ):
+        redacted_texts = [
+            redact(text_str, pii_entities)
+            for text_str, pii_entities in zip(text, pii_entities_lists)
+        ]
+        return pd.Series(redacted_texts, index=text.index)
diff --git a/nemo_curator/modifiers/llm_pii_modifier.py b/nemo_curator/modifiers/llm_pii_modifier.py
new file mode 100644
index 000000000..bbcb3a3a1
--- /dev/null
+++ b/nemo_curator/modifiers/llm_pii_modifier.py
@@ -0,0 +1,176 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import warnings
+from typing import Dict, List, Optional
+
+from openai import OpenAI
+
+from nemo_curator.modifiers import DocumentModifier
+from nemo_curator.utils.distributed_utils import load_object_on_worker
+from nemo_curator.utils.llm_pii_utils import (
+    JSON_SCHEMA,
+    PII_LABELS,
+    SYSTEM_PROMPT,
+    redact,
+    validate_entity,
+)
+
+__all__ = ["LLMPiiModifier"]
+
+
+class LLMInference:
+    """A class for redacting PII via LLM inference"""
+
+    def __init__(
+        self,
+        base_url: str,
+        api_key: Optional[str] = None,
+        model: str = "meta/llama-3.1-70b-instruct",
+        system_prompt: str = SYSTEM_PROMPT,
+        pii_labels: List[str] = PII_LABELS,
+    ):
+        self.client = OpenAI(base_url=base_url, api_key=api_key)
+        self.model = model
+        self.system_prompt = system_prompt[model]
+        self.pii_labels = pii_labels
+
+    def infer(self, text: str) -> List[Dict[str, str]]:
+        """Invoke LLM to get PII entities"""
+
+        text = text.strip()
+        messages = [
+            {"role": "system", "content": self.system_prompt},
+            {"role": "user", "content": text},
+        ]
+
+        response = self.client.chat.completions.create(
+            model=self.model,
+            messages=messages,
+            # The field guided_json is unsupported at the root level
+            # and must be included in the nvext object field
+            extra_body={"nvext": {"guided_json": JSON_SCHEMA}},
+            stream=False,
+            max_tokens=4096,
+        )
+
+        assistant_message = response.choices[0].message.content
+
+        # Parse results
+        try:
+            entities = json.loads(assistant_message)
+            if not entities:
+                # LLM returned valid JSON but no entities discovered
+                return []
+            else:
+                # Check that each entity returned is valid
+                return [e for e in entities if validate_entity(e, text)]
+        except json.decoder.JSONDecodeError:
+            return []
+
+
+class LLMPiiModifier(DocumentModifier):
+    """
+    This class is the entry point to using the LLM-based PII de-identification module.
+    It works with the `Modify` functionality as shown below:
+
+    dataframe = pd.DataFrame({"text": ["Sarah and Ryan went out to play", "Jensen is the CEO of NVIDIA"]})
+    dd = dask.dataframe.from_pandas(dataframe, npartitions=1)
+    dataset = DocumentDataset(dd)
+
+    modifier = LLMPiiModifier(
+        # Endpoint for the user's NIM
+        base_url="http://0.0.0.0:8000/v1",
+        api_key="API KEY (if needed)",
+        model="meta/llama-3.1-70b-instruct",
+        # The user may engineer a custom prompt if desired
+        system_prompt=SYSTEM_PROMPT,
+        pii_labels=PII_LABELS,
+        language="en",
+    )
+
+    modify = Modify(modifier)
+    modified_dataset = modify(dataset)
+    modified_dataset.df.to_json("output_files/*.jsonl", lines=True, orient="records")
+
+    """
+
+    def __init__(
+        self,
+        base_url: str,
+        api_key: Optional[str] = None,
+        model: str = "meta/llama-3.1-70b-instruct",
+        system_prompt: str = SYSTEM_PROMPT,
+        pii_labels: List[str] = PII_LABELS,
+        language: str = "en",
+    ):
+        """
+        Initialize the LLMPiiModifier
+
+        Args:
+            base_url (str): The base URL for the user's NIM
+            api_key (Optional[str]): The API key for the user's NIM, if needed.
+                Default is None.
+            model (str): The model to use for the LLM.
+                Default is "meta/llama-3.1-70b-instruct".
+            system_prompt (str): The system prompt to feed into the LLM.
+                Default prompt has been fine-tuned for "meta/llama-3.1-70b-instruct".
+            pii_labels (List[str]): The PII labels to identify and remove from the text.
+                See documentation for full list of PII labels.
+            language (str): The language to use for the LLM.
+                Default is "en" for English. If non-English, it is recommended
+                to provide a custom system prompt.
+
+        """
+        super().__init__()
+
+        self.base_url = base_url
+        self.api_key = api_key
+        self.model = model
+        self.system_prompt = system_prompt
+        self.pii_labels = pii_labels
+        self.language = language
+
+        if self.language != "en" and self.system_prompt is SYSTEM_PROMPT:
+            warnings.warn(
+                "The default system prompt is only available for English. "
+                "For other languages, please provide a custom system prompt."
+            )
+        if self.model not in SYSTEM_PROMPT:
+            warnings.warn(
+                f"No system prompt has been defined for model {model}. "
+                "Default system prompt will be used."
+            )
+            self.system_prompt[self.model] = SYSTEM_PROMPT[
+                "meta/llama-3.1-70b-instruct"
+            ]
+
+    def modify_document(self, text: str):
+        inferer = load_object_on_worker("inferer", self.load_inferer, {})
+        pii_entities = inferer.infer(text)
+        text_redacted = redact(text, pii_entities)
+        return text_redacted
+
+    def load_inferer(self):
+        """Helper function to load the LLM"""
+        inferer: LLMInference = LLMInference(
+            base_url=self.base_url,
+            api_key=self.api_key,
+            model=self.model,
+            system_prompt=self.system_prompt,
+            pii_labels=self.pii_labels,
+        )
+
+        return inferer
diff --git a/nemo_curator/utils/llm_pii_utils.py b/nemo_curator/utils/llm_pii_utils.py
new file mode 100644
index 000000000..ae693b870
--- /dev/null
+++ b/nemo_curator/utils/llm_pii_utils.py
@@ -0,0 +1,222 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import re
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import DefaultDict, Dict, List, Tuple
+
+
+@dataclass
+class EntitySpan:
+    entity_type: str
+    start_position: int
+    end_position: int
+
+
+JSON_SCHEMA = {
+    "type": "array",
+    "items": {
+        "type": "object",
+        "required": ["entity_type", "entity_text"],
+        "properties": {
+            "entity_type": {"type": "string"},
+            "entity_text": {"type": "string"},
+        },
+    },
+}
+
+# Recommended PII entities to redact
+PII_LABELS = [
+    "medical_record_number",
+    "location",
+    "address",
+    "ssn",
+    "date_of_birth",
+    "date_time",
+    "name",
+    "email",
+    "customer_id",
+    "employee_id",
+    "phone_number",
+    "ip_address",
+    "credit_card_number",
+    "user_name",
+    "device_identifier",
+    "bank_routing_number",
+    "company_name",
+    "unique_identifier",
+    "biometric_identifier",
+    "account_number",
+    "certificate_license_number",
+    "license_plate",
+    "vehicle_identifier",
+    "api_key",
+    "password",
+    "health_plan_beneficiary_number",
+    "national_id",
+    "tax_id",
+    "url",
+    "swift_bic",
+    "cvv",
+    "pin",
+]
+
+SYSTEM_PROMPT = {
+    "meta/llama-3.1-70b-instruct": (
+        "You are an expert redactor. The user is going to provide you with "
+        "some text. Please find all personally identifying information from "
+        "this text. Return results according to this JSON schema: "
+        f"{str(JSON_SCHEMA)}\nOnly return results for entities which actually "
+        "appear in the text. It is very important that you return the  "
+        "entity_text by copying it exactly from the input. Do not perform any "
+        "modification or normalization of the text. The entity_type should be "
+        f"one of these: {', '.join(PII_LABELS)}\n"
+    ),
+}
+
+
+def validate_entity(
+    entity: Dict[str, str],
+    text: str,
+    min_length: int = 2,
+) -> bool:
+    """Validate entity"""
+
+    if not validate_keys(entity):
+        return False
+
+    if entity["entity_text"] not in text:
+        return False
+
+    if len(entity["entity_text"]) < min_length:
+        return False
+
+    return True
+
+
+def validate_keys(entity_dict: Dict) -> bool:
+    """Validate that keys in entity dict match schema"""
+
+    required_keys = JSON_SCHEMA["items"]["required"]
+
+    entity_keys = list(entity_dict.keys())
+    if len(entity_keys) != len(required_keys):
+        return False
+
+    for k in required_keys:
+        if k not in entity_keys:
+            return False
+
+    return True
+
+
+def redact(
+    full_text: str,
+    pii_entities: List[Dict[str, str]],
+) -> str:
+    """Redact given entities from the original text"""
+
+    entity_spans = find_entity_spans(full_text, pii_entities)
+
+    # Initialize an offset to track the changes in the text
+    redacted_text = full_text
+    offset = 0
+    for entity in fix_overlaps(entity_spans):
+        entity_type = entity.entity_type
+        start_position = entity.start_position + offset
+        end_position = entity.end_position + offset
+
+        # Replace the entity value with its type
+        replacement = "{{" + entity_type + "}}"
+        redacted_text = (
+            redacted_text[:start_position] + replacement + redacted_text[end_position:]
+        )
+
+        # Update the offset
+        offset += len(replacement) - (end_position - start_position)
+
+    return redacted_text
+
+
+def find_entity_spans(text: str, entities: List[Dict[str, str]]) -> List[EntitySpan]:
+    """
+    Find the start and end indexes for each entity in the given text.
+
+    Args:
+        text (str): The input text string.
+        entities (list): A list of entities, where each entity is a dictionary
+            containing the entity text and its type.
+
+    Returns:
+        list: A list of EntitySpan objects, where each contains the entity
+            type, start position, and end position.
+
+    """
+    result = []
+    seen: DefaultDict[Tuple, int] = defaultdict(int)
+    for entity in entities:
+        entity_text = entity["entity_text"]
+        entity_type = entity["entity_type"]
+        if (entity_text, entity_type) in seen:
+            continue
+
+        matches = re.finditer(re.escape(entity_text), text, re.IGNORECASE)
+        if matches:
+            result.extend(
+                [
+                    EntitySpan(entity_type, match.start(), match.end())
+                    for match in matches
+                ]
+            )
+            seen[(entity_text, entity_type)] += 1
+
+    return result
+
+
+def fix_overlaps(spans: List[EntitySpan]) -> List[EntitySpan]:
+    """Handle overlaps in entity spans"""
+
+    if not spans:
+        return spans
+
+    # Sort spans and iterate to check for overlap
+    sorted_spans = sorted(spans, key=lambda x: (x.start_position, -x.end_position))
+    results = [sorted_spans[0]]
+
+    for span in sorted_spans[1:]:
+        if span.start_position < results[-1].end_position:
+            # Skip candidate span if sub-interval (regardless of category)
+            if span.end_position < results[-1].end_position:
+                continue
+
+            else:
+                if span.entity_type == results[-1].entity_type:
+                    # If overlapping spans have same type, then extend boundary
+                    results[-1].end_position = span.end_position
+
+                else:
+                    # If overlapping spans have different type, add new span
+                    results.append(
+                        EntitySpan(
+                            span.entity_type,
+                            results[-1].end_position + 1,
+                            span.end_position,
+                        )
+                    )
+        else:
+            # Add non-overlapping span
+            results.append(span)
+
+    return results