Construir RAG con Milvus + PII Masker

PII (Personally Identifiable Information) es un tipo de datos sensibles que pueden utilizarse para identificar a individuos.

PII Masker, desarrollado por HydroX AI, es una herramienta avanzada de código abierto diseñada para proteger sus datos sensibles aprovechando modelos de IA de vanguardia. Tanto si maneja datos de clientes como si realiza análisis de datos o garantiza el cumplimiento de las normativas de privacidad, PII Masker proporciona una solución sólida y escalable para mantener su información segura.

En este tutorial, mostraremos cómo utilizar PII Masker con Milvus para proteger datos privados en aplicaciones RAG(Retrieval-Augmented Generation). Al combinar las fortalezas de las capacidades de enmascaramiento de datos de PII Masker con la eficiente recuperación de datos de Milvus, puede crear canalizaciones seguras y conformes con la privacidad para manejar información sensible con confianza. Este enfoque garantiza que sus aplicaciones estén equipadas para cumplir las normas de privacidad y proteger los datos de los usuarios de forma eficaz.

Preparación

Introducción a PII Masker

Siga la guía de instalación de PII Masker para instalar las dependencias necesarias y descargar el modelo. He aquí una guía sencilla:

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

Descargue el modelo dehttps://huggingface.co/hydroxai/pii_model_weight, y sustitúyalo por los archivos en: pii-masker/output_model/deberta3base_1024/

Dependencias y Entorno

$ pip install --upgrade pymilvus openai requests tqdm dataset

Usaremos OpenAI como LLM en este ejemplo. Debes preparar la clave api OPENAI_API_KEY como variable de entorno.

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

Luego puedes crear un notebook python o jupyter para ejecutar el siguiente código.

Preparar los datos

Vamos a generar algunas líneas falsas que contengan información PII con fines de prueba o demostración.

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.",
]

Enmascarar los datos con PIIMasker

Inicializamos el objeto PIIMasker y cargamos el modelo.

from model import PIIMasker

masker = PIIMasker()

A continuación, enmascaramos PII de una lista de líneas de texto e imprimimos los resultados enmascarados.

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

Preparar el modelo de incrustación

Inicializamos el cliente OpenAI para preparar el modelo de incrustación.

from openai import OpenAI

openai_client = OpenAI()

Definimos una función para generar incrustaciones de texto utilizando el cliente OpenAI. Usamos el modelo text-embedding-3-small como ejemplo.

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

Generamos una incrustación de prueba e imprimimos su dimensión y sus primeros elementos.

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]

Cargar los datos en Milvus

Crear la colección

from pymilvus import MilvusClient

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

Compruebe si la colección ya existe y elimínela en caso afirmativo.

collection_name = "my_rag_collection"

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

Crear una nueva colección con los parámetros especificados.

Si no especificamos ninguna información de campo, Milvus creará automáticamente un campo id por defecto para la clave primaria, y un campo vector para almacenar los datos vectoriales. Se utiliza un campo JSON reservado para almacenar campos no definidos por el esquema y sus valores.

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

Insertar datos

Itere a través de las líneas de texto enmascaradas, cree incrustaciones y luego inserte los datos en Milvus.

Aquí hay un nuevo campo text, que es un campo no definido en el esquema de la colección. Se añadirá automáticamente al campo dinámico JSON reservado, que puede tratarse como un campo normal a alto nivel.

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}

Construir RAG

Recuperar datos para una consulta

Especifiquemos una pregunta sobre los documentos.

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

Busquemos la pregunta en la colección y recuperemos la coincidencia semántica top-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
)

Veamos los resultados de la consulta

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
    ]
]

Utilizar LLM para obtener una respuesta RAG

Convertir los documentos recuperados a un formato de cadena.

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

Definir los mensajes de sistema y de usuario para el modelo de lenguaje.

Nota: Le decimos a LLM que si no hay información útil en los fragmentos, simplemente diga "No lo sé".

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>
"""

Usar OpenAI ChatGPT para generar una respuesta basada en los prompts.

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.

Aquí podemos ver, como hemos reemplazado la PII por máscaras, el LLM no puede obtener la información PII en contexto. De esta forma, podemos proteger de forma efectiva la privacidad de los usuarios.