File size: 3,476 Bytes
57f002f |
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 |
import torch
from transformers import AutoTokenizer, AutoModelForTokenClassification
import gradio as gr
# Load the tokenizer and model
model_name = "iiiorg/piiranha-v1-detect-personal-information"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForTokenClassification.from_pretrained(model_name)
# Set device to GPU if available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
def apply_redaction(masked_text, start, end, pii_type, aggregate_redaction):
for j in range(start, end):
masked_text[j] = ''
if aggregate_redaction:
masked_text[start] = '[redacted]'
else:
masked_text[start] = f'[{pii_type}]'
def mask_pii(text, aggregate_redaction=True):
# Tokenize input text
inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
inputs = {k: v.to(device) for k, v in inputs.items()}
# Get the model predictions
with torch.no_grad():
outputs = model(**inputs)
# Get the predicted labels
predictions = torch.argmax(outputs.logits, dim=-1)
# Convert token predictions to word predictions
encoded_inputs = tokenizer.encode_plus(text, return_offsets_mapping=True, add_special_tokens=True)
offset_mapping = encoded_inputs['offset_mapping']
masked_text = list(text)
is_redacting = False
redaction_start = 0
current_pii_type = ''
for i, (start, end) in enumerate(offset_mapping):
if start == end: # Special token
continue
label = predictions[0][i].item()
if label != model.config.label2id['O']: # Non-O label
pii_type = model.config.id2label[label]
if not is_redacting:
is_redacting = True
redaction_start = start
current_pii_type = pii_type
elif not aggregate_redaction and pii_type != current_pii_type:
# End current redaction and start a new one
apply_redaction(masked_text, redaction_start, start, current_pii_type, aggregate_redaction)
redaction_start = start
current_pii_type = pii_type
else:
if is_redacting:
apply_redaction(masked_text, redaction_start, end, current_pii_type, aggregate_redaction)
is_redacting = False
# Handle case where PII is at the end of the text
if is_redacting:
apply_redaction(masked_text, redaction_start, len(masked_text), current_pii_type, aggregate_redaction)
return ''.join(masked_text)
# Define the function for Gradio interface
def redact_text(text, aggregate_redaction):
return mask_pii(text, aggregate_redaction)
# Create Gradio Interface
demo = gr.Interface(
fn=redact_text,
inputs=[
gr.Textbox(lines=5, label="Enter Text with Potential PII"),
gr.Checkbox(label="Aggregate Redaction", value=True)
],
outputs="text",
title="PII Detection and Redaction",
description="This application detects personal identifiable information (PII) and redacts it from the provided text. You can choose to either aggregate all PII redaction into a single '[redacted]' label or keep each PII type labeled individually.",
examples=[
["John Doe's phone number is 123-456-7890, and his email is [email protected]."],
["Jane was born on 12th August, 1990 and her SSN is 987-65-4321."]
]
)
if __name__ == "__main__":
demo.launch()
|