Proposed fix: Enable LoRA PII auto-detection for Issue #647

[yossiovadia]

Context

While investigating Issue #647 (PII detection confidence issues), I discovered that PII classification appears to be hardcoded to ModernBERT, even though:

• LoRA PII models exist in the model directory
• The Rust layer already has auto-detection infrastructure via lora_config.json checks
• Other classifiers (intent, jailbreak) use the auto-detecting code path

Observations

Current behavior:

• PII uses: InitModernBertPIITokenClassifier() / ClassifyModernBertPIITokens()
• Testing with ModernBERT: 27% success rate (10/37 test cases)

After switching to auto-detection:

• PII uses: InitCandleBertTokenClassifier() / ClassifyCandleBertTokens() (same as intent classifier)
• Testing with LoRA: 73% success rate (27/37 test cases)

Proposed Changes

I wanted to share this potential fix for your review. The changes are minimal (23 lines in classifier.go) and leverage existing auto-detection infrastructure:

Current (Hardcoded):

success := candle_binding.InitModernBertPIITokenClassifier(modelID, useCPU)
result := candle_binding.ClassifyModernBertPIITokens(text, configPath)

Proposed (Auto-Detection):

success := candle_binding.InitCandleBertTokenClassifier(modelID, numClasses, useCPU)
result := candle_binding.ClassifyCandleBertTokens(text)

The Rust layer handles auto-detection by checking for lora_config.json presence.

Test Results

Created comprehensive test suite (37 diverse PII cases):

Approach	Success Rate	Notes
ModernBERT (current)	27% (10/37)	Low confidence, wrong entity types
LoRA (proposed)	73% (27/37)	Higher confidence, correct entity types

Example improvements:

• ✅ Email detection: EMAIL_ADDRESS (0.9) instead of PERSON (0.52)
• ✅ SSN detection: US_SSN (0.9) instead of failed/low confidence
• ✅ Credit Card: CREDIT_CARD (0.9) - previously failed
• ✅ Phone: PHONE_NUMBER (0.9) - previously failed

Files in This PR

1. classifier.go - Switch to auto-detecting functions
2. config.e2e.yaml - Update test config to use LoRA model
3. 06-a-test-pii-direct.py - New: Comprehensive PII test suite
4. pii-confidence-benchmark.py - New: Statistical benchmark tool

Questions for Maintainers

1. Was the ModernBERT hardcoding intentional? Or just an oversight that could be updated?
2. Is the auto-detection approach acceptable? It's already used for intent/jailbreak classifiers
3. Confidence uniformity: LoRA returns exactly 0.9 for all detections (vs ModernBERT's varied scores). Is this expected behavior for LoRA models? See detailed analysis

Testing

# Run direct PII tests (37 cases)
python3 e2e-tests/06-a-test-pii-direct.py

# Run comprehensive benchmark (84 prompts with statistics)
python3 e2e-tests/pii-confidence-benchmark.py

Backward Compatibility

• ✅ Falls back to Traditional/ModernBERT when LoRA not available
• ✅ No breaking config changes
• ✅ use_modernbert field behavior unchanged (was already ignored)

---

Happy to adjust this approach based on your architectural preferences. The main goal is enabling better PII detection for Issue #647.

[netlify]

✅ Deploy Preview for vllm-semantic-router ready!

Name	Link
🔨 Latest commit	20ef032
🔍 Latest deploy log	https://app.netlify.com/projects/vllm-semantic-router/deploys/691d484fa4c9820009ae2fa1
😎 Deploy Preview	https://deploy-preview-648--vllm-semantic-router.netlify.app
📱 Preview on mobile	Toggle QR Code... Use your smartphone camera to open QR code link.

---

To edit notification comments on pull requests, go to your Netlify project configuration.

[github-actions]

👥 vLLM Semantic Team Notification

The following members have been identified for the changed files in this PR and have been automatically assigned:

📁 e2e-tests

Owners: @yossiovadia
Files changed:

• e2e-tests/06-a-test-pii-direct.py
• e2e-tests/pii-confidence-benchmark.py

📁 config

Owners: @rootfs, @Xunzhuo
Files changed:

• config/testing/config.e2e.yaml

📁 e2e

Owners: @Xunzhuo
Files changed:

• e2e/go.mod

📁 src

Owners: @rootfs, @Xunzhuo, @wangchen615
Files changed:

• src/semantic-router/pkg/classification/classifier.go
• src/semantic-router/pkg/classification/classifier_test.go
• src/semantic-router/pkg/extproc/extproc_test.go
• src/semantic-router/pkg/extproc/router.go
• src/semantic-router/pkg/services/classification.go

---

🎉 Thanks for your contributions!

This comment was automatically generated based on the OWNER files in the repository.

[yossiovadia]

Further Investigation on 0.9 Confidence

I conducted additional investigation to understand where the uniform 0.9 confidence originates.

Test: Pure Python Model Inference

Created a standalone Python test that loads the LoRA PII model directly using HuggingFace transformers (bypassing ALL semantic-router code - no Go, no Rust FFI).

Test script: https://gist.github.com/yossiovadia/c1a0e822e836d73db68ea9fe9e321adc

Results:

Total PII detections: 23
Unique confidence values: 23 (ALL DIFFERENT\!)

Confidence distribution:
  Min: 0.203108
  Max: 0.996829
  Mean: 0.633599
  Std Dev: 0.278768

Sample detections:
  Email "john": 0.985991
  Phone "(": 0.996829
  SSN "123": 0.991570
  SSN "45": 0.619693

Conclusion: The LoRA PII model itself produces varied probabilistic confidence scores, not uniform 0.9.

Code Path Analysis

Traced the inference path through the codebase:

1. Go Layer (classifier.go:182-184):

   func (c *PIIInferenceImpl) ClassifyTokens(...) (...) {
       return candle_binding.ClassifyCandleBertTokens(text)  // Direct passthrough
   }

   No confidence manipulation ✓

2. Rust FFI (classify.rs:621):

   .map(|r| (r.token.clone(), r.label_name.clone(), r.confidence))  // Direct passthrough

   No confidence manipulation ✓

3. LoRA Classifier (bert_lora.rs:844):

   results.push((token.clone(), predicted_class, confidence));  // Direct from softmax

   No confidence manipulation ✓

4. PII Aggregation (pii_lora.rs:133,144):

   occurrences.push(PIIOccurrence {
       confidence: *confidence,  // Individual token confidence
   });
   // Final confidence = average of all token confidences
   confidence_scores.iter().sum::<f32>() / confidence_scores.len() as f32

   Uses averaging of token-level scores ✓

Hypothesis

The uniform 0.9 likely comes from aggregation behavior in pii_lora.rs. When the Rust LoRA inference processes tokens through the model, the token-level confidences (which the Python model shows are varied) get aggregated into an entity-level confidence.

The mystery: Why does this averaging consistently produce exactly 0.9? This suggests either:

1. Token-level confidences from the Rust LoRA path are different from Python's inference
2. The aggregation logic has some normalization we haven't identified
3. There's a threshold or post-processing step in the Candle framework's LoRA handling

How to Reproduce

# Install dependencies
pip install peft transformers torch

# Run the test
python3 test_lora_pii_pure_python.py

The test will show the model produces varied confidence scores when accessed directly via Python/HuggingFace.

Bottom Line

The uniform 0.9 is not from our Issue #647 changes or any hardcoded value in semantic-router. It appears to be an artifact of how the Rust/Candle LoRA implementation processes and aggregates token classifications, which differs from the Python/HuggingFace inference path.

Despite this quirk, the core improvement remains: 27% → 73% success rate with correct entity types.

[yossiovadia]

Further Investigation: Jailbreak LoRA Confidence Comparison

To isolate the uniform 0.9 confidence issue, I tested the jailbreak LoRA model using the same semantic-router pathway (Go → Rust → Candle).

Test Results:

Jailbreak LoRA Model via Classification API (/api/v1/classify/security):

• ✅ 14 unique confidence values from 15 test cases
• ✅ Confidence range: 0.9917 to 0.9999
• ✅ NO uniform 0.9 issue

Test script: https://gist.github.com/yossiovadia/3c016171b776d2ed1b62cacdeb452e7a

Sample output:

Total tests: 15
Unique confidence values: 14
Min confidence: 0.9917060137
Max confidence: 0.9999994040

🔍 Uniform 0.9 issue: NO ✅
✅ Jailbreak model shows VARIED confidence scores

Comparison:

Model	Architecture	Pathway	Confidence Behavior
PII LoRA	Token Classification	Go → Rust → Candle	❌ Uniform 0.9 (121/121)
Jailbreak LoRA	Sequence Classification	Go → Rust → Candle	✅ Varied (0.9917-0.9999)
PII LoRA	Token Classification	Python direct	✅ Varied (0.203-0.997) [gist]

Root Cause:

The uniform 0.9 confidence is specific to PII token classification in the Rust/Candle pathway:

Jailbreak implementation (security_lora.rs:81-82):

let (predicted_class, confidence) = self.bert_classifier.classify_text(text)

• Sequence classification → returns raw model confidence

PII implementation (pii_lora.rs:87-89, 142-148):

let token_results = self.bert_token_classifier.classify_tokens(text)
let final_confidence = confidence_scores.iter().sum::<f32>() / confidence_scores.len() as f32

• Token classification → averages confidence scores across detected PII tokens

The uniform 0.9 originates from the Rust token classification implementation (classify_tokens() or averaging logic).

[Xunzhuo]

https://github.com/vllm-project/semantic-router/actions/runs/19407606041?pr=648

[Xunzhuo]

can you change e2e config for ai-gateway as well? deploy/kubernetes/ai-gateway/semantic-router-values/values.yaml. but this is still weird for me, this should be backward compatible right? But this makes the previous full param classification not work