用于PII检测的LLM网关

作者：Anthony Susevski

在企业用例中采用LLM时，一个常见的抱怨是关于数据隐私的，特别是对于处理敏感数据的团队而言。虽然开放权重模型始终是一个不错的选择，并且*如果可能的话应该进行试验*，但有时我们只想快速演示一些东西，或者有充分的理由使用LLM API。在这些情况下，最好有一个网关来处理个人身份信息（PII）数据的清理，以降低PII泄露的风险。

总部位于加拿大多伦多的金融科技公司Wealthsimple，开源了一个仓库，它就是为此目的而创建的。在本笔记本中，我们将探索如何利用这个仓库在向LLM提供者发出API调用之前清理我们的数据。为此，我们将查看来自AI4Privacy的PII数据集，并利用Cohere的免费试用API来展示Wealthsimple仓库用于PII清理的Command R+模型。

首先，按照自述文件中的说明进行安装

1. 安装Poetry和Pyenv
2. 安装pyenv install 3.11.3
3. 安装项目要求

brew install gitleaks
poetry install
poetry run pre-commit install

4. 运行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包装器

包装器对象是一个简单的包装器，它在发出API调用之前将“清理器”应用于提示。使用包装器发出请求后，我们会收到一个响应和一个db_record对象。在我们深入了解更多细节之前，让我们看看它的实际操作。

wrapper = CohereWrapper()

example = "Michael Smith ([email protected], (+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使用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数据库用于存储网关的聊天历史记录。但是，它也非常有用，因为它向我们展示了每次请求中实际发送了哪些数据。正如我们所见，提示被清理了，并发送了以下内容

迈克尔·史密斯（[REDACTED EMAIL ADDRESS]，(+1) [REDACTED PHONE NUMBER]）在使用 PyTorch Trainer 而不是 HF Trainer 时犯了一个错误。\n\n请用 1-2 句话概括以上文字。

等等，我听到你在想，迈克尔·史密斯的个人信息是不是很敏感？可能是的。但这个代码库实际上没有实现姓名清除器。接下来，我们将调查对提示应用了哪些清除器。

清除器!

来自他们代码库的实现的清除器。

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

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

这相当 hacky，但如果你确实需要实现另一个清除器，你可以通过修改调用清除器的包装器方法来实现。下面我们将演示。

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

总结一下，在这个笔记本当中，我们演示了如何使用 Wealthsimple 开源的 PII 检测示例网关，并通过添加自定义清除器来扩展它。如果你确实需要可靠的 PII 检测，请确保运行你自己的测试以验证你使用的任何清除算法是否真正涵盖了你的用例。最重要的是，只要有可能，将开源模型部署在你自己的基础设施上将始终是使用 LLM 进行构建最安全、最可靠的选择 :)