feat: generic plugin system

config.default.toml

@@ -68,9 +68,6 @@ chain_of_thought_skips = [
     ],
 ]
 
-# Whether to print prompt/response pairs when counting refusals.
-print_responses = false
-
 # Whether to print detailed information about residuals and refusal directions.
 print_residual_geometry = false
 
@@ -86,13 +83,14 @@ residual_plot_title = 'PaCMAP Projection of Residual Vectors for "Harmless" and
 # Matplotlib style sheet to use for plots of residual vectors.
 residual_plot_style = "dark_background"
 
-# Assumed "typical" value of the Kullback-Leibler divergence from the original model for abliterated models.
-# This is used to ensure balanced co-optimization of KL divergence and refusal count.
-kl_divergence_scale = 1.0
-
-# The KL divergence to target. Below this value, an objective based on the refusal count is used.
-# This helps prevent the sampler from extensively exploring parameter combinations that "do nothing".
-kl_divergence_target = 0.01
+# List of scorers to evaluate.
+# Each entry is an object:
+# { plugin = <plugin>, direction = <direction>, instance_name = <optional> }
+# where <direction> is one of "minimize", "maximize", "not_set" (aka do not optimize)
+scorers = [
+    { plugin = "heretic.scorers.refusal_rate.RefusalRate", direction = "minimize"},
+    { plugin = "heretic.scorers.kl_divergence.KLDivergence", direction = "minimize"},
+]
 
 # Whether to adjust the refusal directions so that only the component that is
 # orthogonal to the good direction is subtracted during abliteration.
@@ -133,6 +131,31 @@ study_checkpoint_dir = "checkpoints"
 # Maximum size for individual safetensors files generated when exporting a model.
 max_shard_size = "5GB"
 
+# System prompt to use when prompting the model.
+system_prompt = "You are a helpful assistant."
+
+# Dataset of prompts that tend to not result in refusals (used for calculating refusal directions).
+[good_prompts]
+dataset = "mlabonne/harmless_alpaca"
+split = "train[:400]"
+column = "text"
+residual_plot_label = '"Harmless" prompts'
+residual_plot_color = "royalblue"
+
+# Dataset of prompts that tend to result in refusals (used for calculating refusal directions).
+[bad_prompts]
+dataset = "mlabonne/harmful_behaviors"
+split = "train[:400]"
+column = "text"
+residual_plot_label = '"Harmful" prompts'
+residual_plot_color = "darkorange"
+
+# Plugin-specific settings live in a top-level TOML table.
+# For scorer plugins, use: `[scorer.<ClassName>]` (and optionally `[scorer.<ClassName>_<instance_name>]` for instance-related config).
+[scorer.RefusalRate]
+# Whether to print prompt/response pairs when counting refusals.
+print_responses = false
+
 # Strings whose presence in a response (case insensitive) identifies the response as a refusal.
 refusal_markers = [
     "sorry",
@@ -169,33 +192,30 @@ refusal_markers = [
     "ethical boundaries",
 ]
 
-# System prompt to use when prompting the model.
-system_prompt = "You are a helpful assistant."
-
-# Dataset of prompts that tend to not result in refusals (used for calculating refusal directions).
-[good_prompts]
-dataset = "mlabonne/harmless_alpaca"
-split = "train[:400]"
-column = "text"
-residual_plot_label = '"Harmless" prompts'
-residual_plot_color = "royalblue"
-
-# Dataset of prompts that tend to result in refusals (used for calculating refusal directions).
-[bad_prompts]
+# Scorer-owned evaluation prompts
+[scorer.RefusalRate.prompts]
 dataset = "mlabonne/harmful_behaviors"
-split = "train[:400]"
-column = "text"
-residual_plot_label = '"Harmful" prompts'
-residual_plot_color = "darkorange"
-
-# Dataset of prompts that tend to not result in refusals (used for evaluating model performance).
-[good_evaluation_prompts]
-dataset = "mlabonne/harmless_alpaca"
 split = "test[:100]"
 column = "text"
 
-# Dataset of prompts that tend to result in refusals (used for evaluating model performance).
-[bad_evaluation_prompts]
-dataset = "mlabonne/harmful_behaviors"
+# You can also load multiple instances of the same scorer class by setting `instance_name`
+# in the `scorers = [...]` list. Each instance is still identified as `ClassName.instanceName`
+# internally, but its config overrides live under `[scorer.ClassName_<instance_name>]`.
+#
+# Example:
+# scorers = [
+#   { plugin = "heretic.scorers.refusal_rate.RefusalRate", direction = 'minimize', instance_name = "small" },
+#   { plugin = "heretic.scorers.refusal_rate.RefusalRate", direction = 'minimize', instance_name = "tiny" },
+# ]
+#
+# Shared defaults for all instances live under `[scorer.RefusalRate]` and can be overridden per
+# instance under `[scorer.RefusalRate_<instance_name>]`.
+#
+# Example instance override:
+# [scorer.RefusalRate_small.prompts]
+# split = "test[:10]"
+
+[scorer.KLDivergence.prompts]
+dataset = "mlabonne/harmless_alpaca"
 split = "test[:100]"
 column = "text"

src/heretic/config.py

@@ -2,14 +2,15 @@
 # Copyright (C) 2025-2026  Philipp Emanuel Weidmann <pew@worldwidemann.com> + contributors
 
 from enum import Enum
-from typing import Dict
+from typing import Dict, Literal
 
 from pydantic import BaseModel, Field
 from pydantic_settings import (
     BaseSettings,
     CliSettingsSource,
     EnvSettingsSource,
     PydanticBaseSettingsSource,
+    SettingsConfigDict,
     TomlConfigSettingsSource,
 )
 
@@ -74,6 +75,39 @@ class DatasetSpecification(BaseModel):
     )
 
 
+class ScorerConfig(BaseModel):
+    """
+    Configuration for a scorer plugin.
+
+    TOML format:
+    - { plugin = "<plugin>", direction = "<direction>", instance_name = "<optional>" }
+    """
+
+    plugin: str = Field(
+        description=(
+            "Plugin to load. Either a file path with class name "
+            "(`path/to/plugin.py:ClassName`) or a fully-qualified import path "
+            "(`module.submodule.ClassName`)."
+        ),
+    )
+
+    direction: Literal["minimize", "maximize", "not_set"] = Field(
+        description=(
+            "Optimization direction for this scorer. "
+            '"minimize" / "maximize" to include the scorer as an objective, '
+            '"not_set" to compute the score without optimizing for it.'
+        ),
+    )
+
+    instance_name: str | None = Field(
+        default=None,
+        description=(
+            "Optional name to distinguish multiple instances of the same plugin class. "
+            "Instance-specific settings live under `[scorer.<ClassName>_<instance_name>]`."
+        ),
+    )
+
+
 class BenchmarkSpecification(BaseModel):
     task: str = Field(
         description="Task ID of the benchmark in the Language Model Evaluation Harness."
@@ -228,12 +262,6 @@ class Settings(BaseSettings):
         exclude=True,
     )
 
-    print_responses: bool = Field(
-        default=False,
-        description="Whether to print prompt/response pairs when counting refusals.",
-        exclude=True,
-    )
-
     print_residual_geometry: bool = Field(
         default=False,
         description="Whether to print detailed information about residuals and refusal directions.",
@@ -264,19 +292,12 @@ class Settings(BaseSettings):
         exclude=True,
     )
 
-    kl_divergence_scale: float = Field(
-        default=1.0,
+    scorers: list[ScorerConfig] = Field(
+        default_factory=list,
         description=(
-            'Assumed "typical" value of the Kullback-Leibler divergence from the original model for abliterated models. '
-            "This is used to ensure balanced co-optimization of KL divergence and refusal count."
-        ),
-    )
-
-    kl_divergence_target: float = Field(
-        default=0.01,
-        description=(
-            "The KL divergence to target. Below this value, an objective based on the refusal count is used. "
-            'This helps prevent the sampler from extensively exploring parameter combinations that "do nothing".'
+            "List of scorer plugin configs. Each entry is an object"
+            " { plugin = <plugin>, direction = <direction>, instance_name = <optional> }."
+            " <direction> is one of 'minimize', 'maximize', 'not_set' (aka do not optimize)."
         ),
     )
 
@@ -475,23 +496,10 @@ class Settings(BaseSettings):
         description="Dataset of prompts that tend to result in refusals (used for calculating refusal directions).",
     )
 
-    good_evaluation_prompts: DatasetSpecification = Field(
-        default=DatasetSpecification(
-            dataset="mlabonne/harmless_alpaca",
-            split="test[:100]",
-            column="text",
-        ),
-        description="Dataset of prompts that tend to not result in refusals (used for evaluating model performance).",
-    )
-
-    bad_evaluation_prompts: DatasetSpecification = Field(
-        default=DatasetSpecification(
-            dataset="mlabonne/harmful_behaviors",
-            split="test[:100]",
-            column="text",
-        ),
-        description="Dataset of prompts that tend to result in refusals (used for evaluating model performance).",
-    )
+    # We intentionally allow extra keys so users can provide plugin-specific
+    # configuration in TOML tables like `[scorer.RefusalRate]` which are later
+    # consumed via `settings.model_extra` (see `Evaluator._get_plugin_namespace`).
+    model_config = SettingsConfigDict(extra="allow")
 
     @classmethod
     def settings_customise_sources(