PII Guardian PHI3 Mini LORA

This repository contains a fine-tuned model for detecting Personally Identifiable Information (PII) using PHI3 Mini with LoRA applied to the query, key, and value layers. The model is optimized for accuracy and efficiency in identifying various PII entities.

How to Use

You can use the following Python code to load and run the model:

import torch
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

# Check the available device
if torch.cuda.is_available():
    device = "cuda"
elif torch.backends.mps.is_available():
    device = "mps"
else:
    device = "cpu"

print(f"Using device: {device}")

# Check for bf16 support
is_bf16_support = False
try:
    tensor_bf16 = torch.tensor([1.0], dtype=torch.bfloat16, device=device)
    is_bf16_support = True
    print("bf16 tensors are supported.")
except TypeError:
    print("bf16 tensors are not supported.")

# Load the base model and tokenizer
base_model = "microsoft/Phi-3-mini-4k-instruct"
tokenizer = AutoTokenizer.from_pretrained(base_model)

# Load the fine-tuned PII detection model with LoRA
model = AutoModelForCausalLM.from_pretrained(base_model, return_dict=True, device_map=device, torch_dtype=torch.bfloat16 if is_bf16_support else torch.float16)
pii_model = PeftModel.from_pretrained(model, "ab-ai/PII-Guardian-Phi3-Mini-LORA")

# Example input text
input_text = "Hi Abner, just a reminder that your next primary care appointment is on 23/03/1926. Please confirm by replying to this email [email protected]."

# Define the model prompt
model_prompt = f"""### Instruction:
    Identify and extract the following PII entities from the text, if present: companyname, pin, currencyname, email, phoneimei, litecoinaddress, currency, eyecolor, street, mac, state, time, vehiclevin, jobarea, date, bic, currencysymbol, currencycode, age, nearbygpscoordinate, amount, ssn, ethereumaddress, zipcode, buildingnumber, dob, firstname, middlename, ordinaldirection, jobtitle, bitcoinaddress, jobtype, phonenumber, height, password, ip, useragent, accountname, city, gender, secondaryaddress, iban, sex, prefix, ipv4, maskednumber, url, username, lastname, creditcardcvv, county, vehiclevrm, ipv6, creditcardissuer, accountnumber, creditcardnumber. Return the output in JSON format.

    ### Input:
    {input_text}

    ### Output:
    
    """

# Tokenize the input and generate a response
inputs = tokenizer(model_prompt, return_tensors="pt").to(device)
# adjust max_new_tokens according to your need
outputs = pii_model.generate(**inputs, do_sample=True, max_new_tokens=120)
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(response) #  {'firstname': ['Abner'], 'dob': ['23/03/1926'], 'email': ['[email protected]']}