File size: 2,529 Bytes
072d53b b86ee16 |
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 |
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import gradio as gr
# Set up the model and tokenizer
device = "cuda" if torch.cuda.is_available() else "cpu"
model = AutoModelForCausalLM.from_pretrained(
"betterdataai/PII_DETECTION_MODEL",
trust_remote_code=True
).to(device)
tokenizer = AutoTokenizer.from_pretrained(
"betterdataai/PII_DETECTION_MODEL",
trust_remote_code=True
)
classes_list = ['<pin>','<api_key>','<bank_routing_number>','<bban>','<company>','<credit_card_number>','<credit_card_security_code>','<customer_id>','<date>','<date_of_birth>','<date_time>','<driver_license_number>','<email>','<employee_id>','<first_name>','<iban>','<ipv4>','<ipv6>','<last_name>','<local_latlng>','<name>','<passport_number>','<password>','<phone_number>','<social_security_number>','<street_address>','<swift_bic_code>','<time>','<user_name>']
prompt_template = """You are an AI assistant who is responisble for identifying Personal Identifiable information (PII). You will be given a passage of text and you have to \
identify the PII data present in the passage. You should only identify the data based on the classes provided and not make up any class on your own.
```PII Classes```
{classes}
The given text is:
{text}
The PII data are:
"""
def detect_pii(user_input_text):
try:
# 1. Format the prompt
new_prompt = prompt_template.format(classes="\n".join(classes_list), text=user_input_text)
# 2. Tokenize
tokenized_input = tokenizer(new_prompt, return_tensors="pt").to(device)
# 3. Generate output
output = model.generate(**tokenized_input, max_new_tokens=250)
# 4. Decode the PII part
# Use rsplit to be safer, splitting only on the last occurrence
decoded_output = tokenizer.decode(output[0], skip_special_tokens=True)
if "The PII data are:\n" in decoded_output:
pii_classes = decoded_output.rsplit("The PII data are:\n", 1)[1]
else:
pii_classes = "Could not parse model output."
return pii_classes
except Exception as e:
return f"An error occurred: {str(e)}"
# 3. Create the Gradio app
iface = gr.Interface(
fn=detect_pii,
inputs=gr.Textbox(lines=5, label="Enter Text Here"),
outputs=gr.Textbox(label="Detected PII"),
title="PII Detection Model",
description="This app uses 'betterdataai/PII_DETECTION_MODEL' to find PII in text."
)
iface.launch()
|