{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "8f74e1c4",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/node_postprocessor/PII.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "c04ffe8e-6573-470f-aef5-348522a0de15",
   "metadata": {},
   "source": [
    "# PII Masking"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "94cf2040",
   "metadata": {},
   "source": [
    "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83660e9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-llms-openai\n",
    "%pip install llama-index-llms-huggingface"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "254fff73",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install llama-index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "efa2a242-27bc-478f-8939-18a7f8153d4f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:numexpr.utils:Note: NumExpr detected 16 cores but \"NUMEXPR_MAX_THREADS\" not set, so enforcing safe limit of 8.\n",
      "Note: NumExpr detected 16 cores but \"NUMEXPR_MAX_THREADS\" not set, so enforcing safe limit of 8.\n",
      "INFO:numexpr.utils:NumExpr defaulting to 8 threads.\n",
      "NumExpr defaulting to 8 threads.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/loganm/miniconda3/envs/llama-index/lib/python3.11/site-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 logging\n",
    "import sys\n",
    "\n",
    "logging.basicConfig(stream=sys.stdout, level=logging.INFO)\n",
    "logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))\n",
    "\n",
    "from llama_index.core.postprocessor import (\n",
    "    PIINodePostprocessor,\n",
    "    NERPIINodePostprocessor,\n",
    ")\n",
    "from llama_index.llms.huggingface import HuggingFaceLLM\n",
    "from llama_index.core import Document, VectorStoreIndex\n",
    "from llama_index.core.schema import TextNode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "216e951a-42c4-4e6b-b16d-6a6064829ebf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load documents\n",
    "text = \"\"\"\n",
    "Hello Paulo Santos. The latest statement for your credit card account \\\n",
    "1111-0000-1111-0000 was mailed to 123 Any Street, Seattle, WA 98109.\n",
    "\"\"\"\n",
    "node = TextNode(text=text)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "24495d69-d568-4cc7-9445-87692bf77863",
   "metadata": {},
   "source": [
    "### Option 1: Use NER Model for PII Masking\n",
    "\n",
    "Use a Hugging Face NER model for PII Masking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "003f66f0-f67f-47f2-88eb-b2bbb6d33791",
   "metadata": {},
   "outputs": [],
   "source": [
    "processor = NERPIINodePostprocessor()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e76c995c-57ee-4d1b-a771-6626ef93e8cd",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "No model was supplied, defaulted to dbmdz/bert-large-cased-finetuned-conll03-english and revision f2482bf (https://huggingface.co/dbmdz/bert-large-cased-finetuned-conll03-english).\n",
      "Using a pipeline without specifying a model name and revision in production is not recommended.\n",
      "/home/loganm/miniconda3/envs/llama-index/lib/python3.11/site-packages/transformers/pipelines/token_classification.py:169: UserWarning: `grouped_entities` is deprecated and will be removed in version v5.0.0, defaulted to `aggregation_strategy=\"AggregationStrategy.SIMPLE\"` instead.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "from llama_index.core.schema import NodeWithScore\n",
    "\n",
    "new_nodes = processor.postprocess_nodes([NodeWithScore(node=node)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4783c27-9a55-44f1-be9e-a4fe1fc1e0fa",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Hello [ORG_6]. The latest statement for your credit card account 1111-0000-1111-0000 was mailed to 123 [ORG_108] [LOC_112], [LOC_120], [LOC_129] 98109.'"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# view redacted text\n",
    "new_nodes[0].node.get_text()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "075d45dc-a226-4ba7-8c8a-d9dd536f8560",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'[ORG_6]': 'Paulo Santos',\n",
       " '[ORG_108]': 'Any',\n",
       " '[LOC_112]': 'Street',\n",
       " '[LOC_120]': 'Seattle',\n",
       " '[LOC_129]': 'WA'}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# get mapping in metadata\n",
    "# NOTE: this is not sent to the LLM!\n",
    "new_nodes[0].node.metadata[\"__pii_node_info__\"]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "06ca1e50-eeee-4079-bec6-3621cb760f98",
   "metadata": {},
   "source": [
    "### Option 2: Use LLM for PII Masking\n",
    "\n",
    "NOTE: You should be using a *local* LLM model for PII masking. The example shown is using OpenAI, but normally you'd use an LLM running locally, possibly from huggingface. Examples for local LLMs are [here](https://gpt-index.readthedocs.io/en/latest/how_to/customization/custom_llms.html#example-using-a-huggingface-llm)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a2db8d3-6bb7-4855-852e-a4941abb03bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.openai import OpenAI\n",
    "\n",
    "processor = PIINodePostprocessor(llm=OpenAI())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b834e7a3-8f90-45eb-841a-335b0d33dcab",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.schema import NodeWithScore\n",
    "\n",
    "new_nodes = processor.postprocess_nodes([NodeWithScore(node=node)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca1498f3-34a1-4001-90f9-03feb5532d7d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Hello [NAME]. The latest statement for your credit card account [CREDIT_CARD_NUMBER] was mailed to [ADDRESS].'"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# view redacted text\n",
    "new_nodes[0].node.get_text()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d574d591-c1db-498b-ba32-9ed4190c6b4c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'NAME': 'Paulo Santos',\n",
       " 'CREDIT_CARD_NUMBER': '1111-0000-1111-0000',\n",
       " 'ADDRESS': '123 Any Street, Seattle, WA 98109'}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# get mapping in metadata\n",
    "# NOTE: this is not sent to the LLM!\n",
    "new_nodes[0].node.metadata[\"__pii_node_info__\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6cc87ed4",
   "metadata": {},
   "source": [
    "### Option 3: Use Presidio for PII Masking\n",
    "\n",
    "Use presidio to identify and anonymize PII"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac215117",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load documents\n",
    "text = \"\"\"\n",
    "Hello Paulo Santos. The latest statement for your credit card account \\\n",
    "4095-2609-9393-4932 was mailed to Seattle, WA 98109. \\\n",
    "IBAN GB90YNTU67299444055881 and social security number is 474-49-7577 were verified on the system. \\\n",
    "Further communications will be sent to paulo@presidio.site \n",
    "\"\"\"\n",
    "presidio_node = TextNode(text=text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a745520",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.postprocessor.presidio import PresidioPIINodePostprocessor\n",
    "\n",
    "processor = PresidioPIINodePostprocessor()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89cb17ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.schema import NodeWithScore\n",
    "\n",
    "presidio_new_nodes = processor.postprocess_nodes(\n",
    "    [NodeWithScore(node=presidio_node)]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8fe9cef",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'\\nHello <PERSON_1>. The latest statement for your credit card account <CREDIT_CARD_1> was mailed to <LOCATION_2>, <LOCATION_1>. IBAN <IBAN_CODE_1> and social security number is <US_SSN_1> were verified on the system. Further communications will be sent to <EMAIL_ADDRESS_1> \\n'"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# view redacted text\n",
    "presidio_new_nodes[0].node.get_text()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80203af0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'<EMAIL_ADDRESS_1>': 'paulo@presidio.site',\n",
       " '<US_SSN_1>': '474-49-7577',\n",
       " '<IBAN_CODE_1>': 'GB90YNTU67299444055881',\n",
       " '<LOCATION_1>': 'WA 98109',\n",
       " '<LOCATION_2>': 'Seattle',\n",
       " '<CREDIT_CARD_1>': '4095-2609-9393-4932',\n",
       " '<PERSON_1>': 'Paulo Santos'}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# get mapping in metadata\n",
    "# NOTE: this is not sent to the LLM!\n",
    "presidio_new_nodes[0].node.metadata[\"__pii_node_info__\"]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "3444d895-e2fd-4af9-834a-64acf49f74f8",
   "metadata": {},
   "source": [
    "### Feed Nodes to Index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d33a9c0-efcd-4e79-b1f5-05aca9fc109f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:llama_index.token_counter.token_counter:> [build_index_from_nodes] Total LLM token usage: 0 tokens\n",
      "> [build_index_from_nodes] Total LLM token usage: 0 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [build_index_from_nodes] Total embedding token usage: 30 tokens\n",
      "> [build_index_from_nodes] Total embedding token usage: 30 tokens\n"
     ]
    }
   ],
   "source": [
    "# feed into index\n",
    "index = VectorStoreIndex([n.node for n in new_nodes])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc8b1993-d23b-4db1-8bb9-4f882bded66c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:llama_index.token_counter.token_counter:> [retrieve] Total LLM token usage: 0 tokens\n",
      "> [retrieve] Total LLM token usage: 0 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [retrieve] Total embedding token usage: 8 tokens\n",
      "> [retrieve] Total embedding token usage: 8 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [get_response] Total LLM token usage: 71 tokens\n",
      "> [get_response] Total LLM token usage: 71 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [get_response] Total embedding token usage: 0 tokens\n",
      "> [get_response] Total embedding token usage: 0 tokens\n",
      "\n",
      "[ADDRESS]\n"
     ]
    }
   ],
   "source": [
    "response = index.as_query_engine().query(\n",
    "    \"What address was the statement mailed to?\"\n",
    ")\n",
    "print(str(response))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama-index",
   "language": "python",
   "name": "llama-index"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}