Fix for issue with reading ModernBERT jailbreak classifier model

[srampal]

Summary

The issue with the jailbreak classifier has been successfully resolved. Here's what was causing the problem and
how it was fixed:

Root Cause

The jailbreak classifier model was saved with PyTorch's torch.compile() or similar compilation process, which
added the _orig_mod. prefix to all tensor names. For example:

• Expected: model.embeddings.tok_embeddings.weight
• Actual: _orig_mod.model.embeddings.tok_embeddings.weight

Solution

The fix implements a PrefixedVarBuilder wrapper that:

1. Auto-detects whether a model uses the _orig_mod. prefix
2. Tries both prefixed and unprefixed tensor names for maximum compatibility
3. Maintains context so that child VarBuilders (for head, classifier) use the same namespace as the base model

Key Changes

1. Added PrefixedVarBuilder struct in
   /home/ubuntu/repos/github.com/srampal/semantic_router/candle-binding/src/modernbert.rs
2. Modified model loading to use consistent VarBuilder context for all components
3. Implemented fallback logic that tries _orig_mod.{tensor_name} first, then {tensor_name} if that fails

Result

✅ All three classifiers now initialize successfully:

• Category classifier (models without prefix)
• PII token classifier (models without prefix)
• Jailbreak classifier (model with _orig_mod. prefix)

✅ Server starts successfully: The router now starts on port 50051 and all components are functional.

The fix is backward-compatible and will work with both types of models (with and without the _orig_mod.
prefix), ensuring robust model loading for future updates.

[rootfs]

Thanks for the testing. It looks the regression was introduced in this commit

Per this thread, torhc.compile adds new tensor names.

[rootfs]

don't hard code the model shape

[rootfs]

@srampal can you reduce the scope of this patch? A simple change is to prepends _orig_mod. to pytorch compiled models if the first load call fails.

            match FixedModernBertForTokenClassification::load(vb.clone(), &config) {
                Ok(model) => ModernBertModel::Token(model),
                Err(_) => {
                    // Try with _orig_mod prefix (torch.compile models)
                    ModernBertModel::Token(FixedModernBertForTokenClassification::load(vb.pp("_orig_mod"), &config)?)
                }
            }

[srampal]

Reverted this fix .. working on an alternate fix via
#75