[v0.1]Test and Release: Classifier test framework

[rootfs]

Acceptance:

• Test category/PII/jailbreak detection with expected outputs
• Mock BERT model for fast CI runs
• Captures and validates the behavior of classifiers at different confidence threshold levels

[aeft]

@rootfs Hi, I'm new to this project and would love to contribute.

This issue looks like a good starting point. May I work on it?

After reviewing the code, I think it would be best to introduce a small refactor: enabling dependency injection so that we can mock model calls (e.g., the BERT model) during testing.

I’ve submitted a draft PR with an initial implementation of the ClassifyCategory tests. If there are no concerns, I plan to continue improving these tests and also add coverage for PII/jailbreak cases (using the same dependency injection approach).

Do you think the PII/jailbreak tests should be split into multiple PRs?

Any feedback or suggestions would be greatly appreciated. Thanks!

[Xunzhuo]

@aeft just go ahead, thanks for picking it up : ) Welcome

[aeft]

Hi @Xunzhuo . Thanks for your reply!

I’ve found another potential refactor: determining the model inference implementation at classifier initialization (since we already introduced dependency injection). Would you mind if I also make this change? It should help simplify the code and make the tests cleaner.

A potential concern is whether we can guarantee that the config value will remain unchanged after it’s used to initialize the classifier.
In other words, is it acceptable that the model selection is determined solely based on the initial configuration (load from file)?

---

Current Issue
In ClassifyCategory, the config is checked on every call:

if c.Config.Classifier.CategoryModel.UseModernBERT {
    result, err = c.modelInference.ClassifyModernBertText(text)
} else {
    result, err = c.modelInference.ClassifyText(text)
}

Proposed Changes

1. Create dedicated model inference implementations:

   • LinearModelInference for the linear classifier
   • ModernBertModelInference for the ModernBERT classifier

2. Select the implementation in NewClassifier based on config:

   var modelInference ModelInference
   if cfg.Classifier.CategoryModel.UseModernBERT {
       modelInference = &ModernBertModelInference{}
   } else {
       modelInference = &LinearModelInference{}
   }

3. Simplify ClassifyCategory:

   result, err := c.modelInference.ClassifyText(text)

Benefits

• Performance: removes runtime config checks
• Clean dependency injection: dependencies are resolved at creation time
• Code clarity: simpler inference logic
• Better testability: easier to mock different implementations