开源 AI 食谱文档

用于 PII 检测的 LLM 网关

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

Open In Colab

用于 PII 检测的 LLM 网关

作者:Anthony Susevski

在企业应用场景中采用 LLM 时,一个常见的抱怨是数据隐私问题;特别是对于处理敏感数据的团队而言。虽然开源权重模型总是一个很好的选择,并且 *如果可能的话应该进行试用*,但有时我们只是想快速演示,或者有充分的理由使用 LLM API。在这些情况下,拥有一个可以处理个人可识别信息 (PII) 数据清洗的网关是一种很好的做法,以降低 PII 泄露的风险。

总部位于加拿大多伦多的金融科技公司 Wealthsimple 开源了一个代码仓库,正是为此目的而创建的。在本 notebook 中,我们将探讨如何利用这个仓库在调用 LLM 提供商的 API 之前清洗我们的数据。为此,我们将使用 来自 AI4Privacy 的 PII 数据集,并利用 Cohere 的 免费试用 API,以其 Command R+ 模型来演示 Wealthsimple 用于 PII 清洗的仓库。

首先,请按照 README 中的说明进行安装

  1. 安装 Poetry 和 Pyenv
  2. 安装 pyenv install 3.11.3
  3. 安装项目依赖
brew install gitleaks
poetry install
poetry run pre-commit install
  1. 运行 `cp .envrc.example .envrc` 并用 API 密钥更新
import os
from llm_gateway.providers.cohere import CohereWrapper
from datasets import load_dataset
import cohere
import types
import re
COHERE_API_KEY = os.environ["COHERE_API_KEY"]
DATABASE_URL = os.environ[
    "DATABASE_URL"
]  # default database url: "postgresql://postgres:postgres@postgres:5432/llm_gateway"

LLM 包装器 (Wrapper)

这个包装器对象是一个简单的包装器,它在进行 API 调用之前对提示应用“清洗器 (scrubbers)”。使用包装器发出请求后,我们会返回一个响应和一个 db_record 对象。在深入了解更多细节之前,让我们先看看它的实际效果。

wrapper = CohereWrapper()
example = "Michael Smith (msmith@gmail.com, (+1) 111-111-1111) committed a mistake when he used PyTorch Trainer instead of HF Trainer."
>>> response, db_record = wrapper.send_cohere_request(
...     endpoint="generate",
...     model="command-r-plus",
...     max_tokens=25,
...     prompt=f"{example}\n\nSummarize the above text in 1-2 sentences.",
...     temperature=0.3,
... )

>>> print(response)
{'data': ['Michael Smith made a mistake by using PyTorch Trainer instead of HF Trainer.'], 'return_likelihoods': None, 'meta': {'api_version': {'version': '1'}, 'billed_units': {'input_tokens': 48, 'output_tokens': 14}}}

响应返回了 LLM 的输出;在这种情况下,由于我们要求模型对一个已经很短的句子进行总结,它返回了以下消息

['Michael Smith made a mistake by using PyTorch Trainer instead of HF Trainer.'] (迈克尔·史密斯犯了一个错误,他使用了 PyTorch Trainer 而不是 HF Trainer。)

>>> print(db_record)
{'user_input': 'Michael Smith ([REDACTED EMAIL ADDRESS], (+1) [REDACTED PHONE NUMBER]) committed a mistake when he used PyTorch Trainer instead of HF Trainer.\n\nSummarize the above text in 1-2 sentences.', 'user_email': None, 'cohere_response': {'data': ['Michael Smith made a mistake by using PyTorch Trainer instead of HF Trainer.'], 'return_likelihoods': None, 'meta': {'api_version': {'version': '1'}, 'billed_units': {'input_tokens': 48, 'output_tokens': 14}}}, 'cohere_model': 'command-r-plus', 'temperature': 0.3, 'extras': '{}', 'created_at': datetime.datetime(2024, 6, 10, 2, 16, 7, 666438), 'cohere_endpoint': 'generate'}

返回的第二项是数据库记录。该仓库旨在与 postgres 后端一起使用;事实上,该仓库附带一个用 Docker 构建的完整前端。postgres 数据库用于存储网关的聊天历史记录。然而,它也非常有用,因为它向我们展示了每个请求中实际发送了哪些数据。正如我们所看到的,提示被清洗了,发送了以下内容

Michael Smith ([REDACTED EMAIL ADDRESS], (+1) [REDACTED PHONE NUMBER]) committed a mistake when he used PyTorch Trainer instead of HF Trainer.\n\nSummarize the above text in 1-2 sentences.

但是等等,我听到你在想。Michael Smith 不是 PII 吗?可能是。但这个仓库实际上没有实现名称清洗器。下面,我们将研究哪些清洗器被应用于提示

generate 端点实际上已被 Cohere 弃用,因此为 Cohere 新的 Chat 端点创建并提交一个集成将是一个非常棒的开源贡献。

清洗器 (Scrubbers)!

根据他们的仓库,他们实现了以下清洗器

ALL_SCRUBBERS = [
    scrub_phone_numbers,
    scrub_credit_card_numbers,
    scrub_email_addresses,
    scrub_postal_codes,
    scrub_sin_numbers,
]

网关将按顺序应用每个清洗器。

这有点取巧,但如果你真的需要实现另一个清洗器,你可以通过修改包装器调用清洗器的方法来做到这一点。下面我们将进行演示

作者提到,sin 清洗器特别容易清洗掉一些东西,所以他们最后应用它,以确保其他与数字相关的 PII 先被清洗掉

def my_custom_scrubber(text: str) -> str:
    """
    Scrub Michael Smith in text

    :param text: Input text to scrub
    :type text: str
    :return: Input text with any mentions of Michael Smith scrubbed
    :rtype: str
    """
    return re.sub(r"Michael Smith", "[REDACTED PERSON]", text, re.IGNORECASE)
original_method = wrapper.send_cohere_request


def modified_method(self, **kwargs):
    self._validate_cohere_endpoint(kwargs.get("endpoint", None))  # Unfortunate double validate cohere endpoint call
    prompt = kwargs.get("prompt", None)
    text = my_custom_scrubber(prompt)
    kwargs["prompt"] = text
    return original_method(**kwargs)


# Assign the new method to the instance
wrapper.send_cohere_request = types.MethodType(modified_method, wrapper)
>>> response, db_record = wrapper.send_cohere_request(
...     endpoint="generate",
...     model="command-r-plus",
...     max_tokens=25,
...     prompt=f"{example}\n\nSummarize the above text in 1-2 sentences.",
...     temperature=0.3,
... )

>>> print(response)
{'data': ['[REDACTED PERSON] made an error by using PyTorch Trainer instead of HF Trainer. They can be contacted at [RED'], 'return_likelihoods': None, 'meta': {'api_version': {'version': '1'}, 'billed_units': {'input_tokens': 52, 'output_tokens': 25}}}
>>> print(db_record)
{'user_input': '[REDACTED PERSON] ([REDACTED EMAIL ADDRESS], (+1) [REDACTED PHONE NUMBER]) committed a mistake when he used PyTorch Trainer instead of HF Trainer.\n\nSummarize the above text in 1-2 sentences.', 'user_email': None, 'cohere_response': {'data': ['[REDACTED PERSON] made an error by using PyTorch Trainer instead of HF Trainer. They can be contacted at [RED'], 'return_likelihoods': None, 'meta': {'api_version': {'version': '1'}, 'billed_units': {'input_tokens': 52, 'output_tokens': 25}}}, 'cohere_model': 'command-r-plus', 'temperature': 0.3, 'extras': '{}', 'created_at': datetime.datetime(2024, 6, 10, 2, 59, 58, 733195), 'cohere_endpoint': 'generate'}

如果你真的必须做这样的事情,请确保你记住清洗器是按顺序应用的,所以如果你的自定义清洗器与任何默认清洗器冲突,可能会出现一些奇怪的行为。

例如,特别是对于姓名,有 其他清洗库 你可以探索,它们采用更复杂的算法来清洗 PII。这个仓库涵盖了更多的 PII,例如 IP 地址、主机名等。然而,如果你只需要删除特定的匹配项,你可以回到上面的代码。

数据集

让我们在一个完整的数据集上探索这个包装器的实际应用。

pii_ds = load_dataset("ai4privacy/pii-masking-200k")
pii_ds["train"][36]["source_text"]
>>> example = pii_ds["train"][36]["source_text"]

>>> response, db_record = wrapper.send_cohere_request(
...     endpoint="generate",
...     model="command-r-plus",
...     max_tokens=50,
...     prompt=f"{example}\n\nSummarize the above text in 1-2 sentences.",
...     temperature=0.3,
... )

>>> print(response)
{'data': ["The person is requesting an update on assessment results and is offering Kip 100,000 in exchange for the information and the recipient's account details."], 'return_likelihoods': None, 'meta': {'api_version': {'version': '1'}, 'billed_units': {'input_tokens': 64, 'output_tokens': 33}}}
>>> print(db_record)
{'user_input': "I need the latest update on assessment results. Please send the files to V[REDACTED EMAIL ADDRESS]. For your extra time, we'll offer you Kip 100,000 but please provide your лв account details.\n\nSummarize the above text in 1-2 sentences.", 'user_email': None, 'cohere_response': {'data': ["The person is requesting an update on assessment results and is offering Kip 100,000 in exchange for the information and the recipient's account details."], 'return_likelihoods': None, 'meta': {'api_version': {'version': '1'}, 'billed_units': {'input_tokens': 64, 'output_tokens': 33}}}, 'cohere_model': 'command-r-plus', 'temperature': 0.3, 'extras': '{}', 'created_at': datetime.datetime(2024, 6, 10, 3, 10, 51, 416091), 'cohere_endpoint': 'generate'}

常规输出

如果我们简单地将文本原样发送到端点,摘要会是这个样子

 co = cohere.Client(
    api_key=os.environ['COHERE_API_KEY']


 esponse_vanilla = co.generate(
    prompt=f"{example}\n\nSummarize the above text in 1-2 sentences.",
    model="command-r-plus",
    max_tokens=50,
    temperature=0.3
response_vanilla

总结一下,在本 notebook 中,我们演示了如何使用一个由 Wealthsimple 慷慨开源的 PII 检测网关示例,并通过添加自定义清洗器对其进行了扩展。如果你确实需要可靠的 PII 检测,请确保运行你自己的测试,以验证你所使用的任何清洗算法是否真正覆盖了你的用例。最重要的是,在可能的情况下,将开源模型部署在你自己托管的基础设施上,永远是构建 LLM 应用最安全、最可靠的选择 :)

< > 在 GitHub 上更新