Spaces:
Sleeping
Sleeping
File size: 3,860 Bytes
cd3f41e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 |
from gliner import GLiNER
print("Loading model and tokenizer...")
model = GLiNER.from_pretrained("urchade/gliner_multi_pii-v1")
print("Model and tokenizer loaded.")
label_to_mask_map = {
"name": "[NAME]",
"nric": "[NRIC]",
"phone number": "[PHONE]",
"address": "[ADDRESS]",
"email": "[EMAIL]",
"person": "[PERSON]",
"organization": "[ORGANIZATION]",
"passport number": "[PASSPORT_NUMBER]",
"credit card number": "[CREDIT_CARD]",
"social security number": "[SSN]",
"health insurance id number": "[HEALTH_INS_ID]",
"date of birth": "[DOB]",
"mobile phone number": "[MOBILE_PHONE]",
"bank account number": "[BANK_ACCOUNT]",
"medication": "[MEDICATION]",
"cpf": "[CPF]",
"driver's license number": "[DRIVER_LICENSE]",
"tax identification number": "[TAX_ID]",
"medical condition": "[MEDICAL_CONDITION]",
"identity card number": "[IDENTITY_CARD]",
"national id number": "[NATIONAL_ID]",
"ip address": "[IP]",
"email address": "[EMAIL]",
"iban": "[IBAN]",
"credit card expiration date": "[CREDIT_CARD_EXP]",
"username": "[USERNAME]",
"health insurance number": "[HEALTH_INS_NUM]",
"registration number": "[REG_NUM]",
"student id number": "[STUDENT_ID]",
"insurance number": "[INSURANCE_NUM]",
"flight number": "[FLIGHT_NUM]",
"landline phone number": "[LANDLINE_PHONE]",
"blood type": "[BLOOD_TYPE]",
"cvv": "[CVV]",
"reservation number": "[RESERVATION_NUM]",
"digital signature": "[DIGITAL_SIGNATURE]",
"social media handle": "[SOCIAL_MEDIA]",
"license plate number": "[LICENSE_PLATE]",
"cnpj": "[CNPJ]",
"postal code": "[POSTAL_CODE]",
"passport_number": "[PASSPORT_NUMBER]",
"serial number": "[SERIAL_NUM]",
"vehicle registration number": "[VEHICLE_REG_NUM]",
"credit card brand": "[CREDIT_CARD_BRAND]",
"fax number": "[FAX]",
"visa number": "[VISA]",
"insurance company": "[INSURANCE_COMPANY]",
"identity document number": "[IDENTITY_DOCUMENT]",
"transaction number": "[TRANSACTION_NUM]",
"national health insurance number": "[NATIONAL_HEALTH_INS]",
"cvc": "[CVC]",
"birth certificate number": "[BIRTH_CERT]",
"train ticket number": "[TRAIN_TICKET]",
"passport expiration date": "[PASSPORT_EXP_DATE]",
"social_security_number": "[SSN]",
}
valid_labels = list(label_to_mask_map.keys())
def mask_text(text, entities):
"""
Masks the original text by replacing entities with corresponding labels.
Args:
text (str): The original text.
entities (list): A list of dictionaries where each dictionary contains:
- "text": the extracted entity string.
- "label": the label for the entity.
Returns:
str: The masked text.
"""
# Sort entities by their occurrence in text to avoid conflicts
entities_sorted = sorted(entities, key=lambda x: text.find(x["text"]), reverse=True)
# Replace each entity with its corresponding label
for entity in entities_sorted:
label_placeholder = f"{label_to_mask_map[entity['label']]}"
text = text.replace(entity["text"], label_placeholder)
return text
def pii_masking_pipeline(
input_text,
labels=("name", "nric", "phone number", "address", "email"),
):
"""
Masks the PII entities in the input text.
Args:
input_text (str): The input text to mask.
labels (list): The list of PII entity labels to mask.
Returns:
str: The masked text.
"""
# check that the labels are a subset of valid labels
if not set(labels).issubset(valid_labels):
raise ValueError("Invalid labels provided.")
entities = model.predict_entities(input_text, labels)
masked_text = mask_text(input_text, entities)
return masked_text
|