feat: add modifier base class

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
@@ -137,8 +137,6 @@ system_prompt = "You are a helpful assistant."
 # or a path to a plain text file with one prompt per line (empty lines are ignored).
 # For text files, "column" is ignored and "split" is optional; when given, it selects
 # a subset of the lines using slice notation (e.g. "[:400]").
-# "config" specifies a dataset's specific config/subset name (e.g. "english", "hindi").
-# Leave unset for datasets with a single configuration.
 
 # Dataset of prompts that tend to not result in refusals (used for calculating residual directions).
 [good_prompts]

config.piqa.toml

@@ -2,6 +2,6 @@
 # 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"},
+    { plugin = "heretic.scorers.keyword_rate.KeywordRate", optimization = "minimize" },
+    { plugin = "heretic.scorers.benchmark_score.BenchmarkScore", optimization = "maximize" },
 ]

src/heretic/config.py

@@ -54,14 +54,6 @@ class DatasetSpecification(BaseModel):
         description="Hugging Face commit hash of the dataset.",
     )
 
-    config: str | None = Field(
-        default=None,
-        description=(
-            "Dataset config/subset name. Each config can have its own split. "
-            "Used to load a specific config of a dataset that has multiple configurations."
-        ),
-    )
-
     split: str | None = Field(
         default=None,
         description="Portion of the dataset to use. Required for datasets, optional for plain text files.",

src/heretic/main.py

@@ -39,14 +39,13 @@ import logging
 import math
 import os
 import random
-import re
 import time
 import warnings
 from dataclasses import asdict
 from importlib.metadata import version
 from os.path import commonprefix
 from pathlib import Path
-from typing import Any, cast
+from typing import Any
 
 import huggingface_hub
 import lm_eval
@@ -66,7 +65,6 @@ from optuna.storages.journal import JournalFileBackend, JournalFileOpenLock
 from optuna.trial import FrozenTrial, TrialState, create_trial
 from pydantic import ValidationError
 from questionary import Choice, Style
-from rich.markup import escape
 from rich.table import Table
 from rich.text import Text
 from rich.traceback import install
@@ -479,88 +477,40 @@ def run():
         print()
         print("Checking for common response prefix...")
         prefix_check_prompts = good_prompts[:100] + bad_prompts[:100]
+        responses = model.get_responses_batched(prefix_check_prompts)
 
-        # Detect if the model's chat template inserts a reasoning tag on its own
-        # at the end of user's prompt (e.g. <think>) by using a dummy prompt.
-        # If found, then we use the full closed CoT as the response prefix.
-        # LiquidAI's LFM models do this (Lfm2ForCausalLM).
+        # Despite being located in os.path, commonprefix actually performs
+        # a naive string operation without any path-specific logic,
+        # which is exactly what we need here. Trailing spaces are removed
+        # to avoid issues where multiple different tokens that all start
+        # with a space character lead to the common prefix ending with
+        # a space, which would result in an uncommon tokenization.
+        settings.response_prefix = commonprefix(responses).rstrip(" ")
 
-        # This cast is valid because str is the return type
-        # for a single chat operation with tokenize=False.
-        dummy_prompt = cast(
-            str,
-            model.tokenizer.apply_chat_template(
-                [{"role": "user", "content": "This is a dummy prompt."}],
-                add_generation_prompt=True,
-                tokenize=False,
-            ),
-        )
+        if settings.response_prefix:
+            print(f"* Prefix found: [bold]{settings.response_prefix!r}[/]")
 
-        cot_skip_applied = False
+            for cot_initializer, closed_cot_block in settings.chain_of_thought_skips:
+                if settings.response_prefix.startswith(cot_initializer):
+                    settings.response_prefix = closed_cot_block
+                    print(
+                        f"* Closed Chain-of-Thought block: [bold]{settings.response_prefix!r}[/]"
+                    )
 
-        for cot_initializer, closed_cot_block in settings.chain_of_thought_skips:
-            # Match the tag and ignore any whitespace characters following it at the end
-            # (if any), including spaces, tabs, and linebreaks. This is required for models
-            # having whitespaces after the tags.
-            pattern = rf"{re.escape(cot_initializer)}\s*$"
-            match = re.search(pattern, dummy_prompt)
-
-            if match:
-                # We use only the closed CoT block here. Any whitespaces
-                # will be handled by the 'Rechecking with prefix' logic below.
-                settings.response_prefix = closed_cot_block
-                print(
-                    f"* Closed Chain-of-Thought block: [bold]{escape(repr(settings.response_prefix))}[/]"
-                )
-                cot_skip_applied = True
-                break
-
-        # Fallback to inference for models like mistral-3 which are specifically
-        # instructed to generate thinking tags using the system prompt in their
-        # chat template, instead of inserting a prefix tag (e.g. <think>) at
-        # the end of user prompt like the case above. We expect the model to
-        # generate those tags.
-        if settings.response_prefix is None:
-            responses = model.get_responses_batched(prefix_check_prompts)
-
-            # Despite being located in os.path, commonprefix actually performs
-            # a naive string operation without any path-specific logic,
-            # which is exactly what we need here. Trailing spaces are removed
-            # to avoid issues where multiple different tokens that all start
-            # with a space character lead to the common prefix ending with
-            # a space, which would result in an uncommon tokenization.
-            settings.response_prefix = commonprefix(responses).rstrip(" ")
-
-            if settings.response_prefix:
-                print(
-                    f"* Prefix found: [bold]{escape(repr(settings.response_prefix))}[/]"
-                )
-
-                for (
-                    cot_initializer,
-                    closed_cot_block,
-                ) in settings.chain_of_thought_skips:
-                    if settings.response_prefix.startswith(cot_initializer):
-                        settings.response_prefix = closed_cot_block
+                    # When using a Chain-of-Thought skip, we need to check that the prefix
+                    # is actually complete (e.g. not missing a trailing newline).
+                    print("* Rechecking with prefix...")
+                    responses = model.get_responses_batched(prefix_check_prompts)
+                    additional_prefix = commonprefix(responses).rstrip(" ")
+                    if additional_prefix:
+                        settings.response_prefix += additional_prefix
                         print(
-                            f"* Closed Chain-of-Thought block: [bold]{escape(repr(settings.response_prefix))}[/]"
+                            f"* Extended prefix found: [bold]{settings.response_prefix!r}[/]"
                         )
-                        cot_skip_applied = True
-                        break
-            else:
-                print("* None found")
 
-        if cot_skip_applied:
-            # When using a Chain-of-Thought skip, we need to check that the prefix
-            # is actually complete (e.g. not missing a trailing newline).
-            print("* Rechecking with prefix...")
-            responses = model.get_responses_batched(prefix_check_prompts)
-            additional_prefix = commonprefix(responses).rstrip(" ")
-            if additional_prefix:
-                settings.response_prefix += additional_prefix
-                print(
-                    f"* Extended prefix found: [bold]{escape(repr(settings.response_prefix))}[/]"
-                )
+                    break
+        else:
+            print("* None found")
 
     evaluator = Evaluator(settings, model)
 

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

@@ -41,9 +41,11 @@ class BenchmarkScore(Scorer):
         return self.settings.score_name
 
     def init(self, ctx: Context) -> None:
+        model = ctx.get_model()
+
         self.hflm = HFLM(
-            pretrained=ctx._model.model,  # ty:ignore[invalid-argument-type]
-            tokenizer=ctx._model.tokenizer,  # ty:ignore[invalid-argument-type]
+            pretrained=model.model,  # ty:ignore[invalid-argument-type]
+            tokenizer=model.tokenizer,  # ty:ignore[invalid-argument-type]
             batch_size="auto",
         )
 
@@ -52,8 +54,9 @@ class BenchmarkScore(Scorer):
         # 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.
-        self.hflm.pretrained = ctx._model.model
-        self.hflm._model = ctx._model.model
+        model = ctx.get_model()
+        self.hflm.pretrained = model.model
+        self.hflm._model = model.model
 
         results = lm_eval.simple_evaluate(
             model=self.hflm,

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]:
     """
@@ -208,7 +206,6 @@ def load_prompts(
                     )
             dataset = load_dataset(
                 path,
-                name=specification.config,
                 revision=specification.commit,
                 split=split_str,
             )
@@ -226,7 +223,6 @@ def load_prompts(
             # Path should be a local directory.
             dataset = load_dataset(
                 path,
-                name=specification.config,
                 split=split_str,
                 # Don't require the number of examples (lines) per split to be pre-defined.
                 verification_mode=VerificationMode.NO_CHECKS,

tests/mistral-3/SHA256SUMS.windows

@@ -1,7 +1,7 @@
 72f84af4ea36b82409c35e31b584361534305ef7c0d90fce20d0dc38a7efead8 *chat_template.jinja
 e4c5278b361c57621253c27a2c3db358e1580aec8a14be8e19d4420a224137cf *config.json
 8dde85c000ae807be907421465826c7c63a39f6acf6d04a5a84efaf116ed4ef7 *generation_config.json
-20b5a820b38438202c64e4fc9807bd19e29678bebd678d29b2ee2d2f5bf71587 *model.safetensors
+29aff97d5633dead9e1ccd29a2cc153b4b7431d22f63c8d6cf60bc6547681cc9 *model.safetensors
 20e7a6dcde0a6f60ea3b4fb08f6f7afa62532dda93a3111e28384ba5150575f9 *processor_config.json
 c3a8d92e371b92a2cd6e678e31ebc27d0235e929a51fbf290f74742b341fa96f *tokenizer.json
 60a8042e29b4b20e884e48375aa1b9ac0025547371d50e60f6d55e6a9675e868 *tokenizer_config.json

tests/qwen3.5-moe/SHA256SUMS.ci

@@ -1,7 +1,7 @@
 a4aee8afcf2e0711942cf848899be66016f8d14a889ff9ede07bca099c28f715 *chat_template.jinja
 749b56d1b1e08081981169db6f2c44ab0be4fd6ebb452d15baafa5e09c21586a *config.json
 4625d1d64d41d1fa9dae7af4ba1e1d7e65a194073d4efa58acb266a916eaaa74 *generation_config.json
-2b3e575ac065f11ae5d4a7c3740efccbed294b646f1645239191ee8393354e03 *model.safetensors
+5fb94c65bcd9d736735a45e50c2b0bfafd3bb09a444c49b8cff2e131ed35797e *model.safetensors
 01562eddd6f9e9ec4bc31656a3b7055284cafbf889acc6c4348dca431ae31f68 *processor_config.json
 87a7830d63fcf43bf241c3c5242e96e62dd3fdc29224ca26fed8ea333db72de4 *tokenizer.json
 2e31d1126e81bddf8d15c3f95260fb487b48c5131b24fcbb5bb9d2537e7afac0 *tokenizer_config.json

tests/qwen3.5-moe/SHA256SUMS.windows

@@ -1,7 +1,7 @@
 a92e1dd97cb1cb175c9b70c0828e146bea4371c2643319b661b777e89811972e *chat_template.jinja
 b75e911805663da79fb9fbbbcc917b8f1a285d2da54d95c2c63ea7c1ffe9a05a *config.json
 2cbd9df0e99570efcced23b8d777bdf1fc692efda54b21eb59ad56ade76c9db6 *generation_config.json
-6061519a9595326df41abcdd093892463793d4d026d6fd23548f1792f622a252 *model.safetensors
+5f099b32807d0b84ed90765ca0ed53f8771da4738767bc1940486fec954570cf *model.safetensors
 0c29f9491e769aabbc389ad5912127cf6d9d5fceda2db8767f73d48131348c81 *processor_config.json
 87a7830d63fcf43bf241c3c5242e96e62dd3fdc29224ca26fed8ea333db72de4 *tokenizer.json
 4796e48d790a26d65f167bec8fc742beaa71f79f9468a6cd8b3ffa97f6e2a198 *tokenizer_config.json