Skip to content

Commit 8f1701e

Browse files
committed
fix(pii): add fallback to ModernBERT for backward compatibility
Implement graceful fallback strategy for PII initialization: 1. Try auto-detecting InitCandleBertTokenClassifier (enables LoRA) 2. Fallback to InitModernBertPIITokenClassifier if auto-detect fails This maintains backward compatibility with existing ModernBERT models that have incomplete configs (e.g., missing hidden_act field) while still enabling LoRA PII models when available. Also disable PII for caching tests (not needed for those test cases). Resolves test failures while preserving the 27% → 73% improvement. Signed-off-by: Yossi Ovadia <[email protected]>
1 parent 5298c30 commit 8f1701e

File tree

2 files changed

+21
-5
lines changed

2 files changed

+21
-5
lines changed

src/semantic-router/pkg/classification/classifier.go

Lines changed: 19 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -143,15 +143,29 @@ type PIIInitializer interface {
143143
Init(modelID string, useCPU bool, numClasses int) error
144144
}
145145

146-
type PIIInitializerImpl struct{}
146+
type PIIInitializerImpl struct {
147+
usedModernBERT bool // Track which init path succeeded for inference routing
148+
}
147149

148150
func (c *PIIInitializerImpl) Init(modelID string, useCPU bool, numClasses int) error {
149-
// Use auto-detecting Candle BERT init - checks for lora_config.json
151+
// Try auto-detecting Candle BERT init first - checks for lora_config.json
152+
// This enables LoRA PII models when available
150153
success := candle_binding.InitCandleBertTokenClassifier(modelID, numClasses, useCPU)
151-
if !success {
152-
return fmt.Errorf("failed to initialize PII token classifier")
154+
if success {
155+
c.usedModernBERT = false
156+
logging.Infof("Initialized PII token classifier with auto-detection (LoRA or Traditional BERT)")
157+
return nil
158+
}
159+
160+
// Fallback to ModernBERT-specific init for backward compatibility
161+
// This handles models with incomplete configs (missing hidden_act, etc.)
162+
logging.Infof("Auto-detection failed, falling back to ModernBERT PII initializer")
163+
err := candle_binding.InitModernBertPIITokenClassifier(modelID, useCPU)
164+
if err != nil {
165+
return fmt.Errorf("failed to initialize PII token classifier (both auto-detect and ModernBERT): %w", err)
153166
}
154-
logging.Infof("Initialized PII token classifier with auto-detection (LoRA or Traditional)")
167+
c.usedModernBERT = true
168+
logging.Infof("Initialized ModernBERT PII token classifier (fallback mode)")
155169
return nil
156170
}
157171

src/semantic-router/pkg/extproc/extproc_test.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2030,6 +2030,8 @@ var _ = Describe("Caching Functionality", func() {
20302030
BeforeEach(func() {
20312031
cfg = CreateTestConfig()
20322032
cfg.Enabled = true
2033+
// Disable PII detection for caching tests (not needed and avoids model loading issues)
2034+
cfg.InlineModels.Classifier.PIIModel.ModelID = ""
20332035

20342036
var err error
20352037
router, err = CreateTestRouter(cfg)

0 commit comments

Comments
 (0)