用於 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

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 包裝器 (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 嗎？可能是。但這個倉庫實際上沒有實現名稱清洗器。下面，我們將研究哪些清洗器被應用於提示

清洗器 (Scrubbers)！

根據他們的倉庫，他們實現了以下清洗器

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

閘道器將按順序應用每個清洗器。

這有點取巧，但如果你真的需要實現另一個清洗器，你可以透過修改包裝器呼叫清洗器的方法來做到這一點。下面我們將進行演示

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 應用最安全、最可靠的選擇 :)