A privacy-aware LLM routing framework that dynamically routes sensitive and non-sensitive text chunks between local and remote LLM inference.
Inspired by the Privacy-R1 approach, this project combines PII detection, privacy-aware preprocessing, transformer-based actor-critic policy learning, supervised warm-up, and PPO reinforcement learning.
The routing objective is:
- PII-containing chunks → LOCAL
- Non-PII chunks → REMOTE
The system is designed to balance privacy protection, task performance, routing compliance, and practical remote LLM utilization.
Sending prompts to remote LLM providers can expose sensitive information when inputs contain personally identifiable information (PII).
This project addresses the problem through privacy-aware, chunk-level routing. Sensitive chunks are identified before routing, and a hard privacy constraint ensures that PII-sensitive chunks remain local while the learned policy determines routing behavior for suitable non-sensitive content.
The framework optimizes multiple objectives:
- Privacy leakage minimization
- Task performance
- Routing compliance
- Remote model utilization
- Designed a chunk-level privacy-aware routing framework for hybrid local/remote LLM inference.
- Combined PII detection with a transformer-based actor-critic routing policy.
- Used supervised warm-up to initialize the policy before PPO fine-tuning.
- Designed a reward function incorporating privacy leakage, task gain, routing compliance, and remote usage.
- Implemented a hard privacy constraint that overrides learned actions for PII-sensitive chunks.
- Included evaluation and ablation utilities for analyzing routing behavior.
Sending user prompts to remote LLM providers can expose sensitive information when the input contains personally identifiable information (PII).
This project addresses the problem by processing an input prompt as multiple chunks and making privacy-aware routing decisions for each chunk.
The pipeline:
- Preprocesses and chunks the input.
- Detects PII and sensitive entities.
- Generates embeddings using MiniLM.
- Encodes chunk representations using a transformer-based policy.
- Predicts routing actions using an actor-critic architecture.
- Enforces a hard privacy constraint so PII-containing chunks remain local.
- Computes a reward based on privacy leakage, task gain, and routing compliance.
- Performs supervised warm-up.
- Fine-tunes the routing policy using PPO.
- Evaluates the trained policy using privacy and routing metrics.
- PII and sensitive entity detection
- Quasi-identifier detection
- Clinical-style named entity recognition
- Privacy-aware preprocessing and anonymization
- MiniLM-based chunk embeddings
- Transformer-based routing policy
- Actor-critic architecture
- Supervised policy warm-up
- PPO reinforcement learning
- Local/remote hybrid LLM routing
- Hard privacy constraint for PII chunks
- Privacy leakage evaluation
- Routing compliance evaluation
- Remote routing ratio analysis
- Task-gain evaluation
- Ablation analysis
- Training and evaluation visualizations
- Saved PPO routing-agent checkpoint
Input Prompt
│
▼
Text Preprocessing
│
▼
Chunking
│
▼
PII / Sensitive Entity
Detection
│
▼
MiniLM Embeddings
│
▼
Transformer Encoder
│
▼
Actor-Critic Policy
│ │
│ └── Value / Critic
│
▼
Routing Action
│
┌─────────┴─────────┐
│ │
PII / Sensitive Non-PII
│ │
▼ ▼
LOCAL REMOTE
LLM LLM
│ │
└─────────┬─────────┘
▼
Reward Function
│
▼
PPO Training
The transformer policy produces a candidate routing action for each chunk. A hard privacy layer then verifies the routing decision. If a chunk is identified as PII-sensitive, the final action is forced to LOCAL, overriding the learned action.
This separates policy learning from privacy enforcement.
The routing policy is designed around a privacy-aware constraint:
PII-containing chunk → LOCAL
Non-PII chunk → REMOTE
The learned policy optimizes the trade-off between:
- Privacy leakage
- Task gain
- Routing compliance
- Remote model utilization
A hard privacy constraint prevents PII-containing chunks from being intentionally routed to the remote model.
src/phase1_dataset.py
Loads and prepares the training, validation, and test data.
src/phase2_preprocessing.py
Performs privacy-related preprocessing, entity detection, and sensitive-span handling.
src/phase3_agent.py
Defines the transformer-based actor-critic routing policy.
The policy contains:
- Positional encoding
- Transformer encoder layers
- Actor head for routing actions
- Critic head for value estimation
src/phase4_reward.py
Computes the reinforcement-learning reward using privacy and task-related signals.
The reward considers:
- Privacy leakage
- Task gain
- Routing compliance
- Remote usage
src/phase5_rollout.py
Executes routing decisions and interacts with the local and remote LLM components.
src/phase6_warmup.py
Generates heuristic routing labels based on privacy information and provides an initial policy before PPO training.
src/phase7_ppo.py
Fine-tunes the routing policy using Proximal Policy Optimization (PPO).
Privacy-Aware-LLM-Routing/
│
├── data/
│ ├── train.json
│ ├── validation.json
│ ├── test.json
│ └── gold_responses.json
│
├── evaluation/
│ ├── ablation.py
│ ├── evaluate.py
│ ├── metrics.py
│ ├── plot_results.py
│ ├── routing_test.py
│ ├── sanity_test.py
│ └── validate_dataset.py
│
├── models/
│ └── privacy_r1_agent.pt
│
├── results/
│ ├── figures/
│ │ ├── privacy_leak_curve.png
│ │ ├── remote_ratio_curve.png
│ │ ├── reward_curve.png
│ │ ├── routing_compliance_curve.png
│ │ └── task_gain_curve.png
│ │
│ ├── ablation.json
│ ├── ppo_summary.json
│ ├── ppo_training.csv
│ ├── routing_evaluation.json
│ ├── test_evaluation.json
│ ├── validation_evaluation.json
│ └── final_output.json
│
├── scripts/
│ └── prepare_clinical_data.py
│
├── src/
│ ├── phase1_dataset.py
│ ├── phase2_preprocessing.py
│ ├── phase3_agent.py
│ ├── phase4_reward.py
│ ├── phase5_rollout.py
│ ├── phase6_warmup.py
│ └── phase7_ppo.py
│
├── .env.example
├── .gitignore
├── main.py
├── Privacy_R1.ipynb
├── README.md
├── requirements.txt
└── test_remote.py
git clone https://github.com/<your-username>/Privacy-Aware-LLM-Routing.git
cd Privacy-Aware-LLM-Routingpython -m venv .venv
.venv\Scripts\Activate.ps1python3 -m venv .venv
source .venv/bin/activatepip install -r requirements.txtInstall the required spaCy model:
python -m spacy download en_core_web_smThe project uses a Hugging Face token for remote model access.
Copy the provided environment template.
copy .env.example .envcp .env.example .envThen edit .env:
HF_TOKEN=your_huggingface_token_hereReplace the placeholder with your own Hugging Face access token.
Never commit your actual Hugging Face token.
The repository intentionally contains:
.env.example
but excludes:
.env
through .gitignore.
The token is loaded through an environment variable rather than being hard-coded into the source code.
The repository contains:
data/
├── train.json
├── validation.json
├── test.json
└── gold_responses.json
Each training/validation/test record is expected to contain the required fields:
query
pii_units
target_response
The training pipeline validates the dataset schema before training.
The included data is intended for research and experimentation. Do not add real personally identifiable or confidential information to the repository.
Run a minimal training execution:
python main.py --steps 2 --warmup-epochs 1 --ppo-batch-size 2 --ppo-epochs 1This is useful for verifying the pipeline and configuration.
Example:
python main.py --steps 300 --warmup-epochs 1 --ppo-batch-size 8 --ppo-epochs 4--train-file data/train.json
--skip-warmup
--no-debug-rollout
For all available arguments:
python main.py --helppython evaluation/sanity_test.pyTests the privacy routing boundary and verifies that PII-containing chunks are kept local.
python evaluation/routing_test.pypython evaluation/evaluate.py data/validation.json --output results/validation_evaluation.jsonpython evaluation/ablation.py data/test.json --output results/ablation.jsonpython evaluation/validate_dataset.py data/train.jsonpython evaluation/plot_results.pypython test_remote.pyThe project includes visualizations from the recorded training and evaluation runs.
The evaluation analyzes:
- Privacy leakage
- Routing compliance
- Remote routing ratio
- Reward progression
- Task gain
- Ablation behavior
These visualizations are provided to illustrate the behavior of the implemented routing policy across the recorded experiments.
Note: The visualizations represent results from the included experimental runs. They are provided for demonstration and reproducibility purposes and should not be interpreted as formal guarantees of privacy, utility, or generalization.
A trained Privacy-R1 routing-agent checkpoint is included at:
models/privacy_r1_agent.pt
The checkpoint contains the learned parameters of the transformer-based actor-critic routing agent, including:
- Positional encoding
- Transformer encoder layers
- Actor head
- Critic head
The checkpoint can be used for evaluation without retraining the routing policy from scratch.
Training and evaluation outputs include:
models/privacy_r1_agent.pt
results/
├── figures/
│ ├── privacy_leak_curve.png
│ ├── remote_ratio_curve.png
│ ├── reward_curve.png
│ ├── routing_compliance_curve.png
│ └── task_gain_curve.png
│
├── ablation.json
├── ppo_summary.json
├── ppo_training.csv
├── routing_evaluation.json
├── test_evaluation.json
├── validation_evaluation.json
└── final_output.json
- Python
- PyTorch
- Transformers
- Hugging Face Hub
- Sentence Transformers
- spaCy
- Microsoft Presidio
- NumPy
- Pandas
- Scikit-learn
- Matplotlib
This project focuses on privacy-aware LLM routing.
Important security practices include:
- PII detection before routing
- Hard local routing for PII-containing chunks
- Environment-variable based credential management
.envexcluded from version control.env.exampleprovided for configuration- No API tokens embedded in source code
- Synthetic/research data used for repository examples
PII detection systems can produce false positives and false negatives. Therefore, the routing mechanism should not be interpreted as a universal privacy guarantee.
For production deployment, additional privacy controls, auditing, access controls, encryption, and security testing would be required.
Results can vary between runs because of:
- PPO stochasticity
- Random initialization
- GPU execution
- Remote model/provider behavior
- Provider availability and quotas
- Dataset differences
- Model version changes
For meaningful comparisons, use the same:
- Dataset splits
- Model configuration
- Training parameters
- Evaluation configuration
The included checkpoint and evaluation outputs are provided to make the reported experiment easier to reproduce.
If you receive an error indicating that the Hugging Face token is missing:
- Make sure
.envexists. - Check that it contains:
HF_TOKEN=your_actual_token- Make sure the token is valid.
- Run:
python test_remote.pyRun:
python -m spacy download en_core_web_smRemote inference may fail because of:
- Invalid credentials
- Provider availability
- Account limitations
- Model availability
- Network issues
- Provider quotas
Check your Hugging Face account and model/provider access.
This repository contains a research-oriented implementation of privacy-aware LLM routing, including:
- Privacy preprocessing
- Transformer-based routing
- Supervised warm-up
- PPO fine-tuning
- Hybrid local/remote inference
- Evaluation
- Ablation analysis
- Training visualizations
- Trained routing-agent checkpoint
The implementation is intended for research, experimentation, and educational use.
No open-source license has currently been specified.
If this project is released publicly for reuse, an appropriate license should be added to the repository.