PII classification

app.py

@@ -0,0 +1,181 @@
+import streamlit as st
+# from gliner import GLiNER
+from datasets import load_dataset
+from peft import PeftModel, PeftConfig
+import threading
+import time
+import torch
+from torch.profiler import profile, record_function, ProfilerActivity
+from transformers import AutoModelForTokenClassification, AutoTokenizer, pipeline
+
+def predict_entities(text, labels, entity_set):
+    if labels == []:
+        entities = recognizer(text)
+        for entity in entities:
+            if entity['entity'] in entity_set:
+                entity_set[entity['entity']] += 1
+            else:
+                entity_set[entity['entity']] = 1
+    else:
+        # Use Gliner labels
+        entities = model.predict_entities(text, labels, threshold = 0.7)
+        for entity in entities:
+            if entity['label'] in entity_set:
+                entity_set[entity['label']] += 1
+            else:
+                entity_set[entity['label']] = 1
+
+def process_datasets(start, end, unmasked_text, sizes, index, entity_set, labels):
+    size = 0
+    text = ""
+    for i in range(start, end):
+        if len(text) < 700:
+            text = text + " " + unmasked_text[i]
+        else:
+            size += len(text)
+            predict_entities(text, labels, entity_set)
+            
+            text = unmasked_text[i]
+    sizes[index] = size
+
+device = torch.device('cpu')
+print(f"Is CUDA available: {torch.cuda.is_available()}")
+# True
+if torch.cuda.is_available():
+    print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
+    device = torch.device('cuda')
+
+# Load the fine-tuned GLiNER model
+st.write('Loading the pretrained model ...')
+base_model_name = "iiiorg/piiranha-v1-detect-personal-information"
+adapter_model_name = "CarolXia/xia-lora-deberta-v2"
+# config = PeftConfig.from_pretrained(model_name)
+base_model = AutoModelForTokenClassification.from_pretrained(base_model_name)
+if torch.cuda.is_available():
+    base_model = base_model.to("cuda")
+adapter_model = PeftModel.from_pretrained(base_model, adapter_model_name, token=st.secrets["HUGGINGFACE_TOKEN"])
+# Try quantization instead
+# model = AutoModelForTokenClassification.from_pretrained(model_name, device_map="auto", load_in_8bit=True)
+tokenizer = AutoTokenizer.from_pretrained(base_model_name)
+
+model = adapter_model.merge_and_unload()
+model.save_pretrained("./merged_model")
+if torch.cuda.is_available():
+    model = model.to("cuda")
+
+recognizer = pipeline("ner", model=model, tokenizer=tokenizer)
+# model_name = "urchade/gliner_multi_pii-v1"
+# model = GLiNER.from_pretrained(model_name)
+
+# print weights
+pytorch_total_params = sum(p.numel() for p in model.parameters())
+torch_total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+print(f'total params: {pytorch_total_params}. tunable params: {torch_total_params}')
+
+
+
+# Sample text containing PII/PHI entities
+text = """
+Hello Jane Doe. Your AnyCompany Financial Services, LLC credit card account
+4111-0000-1111-0000 has a minimum payment of $24.53 that is due by July 31st.
+Based on your autopay settings, we will withdraw your payment on the due date from
+your bank account XXXXXX1111 with the routing number XXXXX0000.
+
+Your latest statement was mailed to 100 Main Street, Anytown, WA 98121.
+After your payment is received, you will receive a confirmation text message
+at 206-555-0100.
+
+If you have questions about your bill, AnyCompany Customer Service is available by
+phone at 206-555-0199 or email at [email protected].
+"""
+
+# Define the labels for PII/PHI entities
+labels = [
+    "medical_record_number",
+    "date_of_birth",
+    "ssn",
+    "date",
+    "first_name",
+    "email",
+    "last_name",
+    "customer_id",
+    "employee_id",
+    "name",
+    "street_address",
+    "phone_number",
+    "ipv4",
+    "credit_card_number",
+    "license_plate",
+    "address",
+    "user_name",
+    "device_identifier",
+    "bank_routing_number",
+    "date_time",
+    "company_name",
+    "unique_identifier",
+    "biometric_identifier",
+    "account_number",
+    "city",
+    "certificate_license_number",
+    "time",
+    "postcode",
+    "vehicle_identifier",
+    "coordinate",
+    "country",
+    "api_key",
+    "ipv6",
+    "password",
+    "health_plan_beneficiary_number",
+    "national_id",
+    "tax_id",
+    "url",
+    "state",
+    "swift_bic",
+    "cvv",
+    "pin"
+]
+
+st.write('Trying a sample first')
+st.write(text)
+# Predict entities with a confidence threshold of 0.7
+# entities = model.predict_entities(text, labels, threshold=0.7)
+entities = recognizer(text) 
+
+# Display the detected entities
+for entity in entities:
+    st.write(entity)
+
+st.write('Processing the full dataset now ...')
+entity_set=dict()
+dataset = load_dataset("Isotonic/pii-masking-200k", split="train")
+dataset = dataset.select(range(50))
+dataset.set_format("torch", device="cuda")
+unmasked_text = dataset['unmasked_text'] # This will load the entire column inmemory. Must do this to avoid I/O delay later
+
+st.write('Size of the dataset ', dataset.num_rows)
+sizes = [0] * 2
+start = time.time()
+t1 = threading.Thread(target=process_datasets, args=(0, 25, unmasked_text, sizes, 0, entity_set, []))
+t2 = threading.Thread(target=process_datasets, args=(25, 50, unmasked_text, sizes, 1, entity_set, []))
+# with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], profile_memory=True, record_shapes=True) as prof:
+process_datasets(0, 50, unmasked_text, sizes, 0, entity_set, [])
+# recognizer(unmasked_text)
+    # t1.start()
+    # t2.start()
+
+    # t1.join()
+    # t2.join()
+
+end = time.time()
+length = end - start
+
+# Show the results : this can be altered however you like
+st.write('Bytes processed ', sum(sizes))
+st.write("It took", length, "seconds!")
+
+# Display the summary
+st.write('Total entities found')
+for key in entity_set:
+    st.write(key, ' => ', entity_set[key])
+
+# st.write(prof.key_averages().table(sort_by="cpu_time_total", row_limit=10))

requirements.txt

@@ -0,0 +1,9 @@
+datasets
+gliner
+peft
+torch>=2.0.0
+transformers>=4.38.2
+huggingface_hub>=0.21.4
+onnxruntime
+sentencepiece
+tqdm