nisarg6502/Llama3-150M-PII-Redactor

🧠 Custom 150M Parameter LLM: PII Firewall

Model Details

This is a custom 150-million parameter autoregressive language model built entirely from scratch in PyTorch. It has been modernized with Llama-3 architectural features and fine-tuned specifically to act as a local Personally Identifiable Information (PII) redaction microservice.

• Architecture: 12 Transformer blocks, 12 attention heads, 768 hidden dimension.
• Upgrades: Rotary Positional Embeddings (RoPE), RMSNorm, SwiGLU Activations.
• Format: safetensors (Optimized for secure, CPU-bound inference).

Intended Use

This model is designed to intercept and scrub sensitive entities (names, emails, phone numbers, addresses, etc.) from raw text before the data is transmitted to external cloud services (like OpenAI or Anthropic). It acts as a highly constrained, local privacy firewall.

How to Use

Because this is a custom architecture, it does not use the standard transformers auto-classes. You must load the weights into your custom PyTorch class.

from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
import torch

# 1. Download the weights securely
model_path = hf_hub_download(
    repo_id="nisarg6502/Llama3-150M-PII-Redactor", 
    filename="pii_model_epoch_3.safetensors"
)

# 2. Load into your custom GPT PyTorch class
device = torch.device("cpu")
state_dict = load_file(model_path, device=str(device))

# model = GPT(config) # (Instantiate your custom class here)
model.load_state_dict(state_dict)
model.eval()

Training Data & Methodology

• Pre-training: Trained from scratch on WikiText-103 to establish baseline English grammar, syntax, and structural understanding.
• Fine-tuning: Fine-tuned on the ai4privacy/pii-masking-200k dataset.
• Optimization: Utilized strict target-token loss masking during Supervised Fine-Tuning (SFT). The model was only penalized for missing exact PII redaction tags, forcing it to isolate entity redaction pathways without degrading its underlying language capabilities.