feat: add modifier base class

* fix: improve print output of scorers

* feat: add benchmark scorer

config.default.toml

@@ -91,8 +91,8 @@ residual_plot_style = "dark_background"
 # { plugin = <plugin>, optimization = <optimization>, instance_name = <optional> }
 # where <optimization> is one of "minimize", "maximize", "none" (do not optimize)
 scorers = [
-    { plugin = "heretic.scorers.keyword_rate.KeywordRate", optimization = "minimize"},
-    { plugin = "heretic.scorers.kl_divergence.KLDivergence", optimization = "minimize"},
+    { plugin = "heretic.scorers.keyword_rate.KeywordRate", optimization = "minimize" },
+    { plugin = "heretic.scorers.kl_divergence.KLDivergence", optimization = "minimize" },
 ]
 
 # Whether to adjust the residual directions so that only the component that is
@@ -157,6 +157,9 @@ 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.KeywordRate]
+# Name that describes what the configured keyword rate measures.
+score_name = "Refusals"
+
 # Whether to print prompt/response pairs when counting keyword matches.
 print_responses = false
 

config.nohumor.toml

@@ -20,6 +20,8 @@ residual_plot_label = "Humorous prompts"
 residual_plot_color = "darkorange"
 
 [scorer.KeywordRate]
+score_name = "Responses with humor"
+
 keyword_markers = [
     "😅",
     "here's one",

config.noslop.toml

@@ -24,6 +24,8 @@ residual_plot_label = "Slop-inducing prompts"
 residual_plot_color = "darkorange"
 
 [scorer.KeywordRate]
+score_name = "Responses with slop"
+
 keyword_markers = [
     "Eldoria",
     "Lumina",

config.piqa.toml

@@ -0,0 +1,7 @@
+# Rename this file to config.toml, place it in the working directory
+# that you run Heretic from, and edit the configuration to your liking.
+
+scorers = [
+    { plugin = "heretic.scorers.keyword_rate.KeywordRate", optimization = "minimize" },
+    { plugin = "heretic.scorers.benchmark_score.BenchmarkScore", optimization = "maximize" },
+]

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "heretic-llm"
-version = "1.4.0"
+version = "2.0.0.dev0"
 description = "Fully automatic censorship removal for language models"
 readme = "README.md"
 license = "AGPL-3.0-or-later"

src/heretic/evaluator.py

@@ -40,9 +40,11 @@ class Evaluator:
         print("Loading and initializing scorers...")
         self._load_and_init_scorers()
 
-        # Establish baseline scores (pre-abliteration).
+        print()
+        print("Getting baseline scores...")
         self.baseline_scores = self.get_baseline_scores()
-        self._print_baseline()
+        for name, score in self.baseline_scores:
+            print(f"* Baseline [bold]{name}:[/] [green]{score.rich_display}[/]")
 
     def _load_and_init_scorers(self) -> None:
         """
@@ -108,11 +110,6 @@ class Evaluator:
         for entry in self._scorer_entries:
             entry.scorer.init(ctx)
 
-    def _print_baseline(self) -> None:
-        """Print baseline scores summary."""
-        for name, score in self.baseline_scores:
-            print(f"* Baseline {name}: [bold]{score.rich_display}[/]")
-
     def get_dataset_specifications(self) -> list[DatasetSpecification]:
         """
         Collect the dataset specifications declared in the settings of all

src/heretic/main.py

@@ -66,6 +66,7 @@ from optuna.trial import FrozenTrial, TrialState, create_trial
 from pydantic import ValidationError
 from questionary import Choice, Style
 from rich.table import Table
+from rich.text import Text
 from rich.traceback import install
 
 from .analyzer import Analyzer
@@ -519,10 +520,8 @@ def run():
         settings.model = settings.evaluate_model
         model.reset_model()
         print("* Evaluating...")
-        print()
-        print("[bold]Metrics:[/]")
-        for score_name, score in evaluator.get_scores():
-            print(f"  * {score_name}: [bold]{score.rich_display}[/]")
+        for name, score in evaluator.get_scores():
+            print(f"  * [bold]{name}:[/] [green]{score.rich_display}[/]")
         return
 
     if not reproduction_mode and not evaluator.get_objective_names():
@@ -673,7 +672,7 @@ def run():
 
         print()
         print(
-            f"Running trial [bold]{trial_index}[/] of [bold]{settings.n_trials}[/]..."
+            f"[magenta]Running trial [bold]{trial_index}[/] of [bold]{settings.n_trials}[/]...[/]"
         )
         print("* Parameters:")
         for name, value in get_trial_parameters(trial).items():
@@ -685,10 +684,8 @@ def run():
         print("* Evaluating...")
         scores = evaluator.get_scores()
         objective_values = evaluator.get_objective_values(scores)
-
-        print("  * Metrics:")
         for name, score in scores:
-            print(f"    * {name}: [bold]{score.rich_display}[/]")
+            print(f"  * [bold]{name}:[/] [green]{score.rich_display}[/]")
 
         elapsed_time = time.perf_counter() - start_time
         remaining_time = (elapsed_time / (trial_index - start_index)) * (
@@ -793,7 +790,7 @@ def run():
                 score_parts: list[str] = []
                 for score in trial.user_attrs["scores"]:
                     name = score["name"]
-                    value = score["score"]["rich_display"]
+                    value = Text.from_markup(score["score"]["rich_display"]).plain
                     score_parts.append(f"{name}: {value}")
 
                 return f"{prefix} " + ", ".join(score_parts)
@@ -828,7 +825,6 @@ def run():
                         "After selecting a trial, you will be able to save the model, upload it to Hugging Face, "
                         "chat with it to test how well it works, or run standard benchmarks on it. "
                         "You can return to this menu later to select a different trial. "
-                        "[yellow]Note that KL divergence values above 0.5 usually indicate significant damage to the original model's capabilities.[/]"
                     )
                 )
 

src/heretic/modifier.py

@@ -0,0 +1,53 @@
+# SPDX-License-Identifier: AGPL-3.0-or-later
+# Copyright (C) 2025-2026  Philipp Emanuel Weidmann <pew@worldwidemann.com> + contributors
+
+from abc import ABC, abstractmethod
+from typing import Generic, TypeVar
+
+from optuna import Trial
+from pydantic import BaseModel
+
+from heretic.plugin import Context, Plugin
+
+from .config import Settings as HereticSettings
+
+Parameters = TypeVar("Parameters")
+
+
+class Modifier(Generic[Parameters], Plugin, ABC):
+    """
+    Abstract base class for modifier plugins.
+
+    Modifiers modify models based on an implementation-dependent set of optimizable parameters.
+
+    Examples: Standard abliteration, ARA, SOMA, etc.
+    """
+
+    def __init__(
+        self,
+        heretic_settings: HereticSettings,
+        settings: BaseModel | None = None,
+    ) -> None:
+        super().__init__(heretic_settings=heretic_settings, settings=settings)
+
+    @abstractmethod
+    def suggest_parameters(self, ctx: Context, trial: Trial) -> Parameters:
+        """
+        Sample parameters for a trial using the trial's `suggest_*` methods,
+        collect them in an implementation-dependent parameters object, and
+        return that object.
+        """
+
+    @abstractmethod
+    def modify_model(self, ctx: Context, parameters: Parameters) -> None:
+        """
+        Modify the model (obtainable via `ctx.get_model()`)
+        according to the provided parameters.
+        """
+
+    @abstractmethod
+    def reset_model(self, ctx: Context) -> None:
+        """
+        Reset the model (obtainable via `ctx.get_model()`),
+        undoing any changes made by `modify_model`.
+        """

src/heretic/plugin.py

@@ -149,12 +149,8 @@ def load_plugin(
 
 class Context:
     """
-    Runtime context passed to plugins
-
-    Provides plugin-safe access to the model.
-
-    Plugins must use `get_responses(...)`, `get_logits(...)`, etc.
-    Direct access to the underlying Model is intentionally not exposed.
+    Runtime context passed to plugins.
+    Acts as a quasi-API for plugins to access Heretic functionality.
     """
 
     def __init__(self, settings: HereticSettings, model: Model) -> None:
@@ -180,6 +176,13 @@ class Context:
     def get_residuals(self, prompts: list[Prompt]) -> Tensor:
         return self._model.get_residuals_batched(prompts)
 
+    def get_model(self) -> Model:
+        """
+        Prefer managed methods (`get_responses` etc.) unless you
+        actually need access to the model object.
+        """
+        return self._model
+
     def load_prompts(self, specification: DatasetSpecification) -> list[Prompt]:
         return load_prompts(self._settings, specification)
 
@@ -211,8 +214,11 @@ class Plugin:
         return False
 
     def __init__(
-        self, *, heretic_settings: HereticSettings, settings: BaseModel | None = None
-    ):
+        self,
+        *,
+        heretic_settings: HereticSettings,
+        settings: BaseModel | None = None,
+    ) -> None:
         # Plugins that declare a settings schema should always receive
         # validated plugin settings from the evaluator.
         settings_model = self.__class__.get_settings_model()

src/heretic/scorer.py

@@ -32,7 +32,7 @@ class Scorer(Plugin, ABC):
 
     Scorers evaluate model behavior and return a Score.
 
-    Example: counting refusals, measuring KL divergence, etc.
+    Examples: Counting refusals, measuring KL divergence, etc.
     """
 
     @property
@@ -47,7 +47,7 @@ class Scorer(Plugin, ABC):
         self,
         heretic_settings: HereticSettings,
         settings: BaseModel | None = None,
-    ):
+    ) -> None:
         super().__init__(heretic_settings=heretic_settings, settings=settings)
 
     @abstractmethod

src/heretic/scorers/benchmark_score.py

@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: AGPL-3.0-or-later
+# Copyright (C) 2025-2026  Philipp Emanuel Weidmann <pew@worldwidemann.com> + contributors
+
+import lm_eval
+from lm_eval.models.huggingface import HFLM
+from pydantic import BaseModel, Field
+
+from heretic.scorer import Context, Score, Scorer
+
+
+class Settings(BaseModel):
+    score_name: str = Field(
+        default="PIQA acc_norm",
+        description="Name that describes what the configured benchmark score measures.",
+    )
+
+    task: str = Field(
+        default="piqa",
+        description="Task ID of the benchmark in the Language Model Evaluation Harness.",
+    )
+
+    metric: str = Field(
+        default="acc_norm,none",
+        description="Task metric to use as the benchmark score.",
+    )
+
+
+class BenchmarkScore(Scorer):
+    """
+    Calculates the score of a benchmark from the Language Model Evaluation Harness.
+    """
+
+    settings: Settings
+
+    @property
+    def reproducible(self) -> bool:
+        return True
+
+    @property
+    def score_name(self) -> str:
+        return self.settings.score_name
+
+    def init(self, ctx: Context) -> None:
+        model = ctx.get_model()
+
+        self.hflm = HFLM(
+            pretrained=model.model,  # ty:ignore[invalid-argument-type]
+            tokenizer=model.tokenizer,  # ty:ignore[invalid-argument-type]
+            batch_size="auto",
+        )
+
+    def get_score(self, ctx: Context) -> Score:
+        # The purpose of this hack, where we initialize the HFLM object once,
+        # then update its internal model every time we calculate the score,
+        # is to get the benefits of batch size caching while allowing for
+        # model reloads, e.g. when using --evaluate-model.
+        model = ctx.get_model()
+        self.hflm.pretrained = model.model
+        self.hflm._model = model.model
+
+        results = lm_eval.simple_evaluate(
+            model=self.hflm,
+            tasks=[self.settings.task],
+        )
+
+        benchmark_score = float(
+            results["results"][self.settings.task][self.settings.metric]
+        )
+
+        return Score(
+            value=benchmark_score,
+            rich_display=f"[bold]{benchmark_score:.4f}[/]",
+            md_display=f"{benchmark_score:.4f}",
+        )

src/heretic/scorers/keyword_rate.py

@@ -7,7 +7,7 @@ from heretic.config import DatasetSpecification
 from heretic.scorer import Context, Score, Scorer
 from heretic.utils import print
 
-DEFAULT_KEYWORD_MARKERS = [
+REFUSAL_MARKERS = [
     "disclaimer",
     "sorry",
     "i can'",
@@ -45,8 +45,13 @@ DEFAULT_KEYWORD_MARKERS = [
 
 
 class Settings(BaseModel):
+    score_name: str = Field(
+        default="Refusals",
+        description="Name that describes what the configured keyword rate measures.",
+    )
+
     keyword_markers: list[str] = Field(
-        default=DEFAULT_KEYWORD_MARKERS,
+        default=REFUSAL_MARKERS,
         description="Strings whose presence in a response (case insensitive) identifies the response as a keyword match.",
     )
 
@@ -80,12 +85,12 @@ class KeywordRate(Scorer):
 
     @property
     def score_name(self) -> str:
-        return "Keywords"
+        return self.settings.score_name
 
     def init(self, ctx: Context) -> None:
         print()
         print(
-            f"Loading KeywordRate evaluation prompts from [bold]{self.settings.prompts.dataset}[/]..."
+            f"Loading {self.settings.score_name} evaluation prompts from [bold]{self.settings.prompts.dataset}[/]..."
         )
         self.prompts = ctx.load_prompts(self.settings.prompts)
         print(f"* [bold]{len(self.prompts)}[/] prompts loaded")
@@ -113,7 +118,7 @@ class KeywordRate(Scorer):
 
         return Score(
             value=float(match_count / len(self.prompts)),
-            rich_display=f"{match_count}/{len(self.prompts)}",
+            rich_display=f"[bold]{match_count}[/]/{len(self.prompts)}",
             md_display=f"{match_count}/{len(self.prompts)}",
         )
 

src/heretic/scorers/kl_divergence.py

@@ -42,7 +42,7 @@ class KLDivergence(Scorer):
     def init(self, ctx: Context) -> None:
         print()
         print(
-            f"Loading KLDivergence evaluation prompts from [bold]{self.settings.prompts.dataset}[/]..."
+            f"Loading KL divergence evaluation prompts from [bold]{self.settings.prompts.dataset}[/]..."
         )
         self.prompts = ctx.load_prompts(self.settings.prompts)
         print(f"* [bold]{len(self.prompts)}[/] prompts loaded")
@@ -55,21 +55,23 @@ class KLDivergence(Scorer):
     def get_score(self, ctx: Context) -> Score:
         logits = ctx.get_logits(self.prompts)
         logprobs = F.log_softmax(logits, dim=-1)
-        kl = F.kl_div(
+
+        kl_divergence = F.kl_div(
             logprobs,
             self._baseline_logprobs,
             reduction="batchmean",
             log_target=True,
         ).item()
+
         return Score(
-            value=kl,
-            rich_display=f"{kl:.4f}",
-            md_display=f"{kl:.4f}",
+            value=kl_divergence,
+            rich_display=f"[bold]{kl_divergence:.4f}[/]",
+            md_display=f"{kl_divergence:.4f}",
         )
 
     def get_baseline_score(self, ctx: Context) -> Score:
         return Score(
             value=0,
-            rich_display="0 (by definition)",
+            rich_display="[bold]0[/] [italic](by definition)[/]",
             md_display="0 *(by definition)*",
         )

src/heretic/utils.py

@@ -43,8 +43,6 @@ T = TypeVar("T")
 
 print = Console(highlight=False).print
 
-T = TypeVar("T")
-
 
 def deep_merge_dicts(base: dict[str, Any], override: dict[str, Any]) -> dict[str, Any]:
     """

tests/minicpm5/config.toml

@@ -9,7 +9,6 @@ print_debug_information = true
 
 batch_size = 2
 max_response_length = 10
-kl_divergence_target = 0
 n_trials = 2
 n_startup_trials = 1
 
@@ -21,11 +20,6 @@ save_directory = "model"
 
 row_normalization = "none"
 
-scorers = [
-    { plugin = "heretic.scorers.keyword_rate.KeywordRate", optimization = "minimize" },
-    { plugin = "heretic.scorers.kl_divergence.KLDivergence", optimization = "minimize" },
-]
-
 [good_prompts]
 dataset = "mlabonne/harmless_alpaca"
 commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f"

tests/qwen2.5/config.toml

@@ -9,7 +9,6 @@ print_debug_information = true
 
 batch_size = 2
 max_response_length = 10
-kl_divergence_target = 0
 n_trials = 2
 n_startup_trials = 1
 
@@ -21,11 +20,6 @@ save_directory = "model"
 
 row_normalization = "pre"
 
-scorers = [
-    { plugin = "heretic.scorers.keyword_rate.KeywordRate", optimization = "minimize" },
-    { plugin = "heretic.scorers.kl_divergence.KLDivergence", optimization = "minimize" },
-]
-
 [good_prompts]
 dataset = "mlabonne/harmless_alpaca"
 commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f"

tests/run_tests.py

@@ -23,7 +23,9 @@ script_directory = Path(__file__).resolve().parent
 
 project_directory = script_directory.parent
 
-tests_failed = False
+# For tracking failures as (test_name, [failed_files]) and successful runs.
+failed_tests: list[tuple[str, list[str]]] = []
+passed_tests: list[str] = []
 
 for test_directory in script_directory.iterdir():
     if test_directory.is_dir():
@@ -65,6 +67,8 @@ for test_directory in script_directory.iterdir():
 
                             valid_hashes[filename].append(sha256.lower())
 
+            # Track which specific files failed within this test directory.
+            failed_files: list[str] = []
             for filename in valid_hashes:
                 sha256 = get_file_sha256(test_directory / "model" / filename)
 
@@ -79,9 +83,20 @@ for test_directory in script_directory.iterdir():
                             f"{sha256}\n"
                         )
                     )
-                    tests_failed = True
+                    failed_files.append(filename)
 
-if tests_failed:
+            if failed_files:
+                failed_tests.append((test_directory.name, failed_files))
+            else:
+                passed_tests.append(test_directory.name)
+
+if failed_tests:
+    print("#" * 50)
+    print("Summary of test failures:")
+    for test_name, files in failed_tests:
+        files_str = ", ".join(files)
+        print(f"- {test_name} (failed files: {files_str})")
+    print("#" * 50)
     sys.exit("Tests failed.")
 else:
     print("All tests passed.")

uv.lock

@@ -1036,7 +1036,7 @@ wheels = [
 
 [[package]]
 name = "heretic-llm"
-version = "1.4.0"
+version = "2.0.0.dev0"
 source = { editable = "." }
 dependencies = [
     { name = "accelerate" },
@@ -1990,7 +1990,7 @@ wheels = [
 
 [[package]]
 name = "nltk"
-version = "3.10.0"
+version = "3.10.3"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "click" },
@@ -1999,9 +1999,9 @@ dependencies = [
     { name = "regex" },
     { name = "tqdm" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/96/02/df4f105b28a7c16b0e41423bc09cf0f1b8a305df4ef0b10ca74a2e4c648c/nltk-3.10.0.tar.gz", hash = "sha256:4fbac1d98203cbcd1b5d94a2877fb822300072d80604a5e7fae49d2c5f84e8c1", size = 3089244, upload-time = "2026-07-08T02:39:13.562Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/e0/e6/fe51d2bb1a3b446f59c5c8165999a9fee208bc346af90a7cbf7657bc0d75/nltk-3.10.3.tar.gz", hash = "sha256:bb9327a461c3811c2fa4900e03840401f2126adfb30c0072827c433bd2444ea4", size = 5137152, upload-time = "2026-08-12T23:46:37.258Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/6e/89/a0b0f35e2820d6a99d75ea1c11977ee6d5c9e6658eceb45b0c7620881faa/nltk-3.10.0-py3-none-any.whl", hash = "sha256:54ff84d4916d3ef127e8953bee0023f6a6b320b75d634a19e06ef056d3d244bf", size = 1716144, upload-time = "2026-07-08T02:39:09.753Z" },
+    { url = "https://files.pythonhosted.org/packages/b6/6d/ebd2af4640b12168fdf0cb74b6118df2f32a2f62ec7e0c06fbfd80706639/nltk-3.10.3-py3-none-any.whl", hash = "sha256:ff9598a8e20518ee0d557745890cc4435b9578489e2dcbc69c4f81fa060caf7c", size = 1798643, upload-time = "2026-08-12T23:44:13.478Z" },
 ]
 
 [[package]]