feat: Generic BERT Classifier scorer plugin (#500)

* feat: Refusal classifier scorer plugin

Using the generic BERTClassifier scorer plugin

* docs: Add modified_model_type

* fix: Close [red] tag in the error

.gitignore

@@ -21,8 +21,5 @@ wheels/
 # Study checkpoints
 checkpoints/
 
-# Residual plots
-plots/
-
 # Models generated by tests
 /tests/*/model/

config.classifier.toml

@@ -0,0 +1,14 @@
+# Rename this file to config.toml, place it in the working directory
+# that you run Heretic from, and edit the configuration to your liking.
+#
+# BERTClassifier scorer instead of the default KeywordRate.
+# To compare with keyword-based scoring alongside,
+# Use "none" to compute its score for reference only without influencing optimization.
+# Or use "minimize" to use its scores for optimization as well.
+#
+# { plugin = "heretic.scorers.keyword_rate.KeywordRate", optimization = "none" }
+
+scorers = [
+    { plugin = "heretic.scorers.bert_classifier.BERTClassifier", optimization = "minimize" },
+    { plugin = "heretic.scorers.kl_divergence.KLDivergence", optimization = "minimize" },
+]

config.default.toml

@@ -217,6 +217,71 @@ split = "test[:100]"
 column = "text"
 
 
+# Generic text-classification scorer: classifies each response with a
+# text-classification model and counts the responses whose
+# predicted label is one of the `match_labels`, as an alternative to the
+# `keyword_markers` of the KeywordRate scorer.
+#
+# The default settings detect refusals using a ModernBERT model fine-tuned on refusal
+# classification. Any text-classification model and label set can be used by
+# changing the plugin settings.
+#
+# Add or replace an entry in the `scorers` array above:
+# { plugin = "heretic.scorers.bert_classifier.BERTClassifier", optimization = "minimize" }
+#
+[scorer.BERTClassifier]
+# Name that describes what the configured match rate measures.
+score_name = "Refusals (Classifier)"
+
+# Hugging Face model ID, or path to a text-classification model on disk.
+model = "VINAY-UMRETHE/Refusal-Classifier"
+
+# Hugging Face commit hash of the classifier model.
+# Set this to pin a specific revision; unset means the latest revision
+# is auto-pinned and recorded when the classifier is loaded.
+# When reproducing, the recorded revision is restored and used.
+# model_commit = "<commit-hash>"
+
+# List of PyTorch dtypes to try when loading model tensors.
+# If loading with a dtype fails, the next dtype in the list will be tried.
+dtypes = [
+    "auto",
+    "float16",
+    "bfloat16",
+    "float32",
+]
+
+# Device map to pass to Accelerate when loading the model.
+device_map = "auto"
+
+# Labels predicted by the classifier that identify a response as a match.
+# Defaults to the refusal labels of the default model, which should
+# be counted as refusals.
+# Valid labels: compliant, full_refusal, partial_refusal, capability_limitation, clarification.
+match_labels = ["full_refusal", "partial_refusal"]
+
+# Number of prompt/response pairs to classify in parallel.
+batch_size = 16
+
+# Maximum number of tokens per classifier input (longer inputs are truncated).
+max_length = 8192
+
+# Whether to pass inputs as prompt/response text pairs.
+# Requires a model trained on pairs (e.g. the default refusal classifier).
+# Set to False when using a single-sequence text-classification model
+# to evaluate the response string only (e.g. sentiment analysis).
+use_text_pairs = true
+
+# Whether to print prompt/response pairs and their predicted labels.
+print_responses = false
+
+# Dataset of prompts to evaluate the match rate on.
+[scorer.BERTClassifier.prompts]
+dataset = "mlabonne/harmful_behaviors"
+split = "test[:100]"
+column = "text"
+
+
 [scorer.BenchmarkScore]
 # Name that describes what the configured benchmark score measures.
 score_name = "PIQA acc_norm"

config.emotions.toml

@@ -0,0 +1,120 @@
+# Rename this file to config.toml, place it in the working directory
+# that you run Heretic from, and edit the configuration to your liking.
+#
+# BERTClassifier scorer for personality/emotion steering.
+#
+# This uses a DistilRoBERTa model fine-tuned on emotion classification
+# and allows minimizing or maximizing specific emotions in the model's responses.
+#
+# Example use: configure two scorer instances, one to minimize negative emotions
+# (e.g. anger, disgust, fear, sadness) and one to maximize positive emotions
+# (e.g joy, surprise) or vice versa.
+#
+# This can be done in one of the two ways:
+#
+# 1. Multi‑objective steering:
+# Here, we use two instances: `negative_emotions` and `positive_emotions` with
+# `optimization = "minimize"` and `optimization = "maximize"` respectively.
+#
+# Minimizing the first score pushes the model *away* from those negative emotions.
+# Maximizing the second score pushes the model *towards* those emotions.
+#
+# 2. Single‑objective steering:
+# With two instances, heretic builds a Pareto front that **balances** both objectives above.
+# Due to which, the model learns to be less angry *and* more joyful **simultaneously**.
+#
+# With only one instance, we have only one objective which can be either minimized or maximized. 
+# Due to which, the model will reduce anger, but it *may* drift towards **anything** that lowers
+# the other most natural instincts (e.g. neutral, overly polite, or wierdly happy).
+#
+# With the first method, you can expect to control both objectives effectively,
+# which is the recommended method and is the default.
+# With the second method, you can only *hope* and *assume* the most expected common/natural
+# behavior, depending upon what you tried to "minimize" or "maximize".
+scorers = [
+    { plugin = "heretic.scorers.kl_divergence.KLDivergence", optimization = "minimize" },
+    # Scorer to minimize negative emotions.
+    { plugin = "heretic.scorers.bert_classifier.BERTClassifier", optimization = "minimize", instance_name = "negative_emotions" },
+    # Scorer to maximize positive emotions.
+    { plugin = "heretic.scorers.bert_classifier.BERTClassifier", optimization = "maximize", instance_name = "positive_emotions" },
+]
+
+modified_model_type = "happiness-increased"
+
+upload_tags = [
+    "personality-steering",
+    "emotion-steering",
+]
+
+# Dataset of prompts used for steering, which tends to produce
+# the target emotion you want to influence.
+# For example:
+# - For minimizing anger: prompts about creative writing which involves those emotions.
+# - For maximizing joy: prompts about celebrations, happy scenarios, good news from user, etc.
+[modifier.ARA.good_prompts]
+dataset = "VINAY-UMRETHE/Positive-Prompts"
+split = "train[:200]"
+column = "prompt"
+
+[modifier.ARA.bad_prompts]
+dataset = "VINAY-UMRETHE/Negative-Prompts"
+split = "train[:200]"
+column = "prompt"
+
+
+[scorer.BERTClassifier_negative_emotions]
+# Name that describes what the configured match rate measures.
+score_name = "Minimize"
+
+# Hugging Face model ID, or path to a text-classification model on disk.
+model = "j-hartmann/emotion-english-distilroberta-base"
+
+# Hugging Face commit hash of the classifier model.
+# Set this to pin a specific revision; unset means the latest revision
+# is auto-pinned and recorded when the classifier is loaded.
+# When reproducing, the recorded revision is restored and used.
+# model_commit = "<commit-hash>"
+
+# Labels predicted by the classifier that identify a response as a match.
+# We want to minimize these emotions.
+match_labels = ["anger", "disgust", "fear", "sadness"]
+
+# Number of prompt/response pairs to classify in parallel.
+batch_size = 16
+
+# Maximum number of tokens per classifier input (longer inputs are truncated).
+max_length = 512
+
+# Whether to pass inputs as prompt/response text pairs.
+# Requires a model trained on pairs (e.g. the default refusal classifier).
+# Set to False when using a single-sequence text-classification model
+# to evaluate the response string only (e.g. sentiment analysis).
+use_text_pairs = false
+
+# Whether to print prompt/response pairs and their predicted labels.
+print_responses = true
+
+# Dataset of prompts to evaluate the match rate on.
+# Use a dataset that tends to produce the target emotions (anger, disgust, fear, sadness).
+[scorer.BERTClassifier_negative_emotions.prompts]
+dataset = "VINAY-UMRETHE/Negative-Prompts"
+split = "train[:100]"
+column = "prompt"
+
+
+[scorer.BERTClassifier_positive_emotions]
+score_name = "Maximize"
+
+model = "j-hartmann/emotion-english-distilroberta-base"
+
+# We want to maximize these emotions.
+match_labels = ["joy", "surprise"]
+
+batch_size = 16
+max_length = 512
+print_responses = true
+
+[scorer.BERTClassifier_positive_emotions.prompts]
+dataset = "VINAY-UMRETHE/Positive-Prompts"
+split = "train[:100]"
+column = "prompt"

src/heretic/scorers/bert_classifier.py

@@ -0,0 +1,271 @@
+# SPDX-License-Identifier: AGPL-3.0-or-later
+# Copyright (C) 2025-2026  Philipp Emanuel Weidmann <pew@worldwidemann.com> + contributors
+
+# Generic text-classification scorer: classifies each response with a
+# text-classification model and counts the responses whose
+# predicted label is one of the `match_labels`, as an alternative to the
+# `keyword_markers` of the KeywordRate scorer.
+#
+# The default settings detect refusals using a ModernBERT model fine-tuned on refusal
+# classification. Any text-classification model and label set can be used by
+# changing the plugin settings.
+#
+# For more information, see:
+# Implementation: https://github.com/p-e-w/heretic/pull/500
+# Default classifier model: https://huggingface.co/VINAY-UMRETHE/Refusal-Classifier
+
+from typing import cast
+
+import huggingface_hub
+import torch
+from pydantic import BaseModel, Field
+from rich.markup import escape
+from transformers import Pipeline, pipeline
+
+from heretic.config import DatasetSpecification, SingleDatasetSpecification
+from heretic.plugin import Context
+from heretic.scorer import Score, Scorer
+from heretic.system import empty_cache
+from heretic.utils import (
+    format_dataset_specification,
+    format_exception,
+    is_hf_path,
+    print,
+)
+
+DEFAULT_LABELS = ["full_refusal", "partial_refusal"]
+
+
+class Settings(BaseModel):
+    score_name: str = Field(
+        default="Refusals (Classifier)",
+        description="Name that describes what the configured match rate measures.",
+    )
+
+    model: str = Field(
+        default="VINAY-UMRETHE/Refusal-Classifier",
+        description="Hugging Face model ID, or path to a text-classification model on disk.",
+    )
+
+    model_commit: str | None = Field(
+        default=None,
+        description=(
+            "Hugging Face commit hash of the classifier model. "
+            "Set this to pin a specific revision; unset means the latest "
+            "revision is auto-pinned and recorded when the classifier is loaded. "
+            "When reproducing, the recorded revision is restored and used."
+        ),
+    )
+
+    dtypes: list[str] = Field(
+        default=[
+            "auto",
+            "float16",
+            "bfloat16",
+            "float32",
+        ],
+        description=(
+            "List of PyTorch dtypes to try when loading model tensors. "
+            "If loading with a dtype fails, the next dtype in the list will be tried."
+        ),
+    )
+
+    device_map: str | dict[str, int | str] = Field(
+        default="auto",
+        description="Device map to pass to Accelerate when loading the model.",
+    )
+
+    match_labels: list[str] = Field(
+        default=DEFAULT_LABELS,
+        description=(
+            "Labels predicted by the classifier that identify a response as a match. "
+            "Defaults to the refusal labels of the default model, which should "
+            "be counted as refusals. "
+            "Valid labels: compliant, full_refusal, partial_refusal, "
+            "capability_limitation, clarification."
+        ),
+    )
+
+    batch_size: int = Field(
+        default=16,
+        description="Number of prompt/response pairs to classify in parallel.",
+        gt=0,
+    )
+
+    max_length: int = Field(
+        default=8192,
+        description="Maximum number of tokens per classifier input (longer inputs are truncated).",
+        gt=0,
+    )
+
+    use_text_pairs: bool = Field(
+        default=True,
+        description=(
+            "Whether to pass inputs as prompt/response text pairs. "
+            "Requires a model trained on pairs (e.g. the default refusal classifier). "
+            "Set to False when using a single-sequence text-classification model "
+            "to evaluate the response string only (e.g. sentiment analysis)."
+        ),
+    )
+
+    print_responses: bool = Field(
+        default=False,
+        description="Whether to print prompt/response pairs and their predicted labels.",
+    )
+
+    prompts: DatasetSpecification = Field(
+        default=SingleDatasetSpecification(
+            dataset="mlabonne/harmful_behaviors",
+            split="test[:100]",
+            column="text",
+        ),
+        description="Dataset of prompts to evaluate the match rate on.",
+    )
+
+
+class BERTClassifier(Scorer):
+    """
+    Classifies responses using a text-classification model
+    (ModernBERT by default) instead of regex keyword matching.
+    Scores the share of responses whose predicted label is in `match_labels`.
+
+    The default model and labels detect refusals. Any text-classification
+    model and label set can be used by changing the plugin settings.
+
+    Settings live in the [scorer.BERTClassifier] config section.
+    """
+
+    settings: Settings
+
+    @property
+    def reproducible(self) -> bool:
+        # Inference is deterministic: eval mode (no dropout),
+        # argmax over softmax outputs (no sampling), pipeline keeps input order.
+        # All varying inputs are pinned: the classifier model revision is resolved
+        # in init(), the prompts dataset revision is resolved when the prompts are
+        # loaded, the remaining settings come from the config
+        # (match_labels, batch_size, max_length, use_text_pairs).
+        return True
+
+    @property
+    def score_name(self) -> str:
+        return self.settings.score_name
+
+    def init(self, ctx: Context) -> None:
+        print()
+        print(
+            f"Loading {self.settings.score_name} evaluation prompts from [bold]{format_dataset_specification(self.settings.prompts)}[/]..."
+        )
+        self.prompts = ctx.load_prompts(self.settings.prompts)
+        print(f"* [bold]{len(self.prompts)}[/] prompts loaded")
+        print()
+        print(f"Loading text classifier: [bold]{self.settings.model}[/]...")
+
+        if self.settings.model_commit is None and is_hf_path(self.settings.model):
+            try:
+                self.settings.model_commit = huggingface_hub.model_info(
+                    self.settings.model
+                ).sha
+            except Exception as error:
+                raise RuntimeError(
+                    f"[red]Failed to fetch the latest commit hash for the classifier model "
+                    f"[bold]{self.settings.model}[/] ({error}). "
+                    f"Pinning the model commit is required for exact reproducibility.[/]"
+                )
+
+        self.pipeline = None
+        for dtype in self.settings.dtypes:
+            print(f"* Trying dtype [bold]{dtype}[/]...")
+            try:
+                torch_dtype = "auto" if dtype == "auto" else getattr(torch, dtype)
+                self.pipeline = pipeline(
+                    "text-classification",
+                    model=self.settings.model,
+                    revision=self.settings.model_commit,
+                    device_map=self.settings.device_map,
+                    torch_dtype=torch_dtype,
+                )
+                break
+            except Exception as error:
+                self.pipeline = None
+                empty_cache()
+
+                formatted = format_exception(error)
+                if "\n" in formatted:
+                    print(f"* [red]Failed:\n{formatted}[/]")
+                else:
+                    print(f"* [red]Failed ({formatted})[/]")
+                continue
+
+        if self.pipeline is None:
+            raise Exception(
+                "Failed to load text classifier model with all configured dtypes."
+            )
+
+        valid = set((self.pipeline.model.config.id2label or {}).values())
+        unknown = set(self.settings.match_labels) - valid
+        if unknown:
+            raise ValueError(
+                f"Unknown match_labels: {sorted(unknown)}. "
+                f"Valid labels: {sorted(valid)}"
+            )
+
+    def get_score(self, ctx: Context) -> Score:
+        responses = ctx.get_responses(self.prompts)
+        predicted_labels = self._classify(responses)
+
+        match_count = 0
+        for prompt, response, label in zip(self.prompts, responses, predicted_labels):
+            # Evaluate matches based on the classifier's predicted label only.
+            is_match = label in self.settings.match_labels
+            match_count += int(is_match)
+
+            if self.settings.print_responses:
+                print()
+                print(f"[bold]System prompt:[/] {escape(prompt.system)}")
+                print(f"[bold]Prompt:[/] {escape(prompt.user)}")
+                if not response.strip():
+                    response = "[italic]\\[empty][/]"
+                else:
+                    response = escape(response)
+                print(
+                    f"[bold]Response:[/] [{'red' if is_match else 'green'}]{response}[/]"
+                )
+                print(f"[bold]Predicted label:[/] {label}")
+
+        if self.settings.print_responses:
+            print()
+
+        return Score(
+            value=float(match_count / len(self.prompts)),
+            rich_display=f"[bold]{match_count}[/]/{len(self.prompts)}",
+            md_display=f"{match_count}/{len(self.prompts)}",
+        )
+
+    def _classify(self, responses: list[str]) -> list[str]:
+        # Switch input method, considering both model types.
+        if self.settings.use_text_pairs:
+            # Standard way to use the classifier models that are trained on pairs,
+            # (e.g. the default refusal classifier). Otherwise, the
+            # fallback below can degrade the predictions of such models.
+            inputs = [
+                {
+                    "text": prompt.user,
+                    "text_pair": response,
+                }
+                for prompt, response in zip(self.prompts, responses)
+            ]
+        else:
+            # Fallback for single-sequence text-classification models
+            # (e.g. sentiment analysis). Here, we pass only the responses
+            # to avoid string concatenation (e.g. Prompt: {prompt.user}\nResponse: {response}),
+            # since such formatting can degrade the predictions of such models.
+            inputs = responses
+
+        results = cast(Pipeline, self.pipeline)(
+            inputs,
+            batch_size=self.settings.batch_size,
+            truncation=True,
+            max_length=self.settings.max_length,
+        )
+        return [result["label"] for result in results]

tests/README.md

@@ -32,7 +32,7 @@ sha256sum * | Out-File -Encoding utf8NoBOM ../SHA256SUMS.LABEL
 **Verify with**:
 
 ```bash
-Get-Command sha256sum`
+Get-Command sha256sum
 ```
 
 **Expected**:

tests/gemma-3/SHA256SUMS.ci

@@ -0,0 +1,4 @@
+7a79d414471add74c6bb9fc202407c38e59e917119614d6362fdb8cb3ca1ce94 *config.json
+f8f1fd0ef18d709772584f1a4d736fc8f31aafee83a001b34d333c30a606baf8 *model.safetensors
+daab2354f8a74e70d70b4d1f804939b68a8c9624dd06cb7858e52dd8970e9726 *tokenizer.json
+3a194bc6c50615274af46113faa5967e2119b1093da5f305ed151af06d6a4402 *tokenizer_config.json

tests/gemma-3/SHA256SUMS.ci2

@@ -0,0 +1,4 @@
+7a79d414471add74c6bb9fc202407c38e59e917119614d6362fdb8cb3ca1ce94 *config.json
+2b0b38f5cc5b25e424dc944cf0b27432f736cffd59a8753f2ddd23f5434e5ab4 *model.safetensors
+daab2354f8a74e70d70b4d1f804939b68a8c9624dd06cb7858e52dd8970e9726 *tokenizer.json
+3a194bc6c50615274af46113faa5967e2119b1093da5f305ed151af06d6a4402 *tokenizer_config.json

tests/gemma-3/SHA256SUMS.ci3

@@ -0,0 +1,4 @@
+7a79d414471add74c6bb9fc202407c38e59e917119614d6362fdb8cb3ca1ce94 *config.json
+ad8c373f6d6f3f00e646d118116fc1f34c91c5c67a10db6fe9366fd34f534df5 *model.safetensors
+daab2354f8a74e70d70b4d1f804939b68a8c9624dd06cb7858e52dd8970e9726 *tokenizer.json
+3a194bc6c50615274af46113faa5967e2119b1093da5f305ed151af06d6a4402 *tokenizer_config.json

tests/gemma-3/SHA256SUMS.windows

@@ -0,0 +1,4 @@
+5595714af3f245d20831e545ee5ee41fec8f4ad5a89d9ca5abce0bf476debfa7 *config.json
+333322c2521ae1cbe6405f4770d690d8f404d0ce5587308b94520f058b949762 *model.safetensors
+daab2354f8a74e70d70b4d1f804939b68a8c9624dd06cb7858e52dd8970e9726 *tokenizer.json
+43184dd63293fc63deba9c8aab4353b4a15e46a2e7ab691ddcf4fd16d73cf77c *tokenizer_config.json

tests/gemma-3/config.toml

@@ -0,0 +1,76 @@
+# This test case is for the BERTClassifier scorer plugin.
+# After any change related to it, this test should PASS.
+
+model = "tiny-random/gemma-3"
+model_commit = "69a78d1ad0ad66620c43579acd1327553713e22a"
+
+seed = 12345
+print_debug_information = true
+
+batch_size = 2
+max_response_length = 10
+
+n_trials = 2
+n_startup_trials = 1
+
+export_strategy = "merge"
+checkpoint_action = "restart"
+trial_index = 0
+model_action = "save"
+save_directory = "model"
+
+scorers = [
+    { plugin = "heretic.scorers.keyword_rate.KeywordRate", optimization = "none" },
+    { plugin = "heretic.scorers.bert_classifier.BERTClassifier", optimization = "minimize" },
+    { plugin = "heretic.scorers.kl_divergence.KLDivergence", optimization = "minimize" },
+]
+
+[[response_prefix_test_prompts]]
+dataset = "mlabonne/harmless_alpaca"
+commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f"
+split = "train[:5]"
+column = "text"
+
+[[response_prefix_test_prompts]]
+dataset = "mlabonne/harmful_behaviors"
+commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7"
+split = "train[:5]"
+column = "text"
+
+[scorer.KLDivergence.prompts]
+dataset = "mlabonne/harmless_alpaca"
+commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f"
+split = "test[:5]"
+column = "text"
+
+[scorer.KeywordRate]
+print_responses = true
+
+[scorer.KeywordRate.prompts]
+dataset = "mlabonne/harmful_behaviors"
+commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7"
+split = "test[:5]"
+column = "text"
+
+[scorer.BERTClassifier]
+model = "VINAY-UMRETHE/tiny-refusal-classifier"
+model_commit = "ca106855daaa6d695bd2a66934eb1a40cbd3f18a"
+print_responses = true
+
+[scorer.BERTClassifier.prompts]
+dataset = "mlabonne/harmful_behaviors"
+commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7"
+split = "test[:5]"
+column = "text"
+
+[modifier.ARA.good_prompts]
+dataset = "mlabonne/harmless_alpaca"
+commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f"
+split = "train[:5]"
+column = "text"
+
+[modifier.ARA.bad_prompts]
+dataset = "mlabonne/harmful_behaviors"
+commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7"
+split = "train[:5]"
+column = "text"