Milvus + PII MaskerでRAGを構築する

PII（Personally Identifiable Information）とは、個人を特定するために使用できる機密データの一種です。

HydroX AIが開発したPII Maskerは、最先端のAIモデルを活用することで、機密データを保護するように設計された先進的なオープンソースツールです。PII Masker は、顧客データの取り扱い、データ分析の実行、プライバシー規制の遵守のいずれにおいても、情報を安全に保護するための堅牢で拡張可能なソリューションを提供します。

このチュートリアルでは、RAG(Retrieval-Augmented Generation)アプリケーションでプライベートデータを保護するためにPII MaskerとMilvusを使用する方法を紹介します。PII Maskerのデータマスキング機能とMilvusの効率的なデータ検索機能を組み合わせることで、機密情報を安心して取り扱うためのセキュアでプライバシーに準拠したパイプラインを作成することができます。このアプローチにより、お客様のアプリケーションはプライバシー基準を満たし、ユーザーデータを効果的に保護することができます。

準備

PII Masker を使い始める

PII Masker のインストールガイドに従って、必要な依存関係をインストールし、モデルをダウンロードします。以下は簡単なガイドです：

$ git clone https://github.com/HydroXai/pii-masker-v1.git
$ cd pii-masker-v1/pii-masker

https://huggingface.co/hydroxai/pii_model_weight からモデルをダウンロードし、以下のファイルと置き換えます：pii-masker/output_model/deberta3base_1024/

依存関係と環境

$ pip install --upgrade pymilvus openai requests tqdm dataset

この例ではLLMとしてOpenAIを使用します。apiキー OPENAI_API_KEY を環境変数として用意してください。

$ export OPENAI_API_KEY=sk-***********

そして、pythonかjupyterのノートブックを作成し、以下のコードを実行します。

データの準備

テストやデモンストレーションのために、PII情報を含む偽の行をいくつか生成してみましょう。

text_lines = [
    "Alice Johnson, a resident of Dublin, Ireland, attended a flower festival at Hyde Park on May 15, 2023. She entered the park at noon using her digital passport, number 23456789. Alice spent the afternoon admiring various flowers and plants, attending a gardening workshop, and having a light snack at one of the food stalls. While there, she met another visitor, Mr. Thompson, who was visiting from London. They exchanged tips on gardening and shared contact information: Mr. Thompson's address was 492, Pine Lane, and his cell phone number was +018.221.431-4517. Alice gave her contact details: home address, Ranch 16",
    "Hiroshi Tanaka, a businessman from Tokyo, Japan, went to attend a tech expo at the Berlin Convention Center on November 10, 2023. He registered for the event at 9 AM using his digital passport, number Q-24567680. Hiroshi networked with industry professionals, participated in panel discussions, and had lunch with some potential partners. One of the partners he met was from Munich, and they decided to keep in touch: the partner's office address was given as house No. 12, Road 7, Block E. Hiroshi offered his business card with the address, 654 Sakura Road, Tokyo.",
    "In an online forum discussion about culinary exchanges around the world, several participants shared their experiences. One user, Male, with the email 2022johndoe@example.com, shared his insights. He mentioned his ID code 1A2B3C4D5E and reference number L87654321 while residing in Italy but originally from Australia. He provided his +0-777-123-4567 and described his address at 456, Flavorful Lane, Pasta, IT, 00100.",
    "Another user joined the conversation on the topic of international volunteering opportunities. Identified as Female, she used the email 2023janedoe@example.com to share her story. She noted her 9876543210123 and M1234567890123 while residing in Germany but originally from Brazil. She provided her +0-333-987-6543 and described her address at 789, Sunny Side Street, Berlin, DE, 10178.",
]

PIIMasker でデータをマスクする

PIIMasker オブジェクトを初期化し、モデルを読み込みます。

from model import PIIMasker

masker = PIIMasker()

そして、テキスト行のリストからPIIをマスクし、マスクされた結果を表示します。

masked_results = []
for full_text in text_lines:
    masked_text, _ = masker.mask_pii(full_text)
    masked_results.append(masked_text)

for res in masked_results:
    print(res + "\n")

Alice [B-NAME] , a resident of Dublin Ireland attended flower festival at Hyde Park on May 15 2023 [B-PHONE_NUM] She entered the park noon using her digital passport number 23 [B-ID_NUM] [B-NAME] afternoon admiring various flowers and plants attending gardening workshop having light snack one food stalls While there she met another visitor Mr Thompson who was visiting from London They exchanged tips shared contact information : ' s address 492 [I-STREET_ADDRESS] his cell phone + [B-PHONE_NUM] [B-NAME] details home Ranch [B-STREET_ADDRESS]

Hiroshi [B-NAME] [I-STREET_ADDRESS] a businessman from Tokyo Japan went to attend tech expo at the Berlin Convention Center on November 10 2023 . He registered for event 9 AM using his digital passport number Q [B-ID_NUM] [B-NAME] with industry professionals participated in panel discussions and had lunch some potential partners One of he met was Munich they decided keep touch : partner ' s office address given as house No [I-STREET_ADDRESS] [B-NAME] business card 654 [B-STREET_ADDRESS]

In an online forum discussion about culinary exchanges around the world [I-STREET_ADDRESS] several participants shared their experiences [I-STREET_ADDRESS] One user Male with email 2022 [B-EMAIL] his insights He mentioned ID code 1 [B-ID_NUM] [I-PHONE_NUM] reference number L [B-ID_NUM] residing in Italy but originally from Australia provided + [B-PHONE_NUM] [I-PHONE_NUM] described address at 456 [I-STREET_ADDRESS]

Another user joined the conversation on topic of international volunteering opportunities . Identified as Female , she used email 2023 [B-EMAIL] share her story She noted 98 [B-ID_NUM] [I-PHONE_NUM] M [B-ID_NUM] residing in Germany but originally from Brazil provided + [B-PHONE_NUM] [I-PHONE_NUM] described address at 789 [I-STREET_ADDRESS] DE 10 178

埋め込みモデルの準備

埋め込みモデルを準備するために、OpenAIクライアントを初期化します。

from openai import OpenAI

openai_client = OpenAI()

OpenAIクライアントを使って、テキスト埋め込みを生成する関数を定義します。例として、text-embedding-3-small のモデルを使います。

def emb_text(text):
    return (
        openai_client.embeddings.create(input=text, model="text-embedding-3-small")
        .data[0]
        .embedding
    )

テスト埋め込みを生成し、その次元と最初の数要素を表示します。

test_embedding = emb_text("This is a test")
embedding_dim = len(test_embedding)
print(embedding_dim)
print(test_embedding[:10])

1536
[0.009889289736747742, -0.005578675772994757, 0.00683477520942688, -0.03805781528353691, -0.01824733428657055, -0.04121600463986397, -0.007636285852640867, 0.03225184231996536, 0.018949154764413834, 9.352207416668534e-05]

Milvusにデータをロードする。

コレクションの作成

from pymilvus import MilvusClient

milvus_client = MilvusClient(uri="./milvus_demo.db")

コレクションが既に存在するか確認し、存在する場合は削除します。

collection_name = "my_rag_collection"

if milvus_client.has_collection(collection_name):
    milvus_client.drop_collection(collection_name)

指定したパラメータで新しいコレクションを作成します。

フィールド情報を指定しない場合、Milvusは自動的にプライマリキー用のデフォルトid フィールドと、ベクトルデータを格納するためのvector フィールドを作成します。予約されたJSONフィールドは、スキーマで定義されていないフィールドとその値を格納するために使用されます。

milvus_client.create_collection(
    collection_name=collection_name,
    dimension=embedding_dim,
    metric_type="IP",  # Inner product distance
    consistency_level="Strong",  # Strong consistency level
)

データの挿入

マスクされたテキスト行を繰り返し、エンベッディングを作成し、Milvusにデータを挿入します。

ここに新しいフィールドtext 、コレクションスキーマの非定義フィールドです。これは予約されたJSONダイナミックフィールドに自動的に追加され、高レベルでは通常のフィールドとして扱われる。

from tqdm import tqdm

data = []

for i, line in enumerate(tqdm(masked_results, desc="Creating embeddings")):
    data.append({"id": i, "vector": emb_text(line), "text": line})

milvus_client.insert(collection_name=collection_name, data=data)

Creating embeddings: 100%|██████████| 4/4 [00:01<00:00,  2.60it/s]





{'insert_count': 4, 'ids': [0, 1, 2, 3], 'cost': 0}

RAGの構築

クエリのデータを取得する

ドキュメントに関する質問を指定してみましょう。

question = "What was the office address of Hiroshi's partner from Munich?"

コレクションで質問を検索し、セマンティックトップ1マッチを取得する。

search_res = milvus_client.search(
    collection_name=collection_name,
    data=[
        emb_text(question)
    ],  # Use the `emb_text` function to convert the question to an embedding vector
    limit=1,  # Return top 1 results
    search_params={"metric_type": "IP", "params": {}},  # Inner product distance
    output_fields=["text"],  # Return the text field
)

クエリの検索結果を見てみよう

import json

retrieved_lines_with_distances = [
    (res["entity"]["text"], res["distance"]) for res in search_res[0]
]
print(json.dumps(retrieved_lines_with_distances, indent=4))

[
    [
        "Hiroshi [B-NAME] [I-STREET_ADDRESS] a businessman from Tokyo Japan went to attend tech expo at the Berlin Convention Center on November 10 2023 . He registered for event 9 AM using his digital passport number Q [B-ID_NUM] [B-NAME] with industry professionals participated in panel discussions and had lunch some potential partners One of he met was Munich they decided keep touch : partner ' s office address given as house No [I-STREET_ADDRESS] [B-NAME] business card 654 [B-STREET_ADDRESS]",
        0.6544462442398071
    ]
]

LLMを使ってRAGレスポンスを取得する

検索されたドキュメントを文字列形式に変換する。

context = "\n".join(
    [line_with_distance[0] for line_with_distance in retrieved_lines_with_distances]
)

LLMのシステムプロンプトとユーザープロンプトを定義する。

注：スニペットの中に有用な情報がない場合は、"I don't know "と言ってくださいとLLMに伝えます。

SYSTEM_PROMPT = """
Human: You are an AI assistant. You are able to find answers to the questions from the contextual passage snippets provided. If there are no useful information in the snippets, just say "I don't know".
AI:
"""
USER_PROMPT = f"""
Use the following pieces of information enclosed in <context> tags to provide an answer to the question enclosed in <question> tags.
<context>
{context}
</context>
<question>
{question}
</question>
"""

OpenAI ChatGPTを使って、プロンプトに基づいた応答を生成します。

response = openai_client.chat.completions.create(
    model="gpt-4o-mini",
    messages=[
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": USER_PROMPT},
    ],
)
print(response.choices[0].message.content)

I don't know.

ここでは、PIIをマスクに置き換えているため、LLMはPIIの情報を文脈から取得できないことがわかります。このようにして、ユーザーのプライバシーを効果的に保護することができます。