Detecting PII in LLM-Powered Chat Applications¶

by Grayson Adkins, updated April 17, 2024

This notebook demonstrates how to evaluate user prompts and LLM responses for personally identifiable infromation (PII) such as contact information, financial or banking info, digital identifiers, job related data, or other sensitive personal information.

We use the bigcode/starpii PII detection model available on Hugging Face plus LangChain for crafting prompt templates and TruLens for running evaluation and visualizing results.

Open In Colab

Attribution¶

This notebook builds on examples provided by TruLens.

Disclaimer¶

Both TruLens and LangChain are new frameworks with rapidly changing interfaces. I found several deprecated or broken features that I had to resolve while working on this notebook. Be advised that you may similarly find issues with the code here, due to those dependencies.

Install dependencies¶

In [13]:
!pip install -qU trulens_eval langchain

WARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv
In [21]:
!pip install -qU langchain_openai

   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 309.7/309.7 kB 9.1 MB/s eta 0:00:00
WARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv
In [15]:
import os
from dotenv import load_dotenv,find_dotenv

# # Load OPENAI_API_KEY from local .env file
# load_dotenv(find_dotenv())

# Or set it like this
os.environ["OPENAI_API_KEY"] = "sk-..."

## Print key to check
# print(os.environ["OPENAI_API_KEY"])

Set up TruLens¶

In [ ]:
from trulens_eval import Feedback
from trulens_eval import OpenAI as trulens_provider_openai
from trulens_eval import Tru

tru = Tru()
tru.reset_database()
In [22]:
# Imports from langchain to build app. You may need to install langchain first
# with the following:
# ! pip install langchain>=0.0.170
from langchain.chains import LLMChain
from langchain_openai import OpenAI
from langchain.prompts import PromptTemplate
from langchain.prompts.chat import HumanMessagePromptTemplate, ChatPromptTemplate
In [23]:
full_prompt = HumanMessagePromptTemplate(
    prompt=PromptTemplate(
        template=
        "Provide a helpful response with relevant background information for the following: {prompt}",
        input_variables=["prompt"],
    )
)

chat_prompt_template = ChatPromptTemplate.from_messages([full_prompt])

llm = OpenAI(temperature=0.9, max_tokens=128)

chain = LLMChain(llm=llm, prompt=chat_prompt_template, verbose=True)
In [24]:
prompt_input = 'Sam Altman is the CEO at OpenAI, and uses the password: password1234 .'

TruLens feedback functions¶

This feedback function includes chain of thought reasoning.

In [27]:
from trulens_eval.feedback.provider.hugs import Huggingface

# Hugging Face based feedback function collection class
hf_provider = Huggingface()
In [ ]:
# Define a pii_detection feedback function using HuggingFace.
# By default this will check language match on the main app input
f_pii_detection = Feedback(hf_provider.pii_detection_with_cot_reasons).on_input()

Create TruLens recorder¶

In [ ]:
tru_recorder = TruChain(chain,
    app_id='Chain1_ChatApplication',
    feedbacks=[f_pii_detection])

Execute eval¶

In [ ]:
with tru_recorder as recording:
    llm_response = chain(prompt_input)

display(llm_response)

Display results¶

In [ ]:
records, feedback = tru.get_records_and_feedback(app_ids=[])

# Make it a little easier to read
import pandas as pd

pd.set_option("display.max_colwidth", None)
records[["input", "output"] + feedback]
In [ ]:
tru.run_dashboard() # open a local streamlit app to explore

# tru.stop_dashboard() # stop if needed