akshatamadavi/qwen-pii-masker-lora-rank32

Qwen 0.6B — PII Masker (LoRA Fine-tuned)

A lightweight PII (Personally Identifiable Information) masking model fine-tuned from Qwen/Qwen2.5-0.5B using LoRA on the ai4privacy/pii-masking-400k dataset.

Given any natural language text, the model replaces sensitive entities with labeled placeholders like [GIVENNAME], [EMAIL], [DATEOFBIRTH], [CREDITCARDNUMBER] etc. — making it suitable as a privacy layer before text is sent to external LLM APIs.

---

Model Details

Property	Value
Base model	Qwen/Qwen2.5-0.5B
Fine-tuning method	LoRA (Low-Rank Adaptation)
Training paradigm	Instruction fine-tuning (SFT)
LoRA rank (r)	32
LoRA alpha	64
Trainable parameters	6.4M out of 600M (1%)
Training data	20,000 US English records from ai4privacy/pii-masking-400k
Training hardware	Google Colab T4 GPU (16GB)
Training time	~2 hours
Max sequence length	512 tokens

---

Intended Use

Primary use case: Pre-processing user prompts before they are sent to external LLM APIs (ChatGPT, Claude, Gemini etc.) to prevent accidental PII leakage.

Example application: A browser extension that intercepts AI chatbot prompts, screens them through this model locally, and blocks or redacts any sensitive information before the request leaves the organization.

Supported PII entity types:

GIVENNAME · SURNAME · EMAIL · DATEOFBIRTH · TELEPHONENUM · STREET · CITY · ZIPCODE · CREDITCARDNUMBER · IDCARDNUM · DRIVERLICENSENUM · PASSWORD · USERNAME · ACCOUNTNUM · SOCIALNUM · TAXNUM · BUILDINGNUM

---

How to Use

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

# Load base model + LoRA adapter
base_model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-0.5B",
    torch_dtype=torch.float16,
    device_map="auto",
    trust_remote_code=True
)

model = PeftModel.from_pretrained(base_model, "akshatamadavi/qwen-pii-masker-lora")
tokenizer = AutoTokenizer.from_pretrained("akshatamadavi/qwen-pii-masker-lora")
tokenizer.pad_token = tokenizer.eos_token
model.eval()

def mask_pii(text):
    prompt = (
        "### Task: Mask all PII in the following text.\n"
        f"### Input:\n{text}\n"
        "### Output:\n"
    )
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=200,
            do_sample=False,
            pad_token_id=tokenizer.eos_token_id
        )
    new_tokens = outputs[0][inputs['input_ids'].shape[1]:]
    return tokenizer.decode(new_tokens, skip_special_tokens=True).strip()


# Example
text = "Student: John Smith DOB: 01/15/1990 Email: john.smith@gmail.com ZIP: 90210"
print(mask_pii(text))
# → "Student: [GIVENNAME] [SURNAME] DOB: [DATEOFBIRTH] Email: [EMAIL] ZIP: [ZIPCODE]"

---

Training Details

Dataset

• Source: ai4privacy/pii-masking-400k
• Filters applied: US locale, English language, no HTML tags, max 400 characters
• Train split: 20,000 records
• Validation split: 2,000 records
• Label simplification: Numbered labels ([GIVENNAME_1]) simplified to [GIVENNAME] for training stability

Instruction Format

Each training example was formatted as:

### Task: Mask all PII in the following text.
### Input:
my ssn is 123456789, can you tell me my credit score?
### Output:
my ssn is [SSN], can you tell me my credit score?

LoRA Configuration

LoraConfig(
    r=32,
    lora_alpha=64,
    lora_dropout=0.05,
    bias="none",
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                    "gate_proj", "up_proj", "down_proj"],
    task_type="CAUSAL_LM"
)

Training Arguments

TrainingArguments(
    num_train_epochs=2,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,   # effective batch = 16
    learning_rate=3e-4,
    warmup_steps=100,
    fp16=True,
    eval_steps=250,
    save_steps=250,
)

---

Evaluation Results

Evaluated on 500 held-out US English validation samples.

Metric	Score
Entity-level Precision	90.5%
Entity-level Recall	91.5%
Entity-level F1	91.0%
Exact Match	82.2%

Note: Recall is the primary metric for this task. A missed PII entity (false negative) results in sensitive data reaching an external API. A false positive (over-masking) is recoverable.

---

Limitations

• Trained on US English only — may underperform on non-US address formats or non-English text
• Optimized for texts under 400 characters — longer inputs may degrade quality
• Label numbering is simplified ([GIVENNAME] not [GIVENNAME_1]) — cannot distinguish between two people in the same text
• Should not be the sole privacy control in a production system — pair with rule-based fallbacks for structured PII like email regex

---

Project Context

This model was developed as part of an Enterprise LLM Privacy Monitoring Platform — a Chrome extension that intercepts prompts sent to AI chatbots (ChatGPT, Claude, Gemini) in the browser, screens them for PII using this self-hosted model, and blocks the API call if sensitive data is detected. All PII screening happens on-premise so prompts never leave the organization's infrastructure even during the check.

---

Author

Akshata Madavi

MS Software Engineering, San José State University

LinkedIn · GitHub