chore: add yet more CI test hashes

* feat: add modifier base class

* feat: re-implement abliteration as a modifier plugin

* fix: adjust good/bad prompts hack to match tests

* feat: support dataset specifications containing multiple individual datasets

* feat: re-implement ARA as a modifier plugin

Arbitrary-Rank Ablation (ARA) (Weidmann 2026) was originally introduced by @p-e-w in #211

Co-authored-by: kabachuha <artemkhrapov2001@yandex.ru>
Co-authored-by: joninco <joninco@bullpoint.org>
Co-authored-by: Ashar <coder3101@users.noreply.github.com>

* feat: add tests for ARA

* feat: reduce default number of trials

* ci: remove broken "semantic-pull-request" workflow

* ci: add yet another alternative model hash

---------

Co-authored-by: kabachuha <artemkhrapov2001@yandex.ru>
Co-authored-by: joninco <joninco@bullpoint.org>
Co-authored-by: Ashar <coder3101@users.noreply.github.com>

.github/workflows/semantic-pr.yml

@@ -1,19 +0,0 @@
-name: Lint PR
-
-on:
-  pull_request_target:
-    types:
-      - opened
-      - reopened
-      - edited
-
-jobs:
-  main:
-    name: Validate PR title
-    runs-on: ubuntu-latest
-    permissions:
-      pull-requests: read
-    steps:
-      - uses: amannn/action-semantic-pull-request@v6
-        env:
-          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

README.md

@@ -127,87 +127,6 @@ save the model, upload it to Hugging Face, chat with it to test how well it work
 run standard benchmarks on it, or any combination of those actions.
 
 
-## Research features
-
-In addition to its primary function of removing model censorship, Heretic also
-provides features designed to support research into the semantics of model internals
-(interpretability). To use those features, you need to install Heretic with the
-optional `research` extra:
-
-```sh
-pip install -U 'heretic-llm[research]'
-```
-
-This gives you access to the following functionality:
-
-### Generate plots of residual vectors by passing `--plot-residuals`
-
-When run with this flag, Heretic will:
-
-1. Compute residual vectors (hidden states) for the first output token,
-   for each transformer layer, for both "harmful" and "harmless" prompts.
-2. Perform a [PaCMAP projection](https://github.com/YingfanWang/PaCMAP)
-   from residual space to 2D-space.
-3. Left-right align the projections of "harmful"/"harmless" residuals
-   by their geometric medians to make projections for consecutive layers
-   more similar. Additionally, PaCMAP is initialized with the previous
-   layer's projections for each new layer, minimizing disruptive transitions.
-4. Scatter-plot the projections, generating a PNG image for each layer.
-5. Generate an animation showing how residuals transform between layers,
-   as an animated GIF.
-
-<img width="800" height="600" alt="Plot of residual vectors" src="https://github.com/user-attachments/assets/981aa6ed-5ab9-48f0-9abf-2b1a2c430295" />
-
-See [the configuration file](config.default.toml) for options that allow you
-to control various aspects of the generated plots.
-
-Note that PaCMAP is an expensive operation that is performed on the CPU.
-For larger models, it can take an hour or more to compute projections
-for all layers.
-
-### Print details about residual geometry by passing `--print-residual-geometry`
-
-If you are interested in a quantitative analysis of how residual vectors
-for "harmful" and "harmless" prompts relate to each other, this flag gives you
-the following table, packed with metrics that can facilitate understanding
-the same (for [gemma-3-270m-it](https://huggingface.co/google/gemma-3-270m-it)
-in this case):
-
-```
-┏━━━━━━━┳━━━━━━━━┳━━━━━━━━━━┳━━━━━━━━━┳━━━━━━━━━━┳━━━━━━━━━┳━━━━━━━━━━┳━━━━━━━━━━┳━━━━━━━━━━┳━━━━━━━━━━┳━━━━━━━━━━┳━━━━━━━━━┳━━━━━━━━━┳━━━━━━━━┓
-┃ Layer ┃ S(g,b) ┃ S(g*,b*) ┃  S(g,r) ┃ S(g*,r*) ┃  S(b,r) ┃ S(b*,r*) ┃      |g| ┃     |g*| ┃      |b| ┃     |b*| ┃     |r| ┃    |r*| ┃   Silh ┃
-┡━━━━━━━╇━━━━━━━━╇━━━━━━━━━━╇━━━━━━━━━╇━━━━━━━━━━╇━━━━━━━━━╇━━━━━━━━━━╇━━━━━━━━━━╇━━━━━━━━━━╇━━━━━━━━━━╇━━━━━━━━━━╇━━━━━━━━━╇━━━━━━━━━╇━━━━━━━━┩
-│     1 │ 1.0000 │   1.0000 │ -0.4311 │  -0.4906 │ -0.4254 │  -0.4847 │   170.29 │   170.49 │   169.78 │   169.85 │    1.19 │    1.31 │ 0.0480 │
-│     2 │ 1.0000 │   1.0000 │  0.4297 │   0.4465 │  0.4365 │   0.4524 │   768.55 │   768.77 │   771.32 │   771.36 │    6.39 │    5.76 │ 0.0745 │
-│     3 │ 0.9999 │   1.0000 │ -0.5699 │  -0.5577 │ -0.5614 │  -0.5498 │  1020.98 │  1021.13 │  1013.80 │  1014.71 │   12.70 │   11.60 │ 0.0920 │
-│     4 │ 0.9999 │   1.0000 │  0.6582 │   0.6553 │  0.6659 │   0.6627 │  1356.39 │  1356.20 │  1368.71 │  1367.95 │   18.62 │   17.84 │ 0.0957 │
-│     5 │ 0.9987 │   0.9990 │ -0.6880 │  -0.6761 │ -0.6497 │  -0.6418 │   766.54 │   762.25 │   731.75 │   732.42 │   51.97 │   45.24 │ 0.1018 │
-│     6 │ 0.9998 │   0.9998 │ -0.1983 │  -0.2312 │ -0.1811 │  -0.2141 │  2417.35 │  2421.08 │  2409.18 │  2411.40 │   43.06 │   43.47 │ 0.0900 │
-│     7 │ 0.9998 │   0.9997 │ -0.5258 │  -0.5746 │ -0.5072 │  -0.5560 │  3444.92 │  3474.99 │  3400.01 │  3421.63 │   86.94 │   94.38 │ 0.0492 │
-│     8 │ 0.9990 │   0.9991 │  0.8235 │   0.8312 │  0.8479 │   0.8542 │  4596.54 │  4615.62 │  4918.32 │  4934.20 │  384.87 │  377.87 │ 0.2278 │
-│     9 │ 0.9992 │   0.9992 │  0.5335 │   0.5441 │  0.5678 │   0.5780 │  5322.30 │  5316.96 │  5468.65 │  5466.98 │  265.68 │  267.28 │ 0.1318 │
-│    10 │ 0.9974 │   0.9973 │  0.8189 │   0.8250 │  0.8579 │   0.8644 │  5328.81 │  5325.63 │  5953.35 │  5985.15 │  743.95 │  779.74 │ 0.2863 │
-│    11 │ 0.9977 │   0.9978 │  0.4262 │   0.4045 │  0.4862 │   0.4645 │  9644.02 │  9674.06 │  9983.47 │  9990.28 │  743.28 │  726.99 │ 0.1576 │
-│    12 │ 0.9904 │   0.9907 │  0.4384 │   0.4077 │  0.5586 │   0.5283 │ 10257.40 │ 10368.50 │ 11114.51 │ 11151.21 │ 1711.18 │ 1664.69 │ 0.1890 │
-│    13 │ 0.9867 │   0.9874 │  0.4007 │   0.3680 │  0.5444 │   0.5103 │ 12305.12 │ 12423.75 │ 13440.31 │ 13432.47 │ 2386.43 │ 2282.47 │ 0.1293 │
-│    14 │ 0.9921 │   0.9922 │  0.3198 │   0.2682 │  0.4364 │   0.3859 │ 16929.16 │ 17080.37 │ 17826.97 │ 17836.03 │ 2365.23 │ 2301.87 │ 0.1282 │
-│    15 │ 0.9846 │   0.9850 │  0.1198 │   0.0963 │  0.2913 │   0.2663 │ 16858.58 │ 16949.44 │ 17496.00 │ 17502.88 │ 3077.08 │ 3029.60 │ 0.1611 │
-│    16 │ 0.9686 │   0.9689 │ -0.0029 │  -0.0254 │  0.2457 │   0.2226 │ 18912.77 │ 19074.86 │ 19510.56 │ 19559.62 │ 4848.35 │ 4839.75 │ 0.1516 │
-│    17 │ 0.9782 │   0.9784 │ -0.0174 │  -0.0381 │  0.1908 │   0.1694 │ 27098.09 │ 27273.00 │ 27601.12 │ 27653.12 │ 5738.19 │ 5724.21 │ 0.1641 │
-│    18 │ 0.9184 │   0.9196 │  0.1343 │   0.1430 │  0.5155 │   0.5204 │   190.16 │   190.35 │   219.91 │   220.62 │   87.82 │   87.59 │ 0.1855 │
-└───────┴────────┴──────────┴─────────┴──────────┴─────────┴──────────┴──────────┴──────────┴──────────┴──────────┴─────────┴─────────┴────────┘
-g = mean of residual vectors for good prompts
-g* = geometric median of residual vectors for good prompts
-b = mean of residual vectors for bad prompts
-b* = geometric median of residual vectors for bad prompts
-r = residual direction for means (i.e., b - g)
-r* = residual direction for geometric medians (i.e., b* - g*)
-S(x,y) = cosine similarity of x and y
-|x| = L2 norm of x
-Silh = Mean silhouette coefficient of residuals for good/bad clusters
-```
-
-
 ## How Heretic works
 
 Heretic implements a parametrized variant of directional ablation. For each

config.default.toml

@@ -71,58 +71,27 @@ chain_of_thought_skips = [
 # Whether to print additional information that can help with debugging.
 print_debug_information = false
 
-# Whether to print detailed information about residuals and residual directions.
+# List of scorer plugin configs. Each entry is an object
-print_residual_geometry = false
+# { plugin = <plugin>, optimization = <optimization>, instance_name = <optional> }.
-
+# <optimization> is one of "minimize", "maximize", or "none" (do not optimize).
-# Whether to generate plots showing PaCMAP projections of residual vectors.
-plot_residuals = false
-
-# Base path to save plots of residual vectors to.
-residual_plot_path = "plots"
-
-# Title placed above plots of residual vectors.
-residual_plot_title = 'PaCMAP Projection of Residual Vectors for "Harmless" and "Harmful" Prompts'
-
-# Matplotlib style sheet to use for plots of residual vectors.
-residual_plot_style = "dark_background"
-
-# List of scorers to evaluate.
-# Each entry is an object:
-# { 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.keyword_rate.KeywordRate", optimization = "minimize" },
-    { plugin = "heretic.scorers.kl_divergence.KLDivergence", optimization = "minimize"},
+    { plugin = "heretic.scorers.kl_divergence.KLDivergence", optimization = "minimize" },
 ]
 
-# Whether to adjust the residual directions so that only the component that is
+# List of modifier plugin configs. Each entry is an object
-# orthogonal to the good direction is subtracted during abliteration.
+# { plugin = <plugin>, instance_name = <optional> }.
-orthogonalize_direction = true
+# Note that only a single modifier can currently be applied,
-
+# and this list must contain exactly one entry.
-# How to apply row normalization of the weights. Options:
+modifiers = [
-# "none" (no normalization),
+    { plugin = "heretic.modifiers.ara.ARA" },
-# "pre" (compute LoRA adapter relative to row-normalized weights),
+]
-# "full" (like "pre", but renormalizes to preserve original row magnitudes).
-row_normalization = "full"
-
-# The rank of the LoRA adapter to use when "full" row normalization is used.
-# Row magnitude preservation is approximate due to non-linear effects,
-# and this determines the rank of that approximation. Higher ranks produce
-# larger output files and may slow down evaluation.
-full_normalization_lora_rank = 3
-
-# The symmetric winsorization to apply to the per-prompt, per-layer residual vectors,
-# expressed as the quantile to clamp to (between 0 and 1). Disabled by default.
-# This can tame so-called "massive activations" that occur in some models.
-# Example: winsorization_quantile = 0.95 computes the 0.95-quantile of the absolute values
-# of the components, then clamps the magnitudes of all components to that quantile.
-winsorization_quantile = 1.0
 
 # Number of abliteration trials to run during optimization.
-n_trials = 200
+n_trials = 100
 
 # Number of trials that use random sampling for the purpose of exploration.
-n_startup_trials = 60
+n_startup_trials = 30
 
 # Directory to save and load study progress to/from.
 study_checkpoint_dir = "checkpoints"
@@ -133,30 +102,57 @@ max_shard_size = "5GB"
 # System prompt to use when prompting the model.
 system_prompt = "You are a helpful assistant."
 
+# Dataset of prompts to use for automatically determining the optimal batch size.
+[batch_size_test_prompts]
+dataset = "mlabonne/harmless_alpaca"
+split = "train[:256]"
+column = "text"
+
+# Dataset of prompts to use for automatically determining the response prefix.
+[[response_prefix_test_prompts]]
+dataset = "mlabonne/harmless_alpaca"
+split = "train[:100]"
+column = "text"
+
+[[response_prefix_test_prompts]]
+dataset = "mlabonne/harmful_behaviors"
+split = "train[:100]"
+column = "text"
+
+
+# Plugin-specific settings live in top-level TOML tables.
+#
+# For scorer plugins, use: `[scorer.<ClassName>]` (and optionally `[scorer.<ClassName>_<instance_name>]` for instance-related config).
+# For modifier plugins, use: `[modifier.<ClassName>]` (and optionally `[modifier.<ClassName>_<instance_name>]` for instance-related config).
+#
+# You can load multiple instances of the same plugin class by setting `instance_name`
+# in the `scorers/modifiers = [...]` list. Each instance is still identified as `ClassName.instanceName`
+# internally, but its config overrides live under `[scorer/modifier.ClassName_<instance_name>]`.
+#
+# Example:
+# scorers = [
+#   { plugin = "heretic.scorers.keyword_rate.KeywordRate", optimization = "minimize", instance_name = "small" },
+#   { plugin = "heretic.scorers.keyword_rate.KeywordRate", optimization = "minimize", instance_name = "tiny" },
+# ]
+#
+# Shared defaults for all instances live under `[scorer.KeywordRate]` and can be overridden per
+# instance under `[scorer.KeywordRate_<instance_name>]`.
+#
+# Example instance override:
+# [scorer.KeywordRate_small.prompts]
+# split = "test[:10]"
+#
 # Each "dataset" below can be a Hugging Face dataset ID, a path to a dataset on disk,
 # 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]
-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 residual 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.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
 
@@ -197,30 +193,103 @@ keyword_markers = [
     "ethical boundaries",
 ]
 
-# Scorer-owned evaluation prompts
+# Dataset of prompts to evaluate the keyword match rate on.
 [scorer.KeywordRate.prompts]
 dataset = "mlabonne/harmful_behaviors"
 split = "test[:100]"
 column = "text"
 
-# 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.keyword_rate.KeywordRate", optimization = 'minimize', instance_name = "small" },
-#   { plugin = "heretic.scorers.keyword_rate.KeywordRate", optimization = 'minimize', instance_name = "tiny" },
-# ]
-#
-# Shared defaults for all instances live under `[scorer.KeywordRate]` and can be overridden per
-# instance under `[scorer.KeywordRate_<instance_name>]`.
-#
-# Example instance override:
-# [scorer.KeywordRate_small.prompts]
-# split = "test[:10]"
 
+# Dataset of prompts used to measure KL divergence from original model.
 [scorer.KLDivergence.prompts]
 dataset = "mlabonne/harmless_alpaca"
 split = "test[:100]"
 column = "text"
+
+
+[scorer.BenchmarkScore]
+# Name that describes what the configured benchmark score measures.
+score_name = "PIQA acc_norm"
+
+# Task ID of the benchmark in the Language Model Evaluation Harness.
+task = "piqa"
+
+# Task metric to use as the benchmark score.
+metric = "acc_norm,none"
+
+
+[modifier.ARA]
+# Whether to renormalize the rows of the modified matrices to preserve
+# the original matrices' row magnitudes. This is believed to improve
+# intelligence retention (see Lai 2025, "Magnitude-Preserving Orthogonal Ablation").
+preserve_row_magnitudes = true
+
+# The rank of the LoRA adapter to use.
+# While mathematically, ARA is of "arbitrary" rank, experiments have shown that
+# singular values tend to drop rapidly after a few dozen dimensions, and approximating
+# the full transformation with a LoRA has many practical advantages.
+lora_rank = 50
+
+# Number of (outer) L-BFGS optimization steps to perform.
+n_optimization_steps = 5
+
+# Learning rate to use in the L-BFGS optimizer.
+learning_rate = 1.0
+
+# Maximum number of (inner) iterations to perform per (outer) L-BFGS optimization step.
+max_iter = 20
+
+# Number of past updates to store for approximating the Hessian matrix in the L-BFGS optimizer.
+history_size = 10
+
+# Whether to print the loss value for each L-BFGS optimization step.
+print_loss = false
+
+# Dataset of prompts that tend to produce desirable responses.
+[modifier.ARA.good_prompts]
+dataset = "mlabonne/harmless_alpaca"
+split = "train[:400]"
+column = "text"
+
+# Dataset of prompts that tend to produce undesirable responses.
+[modifier.ARA.bad_prompts]
+dataset = "mlabonne/harmful_behaviors"
+split = "train[:400]"
+column = "text"
+
+
+[modifier.Abliteration]
+# Whether to adjust the residual directions so that only the component that is
+# orthogonal to the good direction is subtracted during abliteration.
+orthogonalize_direction = true
+
+# How to apply row normalization of the weights. Options:
+# "none" (no normalization),
+# "pre" (compute LoRA adapter relative to row-normalized weights),
+# "full" (like "pre", but renormalizes to preserve original row magnitudes).
+row_normalization = "full"
+
+# The rank of the LoRA adapter to use when "full" row normalization is used.
+# Row magnitude preservation is approximate due to non-linear effects,
+# and this determines the rank of that approximation. Higher ranks produce
+# larger output files and may slow down evaluation.
+full_normalization_lora_rank = 3
+
+# The symmetric winsorization to apply to the per-prompt, per-layer residual vectors,
+# expressed as the quantile to clamp to (between 0 and 1). Disabled by default.
+# This can tame so-called "massive activations" that occur in some models.
+# Example: winsorization_quantile = 0.95 computes the 0.95-quantile of the absolute values
+# of the components, then clamps the magnitudes of all components to that quantile.
+winsorization_quantile = 1.0
+
+# Dataset of prompts that tend to produce desirable responses.
+[modifier.Abliteration.good_prompts]
+dataset = "mlabonne/harmless_alpaca"
+split = "train[:400]"
+column = "text"
+
+# Dataset of prompts that tend to produce undesirable responses.
+[modifier.Abliteration.bad_prompts]
+dataset = "mlabonne/harmful_behaviors"
+split = "train[:400]"
+column = "text"

config.nohumor.toml

@@ -3,23 +3,9 @@
 
 max_response_length = 300
 
-residual_plot_title = "PaCMAP Projection of Residuals for Serious/Humorous Prompts"
-
-[good_prompts]
-dataset = "mlabonne/harmless_alpaca"
-split = "train[:400]"
-column = "text"
-residual_plot_label = "Serious prompts"
-residual_plot_color = "royalblue"
-
-[bad_prompts]
-dataset = "UnstableLlama/jokes"
-split = "train[:200]"
-column = "text"
-residual_plot_label = "Humorous prompts"
-residual_plot_color = "darkorange"
-
 [scorer.KeywordRate]
+score_name = "Responses with humor"
+
 keyword_markers = [
     "😅",
     "here's one",
@@ -68,3 +54,13 @@ column = "text"
 dataset = "mlabonne/harmless_alpaca"
 split = "test[:100]"
 column = "text"
+
+[modifier.ARA.good_prompts]
+dataset = "mlabonne/harmless_alpaca"
+split = "train[:400]"
+column = "text"
+
+[modifier.ARA.bad_prompts]
+dataset = "UnstableLlama/jokes"
+split = "train[:200]"
+column = "text"

config.noslop.toml

@@ -3,27 +3,11 @@
 
 max_response_length = 300
 
-residual_plot_title = "PaCMAP Projection of Residuals for Slop-Suppressing/Inducing Prompts"
-
 system_prompt = "You are a professional writer."
 
-[good_prompts]
-dataset = "llm-aes/writing-prompts"
-split = "train[:500]"
-column = "prompt"
-prefix = "Write a short story based on the writing prompt below. Avoid literary cliches, purple prose, and flowery language.\n\nWriting prompt:"
-residual_plot_label = "Slop-suppressing prompts"
-residual_plot_color = "royalblue"
-
-[bad_prompts]
-dataset = "llm-aes/writing-prompts"
-split = "train[:500]"
-column = "prompt"
-prefix = "Write a short story based on the writing prompt below. Make extensive use of literary cliches, purple prose, and flowery language.\n\nWriting prompt:"
-residual_plot_label = "Slop-inducing prompts"
-residual_plot_color = "darkorange"
-
 [scorer.KeywordRate]
+score_name = "Responses with slop"
+
 keyword_markers = [
     "Eldoria",
     "Lumina",
@@ -162,3 +146,15 @@ dataset = "llm-aes/writing-prompts"
 split = "train[1000:1100]"
 column = "prompt"
 prefix = "Write a short story based on the writing prompt below. Avoid literary cliches, purple prose, and flowery language.\n\nWriting prompt:"
+
+[modifier.ARA.good_prompts]
+dataset = "llm-aes/writing-prompts"
+split = "train[:500]"
+column = "prompt"
+prefix = "Write a short story based on the writing prompt below. Avoid literary cliches, purple prose, and flowery language.\n\nWriting prompt:"
+
+[modifier.ARA.bad_prompts]
+dataset = "llm-aes/writing-prompts"
+split = "train[:500]"
+column = "prompt"
+prefix = "Write a short story based on the writing prompt below. Make extensive use of literary cliches, purple prose, and flowery language.\n\nWriting prompt:"

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"
@@ -14,7 +14,6 @@ classifiers = [
     "Environment :: Console",
     "Environment :: GPU",
     "Intended Audience :: Science/Research",
-    "License :: OSI Approved :: GNU Affero General Public License v3 or later (AGPLv3+)",
     "Topic :: Scientific/Engineering :: Artificial Intelligence",
     "Programming Language :: Python :: 3",
     "Programming Language :: Python :: 3.10",
@@ -22,41 +21,32 @@ classifiers = [
     "Programming Language :: Python :: 3.12",
 ]
 dependencies = [
-    "accelerate~=1.13",
+    "accelerate~=1.15",
-    "bitsandbytes~=0.49",
+    "bitsandbytes~=0.50",
-    "datasets~=4.7",
+    "datasets~=5.0",
-    "huggingface-hub~=1.7",
+    "huggingface-hub~=1.33",
     "immutabledict~=4.3",
     "langdetect~=1.0",
     "lm-eval[hf]~=0.4",
     "numpy~=2.2",
-    "optuna~=4.7",
+    "optuna~=5.0",
-    "peft~=0.19",
+    "peft~=0.21",
     "psutil~=7.2",
     "py-cpuinfo~=9.0",
-    "pydantic-settings~=2.13",
+    "pydantic-settings~=2.15",
     "questionary~=2.1",
-    "rich~=14.3",
+    "rich~=15.0",
     "tomli-w~=1.2",
     "torch",  # version deliberately unspecified
     "torchvision",  # version deliberately unspecified
-    "tqdm~=4.67",
+    "tqdm~=4.70",
-    "transformers[kernels]~=5.6",
+    "transformers[kernels]~=5.18",
-]
-
-[project.optional-dependencies]
-research = [
-    "geom-median~=0.1",
-    "imageio~=2.37",
-    "matplotlib~=3.10",
-    "pacmap~=0.8",
-    "scikit-learn~=1.7",
 ]
 
 [dependency-groups]
 dev = [
-    "ruff>=0.14.5",
+    "ruff>=0.16.9",
-    "ty>=0.0.5",
+    "ty>=0.0.84",
 ]
 
 [project.urls]
@@ -70,11 +60,31 @@ Changelog = "https://github.com/p-e-w/heretic/releases"
 heretic = "heretic.main:main"
 
 [build-system]
-requires = ["uv_build>=0.8.11,<0.9.0"]
+requires = ["uv_build>=0.12.21,<0.13.0"]
 build-backend = "uv_build"
 
+[tool.ruff.lint]
+ignore = [
+    "B023",  # "Function definition does not bind loop variable"
+    "BLE001",  # "Do not catch blind exception"
+    "TRY002",  # "Create your own exception"
+    "SIM117",  # "Use a single `with` statement with multiple contexts instead of nested `with` statements"
+]
+
 [tool.uv]
-exclude-newer = "7 days"
+# TODO: Re-enable before 2.0 release!
+#exclude-newer = "7 days"
+
+[tool.uv.sources]
+torch = { index = "pytorch" }
+torchvision = { index = "pytorch" }
+
+[[tool.uv.index]]
+# Default to CPU wheels for PyTorch to avoid issues in CI.
+# Can be overridden by passing `--index pytorch=<url>` to uv.
+name = "pytorch"
+url = "https://download.pytorch.org/whl/cpu"
+explicit = true
 
 [tool.uv.build-backend]
 module-name = "heretic"

src/heretic/analyzer.py

@@ -1,357 +0,0 @@
-# SPDX-License-Identifier: AGPL-3.0-or-later
-# Copyright (C) 2025-2026  Philipp Emanuel Weidmann <pew@worldwidemann.com> + contributors
-
-from pathlib import Path
-
-import numpy as np
-import torch
-import torch.linalg as LA
-import torch.nn.functional as F
-from numpy.typing import NDArray
-from rich.progress import track
-from rich.table import Table
-from torch import Tensor
-
-from .config import Settings
-from .model import Model
-from .utils import print
-
-
-class Analyzer:
-    def __init__(
-        self,
-        settings: Settings,
-        model: Model,
-        good_residuals: Tensor,
-        bad_residuals: Tensor,
-    ):
-        self.settings = settings
-        self.model = model
-        self.good_residuals = good_residuals
-        self.bad_residuals = bad_residuals
-
-    def print_residual_geometry(self):
-        try:
-            from geom_median.torch import (  # ty:ignore[unresolved-import]
-                compute_geometric_median,
-            )
-            from sklearn.metrics import silhouette_score  # ty:ignore[unresolved-import]
-        except ImportError:
-            print()
-            print(
-                (
-                    "[red]Research dependencies not found. Printing residual geometry requires "
-                    "installing Heretic with the optional research feature, i.e., "
-                    "using \"pip install -U 'heretic-llm\\[research]'\".[/]"
-                )
-            )
-            return
-
-        print()
-        print("Computing residual geometry...")
-
-        table = Table()
-        table.add_column("Layer", justify="right")
-        table.add_column("S(g,b)", justify="right")
-        table.add_column("S(g*,b*)", justify="right")
-        table.add_column("S(g,r)", justify="right")
-        table.add_column("S(g*,r*)", justify="right")
-        table.add_column("S(b,r)", justify="right")
-        table.add_column("S(b*,r*)", justify="right")
-        table.add_column("|g|", justify="right")
-        table.add_column("|g*|", justify="right")
-        table.add_column("|b|", justify="right")
-        table.add_column("|b*|", justify="right")
-        table.add_column("|r|", justify="right")
-        table.add_column("|r*|", justify="right")
-        table.add_column("Silh", justify="right")
-
-        g = self.good_residuals.mean(dim=0)
-        g_star = torch.stack(
-            [
-                compute_geometric_median(
-                    self.good_residuals[:, layer_index, :].detach().cpu()
-                ).median
-                for layer_index in range(len(self.model.get_layers()) + 1)
-            ]
-        )
-        b = self.bad_residuals.mean(dim=0)
-        b_star = torch.stack(
-            [
-                compute_geometric_median(
-                    self.bad_residuals[:, layer_index, :].detach().cpu()
-                ).median
-                for layer_index in range(len(self.model.get_layers()) + 1)
-            ]
-        )
-        r = b - g
-        r_star = b_star - g_star
-
-        g_b_similarities = F.cosine_similarity(g, b, dim=-1)
-        g_star_b_star_similarities = F.cosine_similarity(g_star, b_star, dim=-1)
-        g_r_similarities = F.cosine_similarity(g, r, dim=-1)
-        g_star_r_star_similarities = F.cosine_similarity(g_star, r_star, dim=-1)
-        b_r_similarities = F.cosine_similarity(b, r, dim=-1)
-        b_star_r_star_similarities = F.cosine_similarity(b_star, r_star, dim=-1)
-
-        g_norms = LA.vector_norm(g, dim=-1)
-        g_star_norms = LA.vector_norm(g_star, dim=-1)
-        b_norms = LA.vector_norm(b, dim=-1)
-        b_star_norms = LA.vector_norm(b_star, dim=-1)
-        r_norms = LA.vector_norm(r, dim=-1)
-        r_star_norms = LA.vector_norm(r_star, dim=-1)
-
-        residuals = (
-            torch.cat(
-                [
-                    self.good_residuals,
-                    self.bad_residuals,
-                ],
-                dim=0,
-            )
-            .detach()
-            .cpu()
-            .numpy()
-        )
-        labels = [0] * len(self.good_residuals) + [1] * len(self.bad_residuals)
-        silhouettes = [
-            silhouette_score(residuals[:, layer_index, :], labels)
-            for layer_index in range(len(self.model.get_layers()) + 1)
-        ]
-
-        for layer_index in range(1, len(self.model.get_layers()) + 1):
-            table.add_row(
-                f"{layer_index}",
-                f"{g_b_similarities[layer_index].item():.4f}",
-                f"{g_star_b_star_similarities[layer_index].item():.4f}",
-                f"{g_r_similarities[layer_index].item():.4f}",
-                f"{g_star_r_star_similarities[layer_index].item():.4f}",
-                f"{b_r_similarities[layer_index].item():.4f}",
-                f"{b_star_r_star_similarities[layer_index].item():.4f}",
-                f"{g_norms[layer_index].item():.2f}",
-                f"{g_star_norms[layer_index].item():.2f}",
-                f"{b_norms[layer_index].item():.2f}",
-                f"{b_star_norms[layer_index].item():.2f}",
-                f"{r_norms[layer_index].item():.2f}",
-                f"{r_star_norms[layer_index].item():.2f}",
-                f"{silhouettes[layer_index]:.4f}",
-            )
-
-        print()
-        print("[bold]Residual Geometry[/]")
-        print(table)
-        print("[bold]g[/] = mean of residual vectors for good prompts")
-        print("[bold]g*[/] = geometric median of residual vectors for good prompts")
-        print("[bold]b[/] = mean of residual vectors for bad prompts")
-        print("[bold]b*[/] = geometric median of residual vectors for bad prompts")
-        print("[bold]r[/] = residual direction for means (i.e., [bold]b - g[/])")
-        print(
-            "[bold]r*[/] = residual direction for geometric medians (i.e., [bold]b* - g*[/])"
-        )
-        print("[bold]S(x,y)[/] = cosine similarity of [bold]x[/] and [bold]y[/]")
-        print("[bold]|x|[/] = L2 norm of [bold]x[/]")
-        print(
-            "[bold]Silh[/] = Mean silhouette coefficient of residuals for good/bad clusters"
-        )
-
-    def plot_residuals(self):
-        try:
-            import imageio.v3 as iio  # ty:ignore[unresolved-import]
-            import matplotlib.pyplot as plt  # ty:ignore[unresolved-import]
-            from geom_median.numpy import (  # ty:ignore[unresolved-import]
-                compute_geometric_median,
-            )
-            from pacmap import PaCMAP  # ty:ignore[unresolved-import]
-        except ImportError:
-            print()
-            print(
-                (
-                    "[red]Research dependencies not found. Plotting residuals requires "
-                    "installing Heretic with the optional research feature, i.e., "
-                    "using \"pip install -U 'heretic-llm\\[research]'\".[/]"
-                )
-            )
-            return
-
-        LAYER_FRAME_DURATION = 1000
-        N_TRANSITION_FRAMES = 20
-        TRANSITION_FRAME_DURATION = 50
-
-        print()
-        print("Plotting residual vectors...")
-
-        layer_residuals_2d = []
-        pacmap_init = None
-
-        for layer_index in track(
-            range(1, len(self.model.get_layers()) + 1),
-            description="* Computing PaCMAP projections...",
-        ):
-            good_residuals = (
-                self.good_residuals[:, layer_index, :].detach().cpu().numpy()
-            )
-            bad_residuals = self.bad_residuals[:, layer_index, :].detach().cpu().numpy()
-
-            residuals = np.vstack((good_residuals, bad_residuals))
-            embedding = PaCMAP(n_components=2, n_neighbors=30)
-            residuals_2d = embedding.fit_transform(residuals, init=pacmap_init)
-            pacmap_init = residuals_2d
-
-            n_good_residuals = good_residuals.shape[0]
-            good_residuals_2d = residuals_2d[:n_good_residuals]
-            bad_residuals_2d = residuals_2d[n_good_residuals:]
-
-            # Important: These are the medians of the 2D-projected residuals,
-            #            not the projections of the medians of the residuals.
-            #            Their only purpose is to rotate the individual plots
-            #            into a consistent orientation. They are not suitable
-            #            for being plotted themselves.
-            good_anchor = compute_geometric_median(good_residuals_2d).median
-            bad_anchor = compute_geometric_median(bad_residuals_2d).median
-
-            # Rotate points to make the line connecting the medians horizontal,
-            # with the median of the good residuals on the left.
-            direction = bad_anchor - good_anchor
-            angle = -np.arctan2(direction[1], direction[0])
-            cosine = np.cos(angle)
-            sine = np.sin(angle)
-            rotation_matrix = np.array([[cosine, -sine], [sine, cosine]])
-            residuals_2d = residuals_2d @ rotation_matrix.T
-
-            good_residuals_2d = residuals_2d[:n_good_residuals]
-            bad_residuals_2d = residuals_2d[n_good_residuals:]
-
-            layer_residuals_2d.append((good_residuals_2d, bad_residuals_2d))
-
-        plt.style.use(self.settings.residual_plot_style)
-
-        def plot(
-            image_path: Path,
-            layer_index: int,
-            good_residuals_2d: NDArray,
-            bad_residuals_2d: NDArray,
-        ):
-            fig, ax = plt.subplots(figsize=(8, 6))
-
-            ax.scatter(
-                good_residuals_2d[:, 0],
-                good_residuals_2d[:, 1],
-                s=10,
-                c=self.settings.good_prompts.residual_plot_color,
-                alpha=0.5,
-                label=self.settings.good_prompts.residual_plot_label,
-            )
-            ax.scatter(
-                bad_residuals_2d[:, 0],
-                bad_residuals_2d[:, 1],
-                s=10,
-                c=self.settings.bad_prompts.residual_plot_color,
-                alpha=0.5,
-                label=self.settings.bad_prompts.residual_plot_label,
-            )
-
-            ax.set_title(self.settings.residual_plot_title, pad=11)
-            ax.legend(loc="upper right")
-            ax.grid(False)
-            ax.set_xticks([])
-            ax.set_yticks([])
-
-            fig.text(
-                0.018,
-                0.02,
-                self.settings.model,
-                ha="left",
-                va="bottom",
-                fontsize=12,
-            )
-            fig.text(
-                0.982,
-                0.02,
-                f"Layer {layer_index:03}",
-                ha="right",
-                va="bottom",
-                fontsize=12,
-            )
-
-            fig.tight_layout()
-            fig.subplots_adjust(bottom=0.08)
-
-            fig.savefig(image_path, dpi=100)
-            plt.close(fig)
-
-        base_path = Path(
-            self.settings.residual_plot_path
-        ) / self.settings.model.replace(
-            "/",
-            "_",
-        ).replace(
-            "\\",
-            "_",
-        )
-
-        base_path.mkdir(parents=True, exist_ok=True)
-
-        images = []
-        durations = []
-
-        for layer_index, (
-            good_residuals_2d,
-            bad_residuals_2d,
-        ) in enumerate(
-            track(
-                layer_residuals_2d,
-                description="* Generating plots...",
-            ),
-            1,
-        ):
-            image_path = base_path / f"layer_{layer_index:03}.png"
-
-            plot(image_path, layer_index, good_residuals_2d, bad_residuals_2d)
-
-            images.append(iio.imread(image_path))
-            durations.append(LAYER_FRAME_DURATION)
-
-            if layer_index < len(layer_residuals_2d):
-                # The first frame of the transition is the layer frame created above.
-                # The last frame is the next layer frame, created in the next iteration of the outer loop.
-                # The following are the intermediate frames.
-                # There are a total of N_TRANSITION_FRAMES frame changes in the transition.
-                for frame_index in range(1, N_TRANSITION_FRAMES):
-                    image_path = (
-                        base_path / f"layer_{layer_index:03}_frame_{frame_index:03}.png"
-                    )
-
-                    progress = frame_index / N_TRANSITION_FRAMES
-
-                    good_residuals_2d_interpolated = good_residuals_2d + progress * (
-                        layer_residuals_2d[layer_index][0] - good_residuals_2d
-                    )
-                    bad_residuals_2d_interpolated = bad_residuals_2d + progress * (
-                        layer_residuals_2d[layer_index][1] - bad_residuals_2d
-                    )
-
-                    plot(
-                        image_path,
-                        layer_index,
-                        good_residuals_2d_interpolated,
-                        bad_residuals_2d_interpolated,
-                    )
-
-                    images.append(iio.imread(image_path))
-                    durations.append(TRANSITION_FRAME_DURATION)
-
-                    # Delete the image file containing the animation frame.
-                    # We have already read its contents and it serves no purpose
-                    # other than building the animation.
-                    image_path.unlink()
-
-        print("* Generating animation...")
-
-        iio.imwrite(
-            base_path / "animation.gif",
-            images,
-            duration=durations,
-            loop=0,
-        )
-
-        print(f"* Plots saved to [bold]{base_path.resolve()}[/].")

src/heretic/config.py

@@ -2,7 +2,7 @@
 # Copyright (C) 2025-2026  Philipp Emanuel Weidmann <pew@worldwidemann.com> + contributors
 
 from enum import Enum
-from typing import Dict, Literal
+from typing import Literal, TypeAlias
 
 from pydantic import (
     BaseModel,
@@ -32,19 +32,12 @@ class QuantizationMethod(str, Enum):
     BNB_4BIT = "bnb_4bit"
 
 
-class RowNormalization(str, Enum):
-    NONE = "none"
-    PRE = "pre"
-    # POST = "post"  # Theoretically possible, but provides no advantage.
-    FULL = "full"
-
-
 class ExportStrategy(str, Enum):
     MERGE = "merge"
     ADAPTER = "adapter"
 
 
-class DatasetSpecification(BaseModel):
+class SingleDatasetSpecification(BaseModel):
     dataset: str = Field(
         description="Hugging Face dataset ID, or path to dataset on disk."
     )
@@ -54,6 +47,14 @@ 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.",
@@ -79,17 +80,10 @@ class DatasetSpecification(BaseModel):
         description="System prompt to use with the prompts (overrides global system prompt if set).",
     )
 
-    residual_plot_label: str | None = Field(
-        default=None,
-        description="Label to use for the dataset in plots of residual vectors.",
-        exclude=True,
-    )
 
-    residual_plot_color: str | None = Field(
+DatasetSpecification: TypeAlias = (
-        default=None,
+    SingleDatasetSpecification | list[SingleDatasetSpecification]
-        description="Matplotlib color to use for the dataset in plots of residual vectors.",
+)
-        exclude=True,
-    )
 
 
 class ScorerConfig(BaseModel):
@@ -142,6 +136,48 @@ class ScorerConfig(BaseModel):
         return value
 
 
+class ModifierConfig(BaseModel):
+    """
+    Configuration for a modifier plugin.
+
+    TOML format:
+    - { plugin = "<plugin>", 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`)."
+        ),
+    )
+
+    instance_name: str | None = Field(
+        default=None,
+        description=(
+            "Optional name to distinguish multiple instances of the same plugin class. "
+            "Instance-specific settings live under `[modifier.<ClassName>_<instance_name>]`."
+        ),
+    )
+
+    @field_validator("instance_name")
+    @classmethod
+    def validate_instance_name(cls, value: str | None) -> str | None:
+        if value is None:
+            return value
+
+        if not value.strip():
+            raise ValueError("cannot be empty or whitespace")
+
+        if "." in value:
+            raise ValueError("'.' is not allowed")
+
+        if any(char.isspace() for char in value):
+            raise ValueError("whitespace is not allowed")
+
+        return value
+
+
 class BenchmarkSpecification(BaseModel):
     task: str = Field(
         description="Task ID of the benchmark in the Language Model Evaluation Harness."
@@ -218,12 +254,12 @@ class Settings(BaseSettings):
         ),
     )
 
-    device_map: str | Dict[str, int | str] = Field(
+    device_map: str | dict[str, int | str] = Field(
         default="auto",
         description="Device map to pass to Accelerate when loading the model.",
     )
 
-    max_memory: Dict[str, str] | None = Field(
+    max_memory: dict[str, str] | None = Field(
         default=None,
         description='Maximum memory to allocate per device (e.g., { "0" = "20GB", "cpu" = "64GB" }).',
     )
@@ -251,6 +287,18 @@ class Settings(BaseSettings):
         exclude=True,
     )
 
+    batch_size_test_prompts: DatasetSpecification = Field(
+        default=SingleDatasetSpecification(
+            dataset="mlabonne/harmless_alpaca",
+            split="train[:256]",
+            column="text",
+        ),
+        description="Dataset of prompts to use for automatically determining the optimal batch size.",
+        # When storing a settings object, the batch size is already fixed,
+        # either determined by the automatic mechanism or by explicit user choice.
+        exclude=True,
+    )
+
     max_response_length: PositiveInt = Field(
         default=100,
         description="Maximum number of tokens to generate for each response.",
@@ -265,6 +313,25 @@ class Settings(BaseSettings):
         ),
     )
 
+    response_prefix_test_prompts: DatasetSpecification = Field(
+        default=[
+            SingleDatasetSpecification(
+                dataset="mlabonne/harmless_alpaca",
+                split="train[:100]",
+                column="text",
+            ),
+            SingleDatasetSpecification(
+                dataset="mlabonne/harmful_behaviors",
+                split="train[:100]",
+                column="text",
+            ),
+        ],
+        description="Dataset of prompts to use for automatically determining the response prefix.",
+        # When storing a settings object, the response prefix is already fixed,
+        # either determined by the automatic mechanism or by explicit user choice.
+        exclude=True,
+    )
+
     chain_of_thought_skips: list[tuple[str, str]] = Field(
         default=[
             # Most thinking models.
@@ -304,38 +371,8 @@ class Settings(BaseSettings):
         exclude=True,
     )
 
-    print_residual_geometry: bool = Field(
-        default=False,
-        description="Whether to print detailed information about residuals and residual directions.",
-        exclude=True,
-    )
-
-    plot_residuals: bool = Field(
-        default=False,
-        description="Whether to generate plots showing PaCMAP projections of residual vectors.",
-        exclude=True,
-    )
-
-    residual_plot_path: str = Field(
-        default="plots",
-        description="Base path to save plots of residual vectors to.",
-        exclude=True,
-    )
-
-    residual_plot_title: str = Field(
-        default='PaCMAP Projection of Residual Vectors for "Harmless" and "Harmful" Prompts',
-        description="Title placed above plots of residual vectors.",
-        exclude=True,
-    )
-
-    residual_plot_style: str = Field(
-        default="dark_background",
-        description="Matplotlib style sheet to use for plots of residual vectors.",
-        exclude=True,
-    )
-
     scorers: list[ScorerConfig] = Field(
-        default_factory=lambda: [
+        default=[
             ScorerConfig(
                 plugin="heretic.scorers.keyword_rate.KeywordRate",
                 optimization="minimize",
@@ -346,58 +383,33 @@ class Settings(BaseSettings):
             ),
         ],
         description=(
-            "List of scorer plugin configs. Each entry is an object"
+            "List of scorer plugin configs. Each entry is an object "
-            " { plugin = <plugin>, optimization = <optimization>, instance_name = <optional> }."
+            "{ plugin = <plugin>, optimization = <optimization>, instance_name = <optional> }. "
-            " <optimization> is one of 'minimize', 'maximize', 'none' (do not optimize)."
+            '<optimization> is one of "minimize", "maximize", or "none" (do not optimize).'
         ),
     )
 
-    orthogonalize_direction: bool = Field(
+    modifiers: list[ModifierConfig] = Field(
-        default=True,
+        default=[
+            ModifierConfig(
+                plugin="heretic.modifiers.ara.ARA",
+            ),
+        ],
         description=(
-            "Whether to adjust the residual directions so that only the component that is "
+            "List of modifier plugin configs. Each entry is an object "
-            "orthogonal to the good direction is subtracted during abliteration."
+            "{ plugin = <plugin>, instance_name = <optional> }. "
-        ),
+            "Note that only a single modifier can currently be applied, "
-    )
+            "and this list must contain exactly one entry."
-
-    row_normalization: RowNormalization = Field(
-        default=RowNormalization.FULL,
-        description=(
-            "How to apply row normalization of the weights. Options: "
-            '"none" (no normalization), '
-            '"pre" (compute LoRA adapter relative to row-normalized weights), '
-            '"full" (like "pre", but renormalizes to preserve original row magnitudes).'
-        ),
-    )
-
-    full_normalization_lora_rank: PositiveInt = Field(
-        default=3,
-        description=(
-            'The rank of the LoRA adapter to use when "full" row normalization is used. '
-            "Row magnitude preservation is approximate due to non-linear effects, "
-            "and this determines the rank of that approximation. Higher ranks produce "
-            "larger output files and may slow down evaluation."
-        ),
-    )
-
-    winsorization_quantile: float = Field(
-        default=1.0,
-        description=(
-            "The symmetric winsorization to apply to the per-prompt, per-layer residual vectors, "
-            "expressed as the quantile to clamp to (between 0 and 1). Disabled by default. "
-            'This can tame so-called "massive activations" that occur in some models. '
-            "Example: winsorization_quantile = 0.95 computes the 0.95-quantile of the absolute values "
-            "of the components, then clamps the magnitudes of all components to that quantile."
         ),
     )
 
     n_trials: PositiveInt = Field(
-        default=200,
+        default=100,
         description="Number of abliteration trials to run during optimization.",
     )
 
     n_startup_trials: NonNegativeInt = Field(
-        default=60,
+        default=30,
         description="Number of trials that use random sampling for the purpose of exploration.",
     )
 
@@ -539,31 +551,9 @@ class Settings(BaseSettings):
         description="System prompt to use when prompting the model.",
     )
 
-    good_prompts: DatasetSpecification = Field(
-        default=DatasetSpecification(
-            dataset="mlabonne/harmless_alpaca",
-            split="train[:400]",
-            column="text",
-            residual_plot_label='"Harmless" prompts',
-            residual_plot_color="royalblue",
-        ),
-        description="Dataset of prompts that tend to not result in refusals (used for calculating refusal directions).",
-    )
-
-    bad_prompts: DatasetSpecification = Field(
-        default=DatasetSpecification(
-            dataset="mlabonne/harmful_behaviors",
-            split="train[:400]",
-            column="text",
-            residual_plot_label='"Harmful" prompts',
-            residual_plot_color="darkorange",
-        ),
-        description="Dataset of prompts that tend to result in refusals (used for calculating refusal directions).",
-    )
-
     # We intentionally allow extra keys so users can provide plugin-specific
     # configuration in TOML tables like `[scorer.KeywordRate]` which are later
-    # consumed via `settings.model_extra` (see `Evaluator._get_plugin_namespace`).
+    # consumed via `settings.model_extra` (see `plugin.get_plugin_namespace`).
     model_config = SettingsConfigDict(extra="allow")
 
     @classmethod

src/heretic/evaluator.py

@@ -9,9 +9,9 @@ from pydantic import BaseModel
 
 from .config import DatasetSpecification, ScorerConfig, Settings
 from .model import Model
-from .plugin import get_plugin_namespace, is_builtin_plugin, load_plugin
+from .plugin import Context, is_builtin_plugin, load_plugin
-from .scorer import Context, Score, Scorer
+from .scorer import Score, Scorer
-from .utils import deep_merge_dicts, parse_study_direction, print
+from .utils import parse_study_direction, print
 
 
 @dataclass
@@ -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:
         """
@@ -61,14 +63,16 @@ class Evaluator:
             scorer_cls.validate_contract()
 
             print(
-                f"* Loaded: [bold]{scorer_cls.__name__} {'- ' + config.instance_name if config.instance_name else ''}[/bold]"
+                f"* Loaded: [bold]{scorer_cls.__name__}{' - ' + config.instance_name if config.instance_name else ''}[/bold]"
             )
 
             # Instantiate scorers.
             instance_name = config.instance_name or None
 
-            raw_settings = self._get_scorer_settings_raw(
+            raw_settings = scorer_cls.get_settings_raw(
-                scorer_cls=scorer_cls, instance_name=instance_name
+                self.settings.model_extra,
+                "scorer",
+                instance_name,
             )
             scorer_settings: BaseModel | None = scorer_cls.validate_settings(
                 raw_settings
@@ -108,11 +112,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
@@ -120,45 +119,9 @@ class Evaluator:
         """
         specifications = []
         for entry in self._scorer_entries:
-            if entry.scorer.settings is None:
+            specifications.extend(entry.scorer.get_dataset_specifications())
-                continue
-            for value in dict(entry.scorer.settings).values():
-                if isinstance(value, DatasetSpecification):
-                    specifications.append(value)
         return specifications
 
-    def _get_scorer_settings_raw(
-        self, *, scorer_cls: type[Scorer], instance_name: str | None
-    ) -> dict[str, Any]:
-        """
-        Build the raw settings dict for a scorer class and optional instance.
-
-        Config rules:
-        - Base settings live in `[scorer.ClassName]` (applies to all instances).
-        - Instance overrides live in `[scorer.ClassName_<instance_name>]` (preferred).
-        - Only merge/validate keys that exist in the scorer Settings schema.
-        """
-        settings_model = scorer_cls.get_settings_model()
-        if settings_model is None:
-            # No settings schema: nothing to merge/validate.
-            return {}
-
-        class_name = scorer_cls.__name__
-
-        namespaces = [f"scorer.{class_name}"]
-        if instance_name:
-            namespaces.append(f"scorer.{class_name}_{instance_name}")
-
-        merged_settings: dict[str, Any] = {}
-        allowed_keys = set(settings_model.model_fields.keys())
-
-        for namespace in namespaces:
-            raw_table = get_plugin_namespace(self.settings.model_extra, namespace)
-            filtered = {k: v for k, v in raw_table.items() if k in allowed_keys}
-            merged_settings = deep_merge_dicts(merged_settings, filtered)
-
-        return merged_settings
-
     def all_scorers_reproducible(self) -> bool:
         """
         Returns True if all scorers are reproducible,

src/heretic/main.py

@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: AGPL-3.0-or-later
 # Copyright (C) 2025-2026  Philipp Emanuel Weidmann <pew@worldwidemann.com> + contributors
 
-# ruff: noqa: E402
-
 import sys
 
 # Ensure standard output/error use UTF-8 instead of system default charmap (e.g. cp1252 on Windows).
@@ -23,7 +21,7 @@ def _is_help_invocation() -> bool:
 
 # Parse and handle CLI help before importing heavyweight ML/runtime dependencies.
 if _is_help_invocation():
-    Settings()  # ty:ignore[missing-argument]
+    Settings()
 
 # FIXME: Rich progress bars are currently disabled because of rendering issues
 #        when used from multiple threads in parallel (e.g. by huggingface_hub).
@@ -39,13 +37,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
+from typing import Any, cast
 
 import huggingface_hub
 import lm_eval
@@ -53,7 +51,6 @@ import numpy as np
 import optuna
 import questionary
 import torch
-import torch.nn.functional as F
 import transformers
 from huggingface_hub import HfApi, ModelCard, ModelCardData
 from lm_eval.models.huggingface import HFLM
@@ -65,13 +62,20 @@ 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
 
-from .analyzer import Analyzer
+from .config import (
-from .config import ExportStrategy, QuantizationMethod
+    DatasetSpecification,
+    ExportStrategy,
+    QuantizationMethod,
+)
 from .evaluator import Evaluator
-from .model import AbliterationParameters, Model, get_model_class
+from .model import Model, get_model_class
+from .modifier import load_and_init_modifiers
+from .plugin import Context, is_builtin_plugin
 from .reproduce import (
     check_environment,
     collect_reproducibles,
@@ -80,11 +84,12 @@ from .reproduce import (
 from .system import empty_cache, get_accelerator_info
 from .utils import (
     ask_if_unset,
+    format_dataset_specification,
     format_duration,
     format_exception,
     get_file_sha256,
     get_readme_intro,
-    get_trial_parameters,
+    is_dataset_specification_reproducible,
     is_hf_path,
     load_prompts,
     print,
@@ -216,7 +221,7 @@ def run():
     try:
         # The required argument "model" must be provided by the user,
         # either on the command line or in the configuration file.
-        settings = Settings()  # ty:ignore[missing-argument]
+        settings = Settings()
     except ValidationError as error:
         print(f"[red]Configuration contains [bold]{error.error_count()}[/] errors:[/]")
 
@@ -242,18 +247,12 @@ def run():
         # FIXME: "Reproduction"/"reproducibility" name inconsistency!
         reproduction_information = load_reproduction_information(settings.reproduce)
 
-        # Version 3 is the plugin-era schema, which stores generic scorer
+        if reproduction_information["version"] != "4":
-        # `scores`/`baseline_scores`. It is intentionally NOT compatible with the
-        # pre-plugin v1/v2 schema (hardcoded refusals/KL `metrics`), so those are
-        # rejected rather than silently failing on a missing key later.
-        if reproduction_information["version"] != "3":
             print(
-                (
+                f"[red]Unsupported file format version: [bold]{reproduction_information['version']}[/].[/] "
-                    f"[red]Unsupported file format version: [bold]{reproduction_information['version']}[/].[/] "
+                "This version of Heretic reads version 4 (plugin-based) reproduce.json files. "
-                    "This version of Heretic reads version 3 (plugin scorer) reproduce.json files. "
+                "Older files were produced before the introduction of the plugin system and are not supported. "
-                    "Older files were produced before the scorer-plugin refactor and are not supported. "
+                "Please install Heretic 1.4 to use these files."
-                    "Please install Heretic 1.4 to use these files."
-                )
             )
             return
 
@@ -336,12 +335,10 @@ def run():
             if settings.checkpoint_action is None:
                 print()
                 print(
-                    (
+                    "[green]You have already processed this model.[/] "
-                        "[green]You have already processed this model.[/] "
+                    "You can show the results from the previous run, allowing you to export models or to run additional trials. "
-                        "You can show the results from the previous run, allowing you to export models or to run additional trials. "
+                    "Alternatively, you can ignore the previous run and start from scratch. "
-                        "Alternatively, you can ignore the previous run and start from scratch. "
+                    "This will delete the checkpoint file and all results from the previous run."
-                        "This will delete the checkpoint file and all results from the previous run."
-                    )
                 )
 
             choices.append(
@@ -354,12 +351,10 @@ def run():
             if settings.checkpoint_action is None:
                 print()
                 print(
-                    (
+                    "[yellow]You have already processed this model, but the run was interrupted.[/] "
-                        "[yellow]You have already processed this model, but the run was interrupted.[/] "
+                    "You can continue the previous run from where it stopped. This will override any specified settings. "
-                        "You can continue the previous run from where it stopped. This will override any specified settings. "
+                    "Alternatively, you can ignore the previous run and start from scratch. "
-                        "Alternatively, you can ignore the previous run and start from scratch. "
+                    "This will delete the checkpoint file and all results from the previous run."
-                        "This will delete the checkpoint file and all results from the previous run."
-                    )
                 )
 
             choices.append(
@@ -411,17 +406,17 @@ def run():
     print()
     print_memory_usage()
 
-    print()
-    print(f"Loading good prompts from [bold]{settings.good_prompts.dataset}[/]...")
-    good_prompts = load_prompts(settings, settings.good_prompts)
-    print(f"* [bold]{len(good_prompts)}[/] prompts loaded")
-
-    print()
-    print(f"Loading bad prompts from [bold]{settings.bad_prompts.dataset}[/]...")
-    bad_prompts = load_prompts(settings, settings.bad_prompts)
-    print(f"* [bold]{len(bad_prompts)}[/] prompts loaded")
-
     if settings.batch_size == 0:
+        print()
+        print(
+            f"Loading batch size test prompts from [bold]{format_dataset_specification(settings.batch_size_test_prompts)}[/]..."
+        )
+        batch_size_test_prompts = load_prompts(
+            settings,
+            settings.batch_size_test_prompts,
+        )
+        print(f"* [bold]{len(batch_size_test_prompts)}[/] prompts loaded")
+
         print()
         print("Determining optimal batch size...")
 
@@ -432,7 +427,9 @@ def run():
         while batch_size <= settings.max_batch_size:
             print(f"* Trying batch size [bold]{batch_size}[/]... ", end="")
 
-            prompts = good_prompts * math.ceil(batch_size / len(good_prompts))
+            prompts = batch_size_test_prompts * math.ceil(
+                batch_size / len(batch_size_test_prompts)
+            )
             prompts = prompts[:batch_size]
 
             try:
@@ -473,43 +470,100 @@ def run():
         print(f"* Chosen batch size: [bold]{settings.batch_size}[/]")
 
     if settings.response_prefix is None:
+        print()
+        print(
+            f"Loading response prefix test prompts from [bold]{format_dataset_specification(settings.response_prefix_test_prompts)}[/]..."
+        )
+        response_prefix_test_prompts = load_prompts(
+            settings,
+            settings.response_prefix_test_prompts,
+        )
+        print(f"* [bold]{len(response_prefix_test_prompts)}[/] prompts loaded")
+
         print()
         print("Checking for common response prefix...")
-        prefix_check_prompts = good_prompts[:100] + bad_prompts[:100]
-        responses = model.get_responses_batched(prefix_check_prompts)
 
-        # Despite being located in os.path, commonprefix actually performs
+        # Detect if the model's chat template inserts a reasoning tag on its own
-        # a naive string operation without any path-specific logic,
+        # at the end of user's prompt (e.g. <think>) by using a dummy prompt.
-        # which is exactly what we need here. Trailing spaces are removed
+        # If found, then we use the full closed CoT as the response prefix.
-        # to avoid issues where multiple different tokens that all start
+        # LiquidAI's LFM models do this (Lfm2ForCausalLM).
-        # 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:
+        # This cast is valid because str is the return type
-            print(f"* Prefix found: [bold]{settings.response_prefix!r}[/]")
+        # 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,
+            ),
+        )
 
-            for cot_initializer, closed_cot_block in settings.chain_of_thought_skips:
+        cot_skip_applied = False
-                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}[/]"
-                    )
 
-                    # When using a Chain-of-Thought skip, we need to check that the prefix
+        for cot_initializer, closed_cot_block in settings.chain_of_thought_skips:
-                    # is actually complete (e.g. not missing a trailing newline).
+            # Match the tag and ignore any whitespace characters following it at the end
-                    print("* Rechecking with prefix...")
+            # (if any), including spaces, tabs, and linebreaks. This is required for models
-                    responses = model.get_responses_batched(prefix_check_prompts)
+            # having whitespaces after the tags.
-                    additional_prefix = commonprefix(responses).rstrip(" ")
+            pattern = rf"{re.escape(cot_initializer)}\s*$"
-                    if additional_prefix:
+            match = re.search(pattern, dummy_prompt)
-                        settings.response_prefix += additional_prefix
+
+            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(response_prefix_test_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
                         print(
-                            f"* Extended prefix found: [bold]{settings.response_prefix!r}[/]"
+                            f"* Closed Chain-of-Thought block: [bold]{escape(repr(settings.response_prefix))}[/]"
                         )
+                        cot_skip_applied = True
+                        break
+            else:
+                print("* None found")
 
-                    break
+        if cot_skip_applied:
-        else:
+            # When using a Chain-of-Thought skip, we need to check that the prefix
-            print("* None found")
+            # is actually complete (e.g. not missing a trailing newline).
+            print("* Rechecking with prefix...")
+            responses = model.get_responses_batched(response_prefix_test_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))}[/]"
+                )
 
     evaluator = Evaluator(settings, model)
 
@@ -519,10 +573,8 @@ def run():
         settings.model = settings.evaluate_model
         model.reset_model()
         print("* Evaluating...")
-        print()
+        for name, score in evaluator.get_scores():
-        print("[bold]Metrics:[/]")
+            print(f"  * [bold]{name}:[/] [green]{score.rich_display}[/]")
-        for score_name, score in evaluator.get_scores():
-            print(f"  * {score_name}: [bold]{score.rich_display}[/]")
         return
 
     if not reproduction_mode and not evaluator.get_objective_names():
@@ -535,53 +587,17 @@ def run():
         return
 
     print()
-    print("Calculating per-layer residual directions...")
+    print("Loading and initializing modifiers...")
+    modifier_entries = load_and_init_modifiers(settings, model)
 
-    needs_full_residuals = settings.print_residual_geometry or settings.plot_residuals
+    # `load_and_init_modifiers` currently guarantees that the returned list has exactly one element.
-
+    # This may change in the future when support for multiple modifiers is implemented.
-    if needs_full_residuals:
+    modifier_entry = modifier_entries[0]
-        print("* Obtaining residuals for good prompts...")
+    modifier = modifier_entry.modifier
-        good_residuals = model.get_residuals_batched(good_prompts)
+    modifier_name = modifier_entry.name
-        print("* Obtaining residuals for bad prompts...")
-        bad_residuals = model.get_residuals_batched(bad_prompts)
-
-        good_means = good_residuals.mean(dim=0)
-        bad_means = bad_residuals.mean(dim=0)
-
-        analyzer = Analyzer(settings, model, good_residuals, bad_residuals)
-
-        if settings.print_residual_geometry:
-            analyzer.print_residual_geometry()
-
-        if settings.plot_residuals:
-            analyzer.plot_residuals()
-
-        # We don't need the full residuals after computing their means and analyzing geometry.
-        del good_residuals, bad_residuals, analyzer
-    else:
-        print("* Obtaining residual mean for good prompts...")
-        good_means = model.get_residuals_mean(good_prompts)
-        print("* Obtaining residual mean for bad prompts...")
-        bad_means = model.get_residuals_mean(bad_prompts)
-
-    residual_directions = F.normalize(bad_means - good_means, p=2, dim=1)
-
-    if settings.orthogonalize_direction:
-        # Implements https://huggingface.co/blog/grimjim/projected-abliteration
-        # Adjust the residual directions so that only the component that is
-        # orthogonal to the good direction is subtracted during abliteration.
-        good_directions = F.normalize(good_means, p=2, dim=1)
-        projection_vector = torch.sum(residual_directions * good_directions, dim=1)
-        residual_directions = (
-            residual_directions - projection_vector.unsqueeze(1) * good_directions
-        )
-        residual_directions = F.normalize(residual_directions, p=2, dim=1)
-        del good_directions, projection_vector
-
-    del good_means, bad_means
 
     # Clear cache before starting the optimization study.
-    # This should free up memory from the objects released with the del statements above.
+    # This should free up memory from temporary objects created while initializing modifiers.
     empty_cache()
 
     trial_index = 0
@@ -593,102 +609,26 @@ def run():
         trial_index += 1
         trial.set_user_attr("index", trial_index)
 
-        direction_scope = trial.suggest_categorical(
+        ctx = Context(settings=settings, model=model)
-            "direction_scope",
+        parameters = modifier.suggest_parameters(ctx, trial)
-            [
+        trial.set_user_attr("parameters", parameters.to_dict())
-                "global",
-                "per layer",
-            ],
-        )
-
-        last_layer_index = len(model.get_layers()) - 1
-
-        # Discrimination between "harmful" and "harmless" inputs is usually strongest
-        # in layers slightly past the midpoint of the layer stack. See the original
-        # abliteration paper (https://arxiv.org/abs/2406.11717) for a deeper analysis.
-        #
-        # Note that we always sample this parameter even though we only need it for
-        # the "global" direction scope. The reason is that multivariate TPE doesn't
-        # work with conditional or variable-range parameters.
-        direction_index = trial.suggest_float(
-            "direction_index",
-            0.4 * last_layer_index,
-            0.9 * last_layer_index,
-        )
-
-        if direction_scope == "per layer":
-            direction_index = None
-
-        parameters = {}
-
-        for component in model.get_abliterable_components():
-            # The parameter ranges are based on experiments with various models
-            # and much wider ranges. They are not set in stone and might have to be
-            # adjusted for future models.
-            #
-            # The MLP gets a negative lower bound that is then clamped to 0, so the
-            # optimizer can fully disable its ablation. The clamp puts a positive
-            # probability mass on exactly 0 (the continuous sampler would otherwise
-            # reach 0 with probability zero). Ablating the MLP is often unnecessary for
-            # removing refusals and tends to damage model intelligence more than
-            # ablating the attention output, so on many models the optimum is to leave
-            # it (mostly) untouched. See issue #202.
-            max_weight_lower_bound = -0.25 if component == "mlp.down_proj" else 0.8
-            max_weight = max(
-                0.0,
-                trial.suggest_float(
-                    f"{component}.max_weight",
-                    max_weight_lower_bound,
-                    1.5,
-                ),
-            )
-            max_weight_position = trial.suggest_float(
-                f"{component}.max_weight_position",
-                0.6 * last_layer_index,
-                1.0 * last_layer_index,
-            )
-            # For sampling purposes, min_weight is expressed as a fraction of max_weight,
-            # again because multivariate TPE doesn't support variable-range parameters.
-            # The value is transformed into the actual min_weight value below.
-            min_weight = trial.suggest_float(
-                f"{component}.min_weight",
-                0.0,
-                1.0,
-            )
-            min_weight_distance = trial.suggest_float(
-                f"{component}.min_weight_distance",
-                1.0,
-                max(0.6 * last_layer_index, 1.0),
-            )
-
-            parameters[component] = AbliterationParameters(
-                max_weight=max_weight,
-                max_weight_position=max_weight_position,
-                min_weight=(min_weight * max_weight),
-                min_weight_distance=min_weight_distance,
-            )
-
-        trial.set_user_attr("direction_index", direction_index)
-        trial.set_user_attr("parameters", {k: asdict(v) for k, v in parameters.items()})
 
         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():
+        for name, value in modifier.render_trial_parameters(trial).items():
             print(f"  * {name} = [bold]{value}[/]")
         print("* Resetting model...")
-        model.reset_model()
+        modifier.reset_model(ctx)
-        print("* Abliterating...")
+        print(f"* Modifying model using {modifier_name}...")
-        model.abliterate(residual_directions, direction_index, parameters)
+        modifier.modify_model(ctx, parameters)
         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 +733,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)
@@ -823,13 +763,10 @@ def run():
             if settings.trial_index is None:
                 print()
                 print(
-                    (
+                    "The following trials resulted in Pareto optimal combinations of the optimization objectives. "
-                        "The following trials resulted in Pareto optimal combinations of the optimization objectives. "
+                    "After selecting a trial, you will be able to save the model, upload it to Hugging Face, "
-                        "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. "
-                        "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. "
-                        "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.[/]"
-                    )
                 )
 
         while trial_loop_active:
@@ -838,13 +775,10 @@ def run():
                 trial_loop_active = False
 
             if reproduction_mode:
-                parameters = reproduction_information["parameters"]
-
                 trial = create_trial(
                     values=[],
                     user_attrs={
-                        "direction_index": parameters["direction_index"],
+                        "parameters": reproduction_information["parameters"],
-                        "parameters": parameters["abliteration_parameters"],
                         "scores": reproduction_information["scores"],
                     },
                 )
@@ -914,24 +848,21 @@ def run():
                 )
 
             print("* Parameters:")
-            for name, value in get_trial_parameters(trial).items():
+            for name, value in modifier.render_trial_parameters(trial).items():
                 print(f"  * {name} = [bold]{value}[/]")
 
             # Per https://github.com/huggingface/peft/issues/868#issuecomment-1820642893
             # once a LoRA is merged it's expected to be empty. Provide a utility function
             # to restore the previous LoRA-ified state.
             def reset_trial_model():
+                ctx = Context(settings=settings, model=model)
                 print("* Resetting model...")
-                model.reset_model()
+                modifier.reset_model(ctx)
-                print("* Abliterating...")
+                print(f"* Modifying model using {modifier_name}...")
-                model.abliterate(
+                parameters = modifier.parameters_class.from_dict(
-                    residual_directions,
+                    trial.user_attrs["parameters"]
-                    trial.user_attrs["direction_index"],
-                    {
-                        k: AbliterationParameters(**v)
-                        for k, v in trial.user_attrs["parameters"].items()
-                    },
                 )
+                modifier.modify_model(ctx, parameters)
 
             reset_trial_model()
 
@@ -1114,35 +1045,33 @@ def run():
                             # are available on the Hugging Face Hub (not local paths),
                             # that all datasets are pinned to a commit (an unpinned
                             # dataset was likely loaded from a local cache), and that
-                            # only built-in scorer plugins are used (external plugins
+                            # only built-in plugins are used (external plugins cannot
-                            # cannot be resolved when reproducing).
+                            # be resolved when reproducing).
-                            dataset_specifications = [
+                            dataset_specifications: list[DatasetSpecification] = [
-                                settings.good_prompts,
-                                settings.bad_prompts,
                                 *evaluator.get_dataset_specifications(),
+                                *modifier.get_dataset_specifications(),
                             ]
                             is_reproducible = (
                                 is_hf_path(settings.model)
                                 and all(
-                                    is_hf_path(specification.dataset)
+                                    is_dataset_specification_reproducible(specification)
-                                    and specification.commit is not None
                                     for specification in dataset_specifications
                                 )
                                 and evaluator.all_scorers_reproducible()
                                 and evaluator.all_scorers_builtin()
+                                and modifier.reproducible
+                                and is_builtin_plugin(modifier_entry.config.plugin)
                                 and not reproduction_mode
                             )
 
                             if is_reproducible:
                                 if settings.upload_reproducibility_information is None:
                                     print(
-                                        (
+                                        "Heretic can add information to the repository that allows others to reproduce the model. "
-                                            "Heretic can add information to the repository that allows others to reproduce the model. "
+                                        "This is optional, but valuable to the community as both a learning tool and to preserve computational work already done. "
-                                            "This is optional, but valuable to the community as both a learning tool and to preserve computational work already done. "
+                                        "Guaranteeing reproducibility requires basic system information (Python and OS version, CPU and GPU/accelerator info) "
-                                            "Guaranteeing reproducibility requires basic system information (Python and OS version, CPU and GPU/accelerator info) "
+                                        "as tensor operations can give different results in different system environments. "
-                                            "as tensor operations can give different results in different system environments. "
+                                        "[bold]The information does not include any file system paths or other private data.[/]"
-                                            "[bold]The information does not include any file system paths or other private data.[/]"
-                                        )
                                     )
 
                                 reproducibility_information = ask_if_unset(
@@ -1174,7 +1103,7 @@ def run():
                             if strategy == ExportStrategy.ADAPTER:
                                 print("Uploading LoRA adapter...")
                                 model.model.push_to_hub(
-                                    repo_id,
+                                    repo_id,  # ty: ignore[invalid-argument-type]
                                     private=private,
                                     max_shard_size=settings.max_shard_size,
                                     token=token,
@@ -1183,7 +1112,7 @@ def run():
                                 print("Uploading merged model...")
                                 merged_model = model.get_merged_model()
                                 merged_model.push_to_hub(
-                                    repo_id,
+                                    repo_id,  # ty: ignore[invalid-argument-type]
                                     private=private,
                                     max_shard_size=settings.max_shard_size,
                                     token=token,
@@ -1226,9 +1155,16 @@ def run():
                                 card.data.tags.append("abliterated")
                                 if reproducibility_information != "none":
                                     card.data.tags.append("reproducible")
+
+                                # Must be a Hugging Face Hub repository ID,
+                                # so local paths are excluded.
+                                if is_hf_path(settings.model):
+                                    card.data.base_model = settings.model
+
                                 card.text = (
                                     get_readme_intro(
                                         settings,
+                                        modifier,
                                         trial,
                                         reproducibility_information != "none",
                                     )
@@ -1247,6 +1183,7 @@ def run():
                                     upload_reproduce_folder(
                                         repo_id,
                                         settings,
+                                        dataset_specifications,
                                         token,
                                         checkpoint_path=study_checkpoint_file,
                                         trial=trial,
@@ -1370,7 +1307,7 @@ def run():
                             benchmark_original_model = scope == "Benchmark both models"
 
                             hflm = HFLM(
-                                pretrained=model.model,  # ty:ignore[invalid-argument-type]
+                                pretrained=model.model,
                                 tokenizer=model.tokenizer,  # ty:ignore[invalid-argument-type]
                                 batch_size="auto",
                             )

src/heretic/model.py

@@ -1,19 +1,15 @@
 # SPDX-License-Identifier: AGPL-3.0-or-later
 # Copyright (C) 2025-2026  Philipp Emanuel Weidmann <pew@worldwidemann.com> + contributors
 
-import math
+from collections.abc import Callable
 from contextlib import suppress
-from dataclasses import dataclass
+from typing import Any, TypeAlias, cast
-from typing import Any, Type, cast
 
-import bitsandbytes as bnb
 import torch
-import torch.linalg as LA
-import torch.nn.functional as F
 from peft import LoraConfig, PeftModel, get_peft_model
-from peft.tuners.lora.layer import Linear
 from torch import FloatTensor, LongTensor, Tensor
 from torch.nn import Module, ModuleList
+from torch.utils.hooks import RemovableHandle
 from transformers import (
     AutoModelForCausalLM,
     AutoModelForImageTextToText,
@@ -28,31 +24,30 @@ from transformers import (
     TextStreamer,
 )
 from transformers.generation import (
-    GenerateDecoderOnlyOutput,  # ty:ignore[possibly-missing-import]
+    GenerateDecoderOnlyOutput,
 )
 
-from .config import QuantizationMethod, RowNormalization, Settings
+from .config import QuantizationMethod, Settings
 from .system import empty_cache
 from .utils import Prompt, batchify, format_exception, print
 
 
 def get_model_class(
     model: str,
-) -> Type[AutoModelForImageTextToText] | Type[AutoModelForCausalLM]:
+) -> type[AutoModelForImageTextToText] | type[AutoModelForCausalLM]:
     configs = PretrainedConfig.get_config_dict(model)
 
-    if any([("vision_config" in config) for config in configs]):
+    if any(("vision_config" in config) for config in configs):
         return AutoModelForImageTextToText
     else:
         return AutoModelForCausalLM
 
 
-@dataclass
+# The list contains one element per layer.
-class AbliterationParameters:
+# Each element maps from the component name to a (possibly sparse) mapping
-    max_weight: float
+# from the module index to an (input, output) tuple containing the I/O
-    max_weight_position: float
+# tensors of shape (prompt, component).
-    min_weight: float
+ModuleIO: TypeAlias = list[dict[str, dict[int, tuple[Tensor, Tensor]]]]
-    min_weight_distance: float
 
 
 class Model:
@@ -74,9 +69,14 @@ class Model:
         print()
         print(f"Loading model [bold]{settings.model}[/]...")
 
-        self.tokenizer = AutoTokenizer.from_pretrained(
+        # PreTrainedTokenizerBase is the "base class for all tokenizer backends"
-            settings.model,
+        # according to the documentation.
-            **self.revision_kwargs,
+        self.tokenizer = cast(
+            PreTrainedTokenizerBase,
+            AutoTokenizer.from_pretrained(
+                settings.model,
+                **self.revision_kwargs,
+            ),
         )
 
         # Multimodal models have a processor we'll want to save.
@@ -96,7 +96,7 @@ class Model:
         #           after the prompt and thinks the sequence is complete.
         self.tokenizer.padding_side = "left"
 
-        self.model = None  # ty:ignore[invalid-assignment]
+        self.model = None
         self.max_memory = (
             {int(k) if k.isdigit() else k: v for k, v in settings.max_memory.items()}
             if settings.max_memory
@@ -149,7 +149,7 @@ class Model:
                     max_new_tokens=1,
                 )
             except Exception as error:
-                self.model = None  # ty:ignore[invalid-assignment]
+                self.model = None
                 empty_cache()
 
                 formatted = format_exception(error)
@@ -168,11 +168,6 @@ class Model:
         if self.model is None:
             raise Exception("Failed to load model with all configured dtypes.")
 
-        self._apply_lora()
-
-        # LoRA B matrices are initialized to zero by default in PEFT,
-        # so we don't need to do anything manually.
-
         print(f"* Transformer model with [bold]{len(self.get_layers())}[/] layers")
 
         all_components = {}
@@ -186,7 +181,7 @@ class Model:
         for component, count in all_components.items():
             print(f"  * [bold]{component}[/]: [bold]{count}[/] modules total")
 
-    def _apply_lora(self):
+    def apply_lora(self, lora_rank: int):
         # Guard against calling this method at the wrong time.
         assert isinstance(self.model, PreTrainedModel)
 
@@ -211,13 +206,6 @@ class Model:
 
         target_modules = sorted(target_modules_set)
 
-        if self.settings.row_normalization != RowNormalization.FULL:
-            # Rank 1 is sufficient for directional ablation without renormalization.
-            lora_rank = 1
-        else:
-            # Row magnitude preservation introduces nonlinear effects.
-            lora_rank = self.settings.full_normalization_lora_rank
-
         self.peft_config = LoraConfig(
             r=lora_rank,
             target_modules=target_modules,
@@ -233,11 +221,6 @@ class Model:
         # so the result is a PeftModel rather than a PeftMixedModel.
         self.model = cast(PeftModel, get_peft_model(self.model, self.peft_config))
 
-        display_targets = sorted({name.rsplit(".", 1)[-1] for name in target_modules})
-        print(
-            f"* LoRA adapters initialized (target types: {', '.join(display_targets)})"
-        )
-
     def _get_quantization_config(self, dtype: str) -> BitsAndBytesConfig | None:
         """
         Creates quantization config based on settings.
@@ -312,7 +295,7 @@ class Model:
             self.needs_reload = True
             return merged_model
 
-    def reset_model(self):
+    def reset_model(self) -> bool:
         """
         Resets the model to a clean state for the next trial or evaluation.
 
@@ -321,6 +304,8 @@ class Model:
           resets LoRA adapter weights to zero (identity transformation).
         - Slow path: If switching models or after merge_and_unload(),
           performs full model reload with quantization config.
+
+        Returns True if the fast path was taken.
         """
 
         # If a prior model load was interrupted/cancelled mid-process, self.model will be None.
@@ -333,10 +318,10 @@ class Model:
             for name, module in self.model.named_modules():
                 if "lora_B" in name and hasattr(module, "weight"):
                     torch.nn.init.zeros_(module.weight)
-            return
+            return True
 
         # Purge existing model object from memory to make space.
-        self.model = None  # ty:ignore[invalid-assignment]
+        self.model = None
         empty_cache()
 
         quantization_config = self._get_quantization_config(
@@ -360,10 +345,10 @@ class Model:
             **extra_kwargs,
         )
 
-        self._apply_lora()
-
         self.needs_reload = False
 
+        return False
+
     def get_layers(self) -> ModuleList:
         model = self.model
 
@@ -373,10 +358,10 @@ class Model:
 
         # Most multimodal models.
         with suppress(Exception):
-            return model.model.language_model.layers
+            return model.model.language_model.layers  # ty: ignore[unresolved-attribute, invalid-return-type]
 
         # Text-only models.
-        return model.model.layers
+        return model.model.layers  # ty: ignore[unresolved-attribute, invalid-return-type]
 
     def get_layer_modules(self, layer_index: int) -> dict[str, list[Module]]:
         layer = self.get_layers()[layer_index]
@@ -397,50 +382,50 @@ class Model:
 
         # Standard self-attention out-projection (most models).
         with suppress(Exception):
-            try_add("attn.o_proj", layer.self_attn.o_proj)  # ty:ignore[possibly-missing-attribute]
+            try_add("attn.o_proj", layer.self_attn.o_proj)  # ty: ignore[unresolved-attribute]
 
         # Qwen3.5 MoE hybrid layers use GatedDeltaNet (linear attention) instead of
         # standard self-attention, so self_attn.o_proj doesn't exist on those layers.
         with suppress(Exception):
-            try_add("attn.o_proj", layer.linear_attn.out_proj)  # ty:ignore[possibly-missing-attribute]
+            try_add("attn.o_proj", layer.linear_attn.out_proj)  # ty: ignore[unresolved-attribute]
 
         # Most dense models.
         with suppress(Exception):
-            try_add("mlp.down_proj", layer.mlp.down_proj)  # ty:ignore[possibly-missing-attribute]
+            try_add("mlp.down_proj", layer.mlp.down_proj)  # ty: ignore[unresolved-attribute]
 
         # Some MoE models (e.g. Qwen3).
         with suppress(Exception):
-            for expert in layer.mlp.experts:  # ty:ignore[possibly-missing-attribute, not-iterable]
+            for expert in layer.mlp.experts:  # ty:ignore[not-iterable, unresolved-attribute]
-                try_add("mlp.down_proj", expert.down_proj)  # ty:ignore[possibly-missing-attribute]
+                try_add("mlp.down_proj", expert.down_proj)  # ty: ignore[unresolved-attribute]
 
         # Phi-3.5-MoE (and possibly others).
         with suppress(Exception):
-            for expert in layer.block_sparse_moe.experts:  # ty:ignore[possibly-missing-attribute, not-iterable]
+            for expert in layer.block_sparse_moe.experts:  # ty:ignore[not-iterable, unresolved-attribute]
-                try_add("mlp.down_proj", expert.w2)  # ty:ignore[possibly-missing-attribute]
+                try_add("mlp.down_proj", expert.w2)  # ty: ignore[unresolved-attribute]
 
         # LFM dense operator blocks.
         with suppress(Exception):
-            try_add("attn.o_proj", layer.conv.out_proj)  # ty:ignore[possibly-missing-attribute]
+            try_add("attn.o_proj", layer.conv.out_proj)  # ty: ignore[unresolved-attribute]
 
         with suppress(Exception):
-            try_add("mlp.down_proj", layer.feed_forward.w2)  # ty:ignore[possibly-missing-attribute]
+            try_add("mlp.down_proj", layer.feed_forward.w2)  # ty: ignore[unresolved-attribute]
 
         # LFM transformer blocks.
         with suppress(Exception):
-            try_add("attn.o_proj", layer.self_attn.out_proj)  # ty:ignore[possibly-missing-attribute]
+            try_add("attn.o_proj", layer.self_attn.out_proj)  # ty: ignore[unresolved-attribute]
 
         with suppress(Exception):
-            for expert in layer.feed_forward.experts:  # ty:ignore[possibly-missing-attribute, not-iterable]
+            for expert in layer.feed_forward.experts:  # ty:ignore[not-iterable, unresolved-attribute]
-                try_add("mlp.down_proj", expert.w2)  # ty:ignore[possibly-missing-attribute]
+                try_add("mlp.down_proj", expert.w2)  # ty: ignore[unresolved-attribute]
 
         # Granite MoE Hybrid - attention layers with shared_mlp.
         with suppress(Exception):
-            try_add("mlp.down_proj", layer.shared_mlp.output_linear)  # ty:ignore[possibly-missing-attribute]
+            try_add("mlp.down_proj", layer.shared_mlp.output_linear)  # ty: ignore[unresolved-attribute]
 
         # Granite MoE Hybrid - MoE layers with experts.
         with suppress(Exception):
-            for expert in layer.moe.experts:  # ty:ignore[possibly-missing-attribute, not-iterable]
+            for expert in layer.moe.experts:  # ty:ignore[not-iterable, unresolved-attribute]
-                try_add("mlp.down_proj", expert.output_linear)  # ty:ignore[possibly-missing-attribute]
+                try_add("mlp.down_proj", expert.output_linear)  # ty: ignore[unresolved-attribute]
 
         # We need at least one module across all components for abliteration to work.
         total_modules = sum(len(mods) for mods in modules.values())
@@ -458,166 +443,6 @@ class Model:
 
         return sorted(components)
 
-    def abliterate(
-        self,
-        residual_directions: Tensor,
-        direction_index: float | None,
-        parameters: dict[str, AbliterationParameters],
-    ):
-        if direction_index is None:
-            residual_direction = None
-        else:
-            # The index must be shifted by 1 because the first element
-            # of residual_directions is the direction for the embeddings.
-            weight, index = math.modf(direction_index + 1)
-            residual_direction = F.normalize(
-                residual_directions[int(index)].lerp(
-                    residual_directions[int(index) + 1],
-                    weight,
-                ),
-                p=2,
-                dim=0,
-            )
-
-        # Note that some implementations of abliteration also orthogonalize
-        # the embedding matrix, but it's unclear if that has any benefits.
-        for layer_index in range(len(self.get_layers())):
-            for component, modules in self.get_layer_modules(layer_index).items():
-                params = parameters[component]
-
-                # Type inference fails here for some reason.
-                distance = cast(float, abs(layer_index - params.max_weight_position))
-
-                # Don't orthogonalize layers that are more than
-                # min_weight_distance away from max_weight_position.
-                if distance > params.min_weight_distance:
-                    continue
-
-                # Interpolate linearly between max_weight and min_weight
-                # over min_weight_distance.
-                weight = params.max_weight + (distance / params.min_weight_distance) * (
-                    params.min_weight - params.max_weight
-                )
-
-                # A weight of 0 disables this component's ablation. reset_model() has
-                # already left the adapter at identity, so abort before the otherwise
-                # wasteful decomposition (which would also be operating on a zero matrix).
-                if weight == 0:
-                    continue
-
-                if residual_direction is None:
-                    # The index must be shifted by 1 because the first element
-                    # of residual_directions is the direction for the embeddings.
-                    layer_residual_direction = residual_directions[layer_index + 1]
-                else:
-                    layer_residual_direction = residual_direction
-
-                for module in modules:
-                    # FIXME: This cast is potentially invalid, because the program logic
-                    #        does not guarantee that the module is of type Linear, and in fact
-                    #        the retrieved modules might not conform to the interface assumed
-                    #        below (though they do in practice). However, this is difficult
-                    #        to fix cleanly, because get_layer_modules is called twice on
-                    #        different model configurations, and PEFT employs different
-                    #        module types depending on the chosen quantization.
-                    module = cast(Linear, module)
-
-                    # LoRA abliteration: delta W = -lambda * v * (v^T W)
-                    # lora_B = -lambda * v
-                    # lora_A = v^T W
-
-                    # Use the FP32 residual direction directly (no downcast/upcast)
-                    # and move to the correct device.
-                    v = layer_residual_direction.to(module.weight.device)
-
-                    # Get W (dequantize if necessary).
-                    #
-                    # FIXME: This cast is valid only under the assumption that the original
-                    #        module wrapped by the LoRA adapter has a weight attribute.
-                    #        See the comment above for why this is currently not guaranteed.
-                    base_weight = cast(Tensor, module.base_layer.weight)
-                    quant_state = getattr(base_weight, "quant_state", None)
-
-                    if quant_state is None:
-                        W = base_weight.to(torch.float32)
-                    else:
-                        # 4-bit quantization.
-                        # This cast is always valid. Type inference fails here because the
-                        # bnb.functional module is not found by ty for some reason.
-                        W = cast(
-                            Tensor,
-                            bnb.functional.dequantize_4bit(  # ty:ignore[possibly-missing-attribute]
-                                base_weight.data,
-                                quant_state,
-                            ).to(torch.float32),
-                        )
-
-                    # Flatten weight matrix to (out_features, in_features).
-                    W = W.view(W.shape[0], -1)
-
-                    if self.settings.row_normalization == RowNormalization.FULL:
-                        # Keep a reference to the original weight matrix so we can subtract it later.
-                        W_org = W
-
-                    if self.settings.row_normalization != RowNormalization.NONE:
-                        # Get the row norms.
-                        W_row_norms = LA.vector_norm(W, dim=1, keepdim=True)
-                        # Normalize the weight matrix along the rows.
-                        W = F.normalize(W, p=2, dim=1)
-
-                    # Calculate lora_A = v^T W
-                    # v is (d_out,), W is (d_out, d_in)
-                    # v @ W -> (d_in,)
-                    lora_A = (v @ W).view(1, -1)
-
-                    # Calculate lora_B = -weight * v
-                    # v is (d_out,)
-                    lora_B = (-weight * v).view(-1, 1)
-
-                    if self.settings.row_normalization == RowNormalization.PRE:
-                        # Make the LoRA adapter apply to the original weight matrix.
-                        lora_B = W_row_norms * lora_B
-                    elif self.settings.row_normalization == RowNormalization.FULL:
-                        # Approximates https://huggingface.co/blog/grimjim/norm-preserving-biprojected-abliteration
-                        W = W + lora_B @ lora_A
-                        # Normalize the adjusted weight matrix along the rows.
-                        W = F.normalize(W, p=2, dim=1)
-                        # Restore the original row norms of the weight matrix.
-                        W = W * W_row_norms
-                        # Subtract the original matrix to turn W into a delta.
-                        W = W - W_org
-                        # Use a low-rank SVD to get an approximation of the matrix.
-                        r = self.peft_config.r
-
-                        # svd_lowrank is randomized:
-                        # https://github.com/pytorch/pytorch/blob/20919052303c0b5ba87f8bf7e19237dc33ab09d3/torch/_lowrank.py#L108-L109
-                        # Reseed immediately before the call so restoring a trial is independent of RNG history.
-                        torch.manual_seed(self.settings.seed)
-                        # "It's safe to call this function if CUDA is not available;
-                        # in that case, it is silently ignored."
-                        torch.cuda.manual_seed_all(self.settings.seed)  # ty:ignore[invalid-argument-type]
-                        U, S, Vh = torch.svd_lowrank(W, q=2 * r + 4, niter=6)
-
-                        # Truncate it to the part we want to store in the LoRA adapter.
-                        # Note: svd_lowrank actually returns V, so transpose it to get Vh.
-                        U = U[:, :r]
-                        S = S[:r]
-                        Vh = Vh[:, :r].T
-                        # Transfer it into the LoRA adapter components. Split the singular values
-                        # evenly between the two components to keep their norms balanced and avoid
-                        # potential issues with numerical stability.
-                        sqrt_S = torch.sqrt(S)
-                        lora_B = U @ torch.diag(sqrt_S)
-                        lora_A = torch.diag(sqrt_S) @ Vh
-
-                    # Assign to adapters. The adapter name is "default", because that's
-                    # what PEFT uses when no name is explicitly specified, as above.
-                    # These casts are therefore valid.
-                    weight_A = cast(Tensor, module.lora_A["default"].weight)
-                    weight_B = cast(Tensor, module.lora_B["default"].weight)
-                    weight_A.data = lora_A.to(weight_A.dtype)
-                    weight_B.data = lora_B.to(weight_B.dtype)
-
     def generate(
         self,
         prompts: list[Prompt],
@@ -691,16 +516,22 @@ class Model:
         skip_special_tokens: bool = False,
     ) -> list[str]:
         responses = []
+
         for batch in batchify(prompts, self.settings.batch_size):
-            for response in self.get_responses(
+            responses.extend(
-                batch,
+                self.get_responses(
-                skip_special_tokens=skip_special_tokens,
+                    batch,
-            ):
+                    skip_special_tokens=skip_special_tokens,
-                responses.append(response)
+                )
+            )
 
         return responses
 
-    def get_residuals(self, prompts: list[Prompt]) -> Tensor:
+    def get_residuals(
+        self,
+        prompts: list[Prompt],
+        winsorization_quantile: float = 1.0,
+    ) -> Tensor:
         # We only generate one token, and we return the residual vectors
         # at that token position, for each prompt and layer.
         _, outputs = self.generate(
@@ -734,13 +565,13 @@ class Model:
         # problems during calculations involving residual vectors.
         residuals = residuals.to(torch.float32)
 
-        if 0 <= self.settings.winsorization_quantile < 1:
+        if 0 <= winsorization_quantile < 1:
             # Apply symmetric winsorization to each layer of the per-prompt residuals.
             abs_residuals = torch.abs(residuals)
             # Get the (prompt, layer, 1) quantiles of the (prompt, layer, component) residuals.
             thresholds = torch.quantile(
                 abs_residuals,
-                self.settings.winsorization_quantile,
+                winsorization_quantile,
                 dim=2,
                 keepdim=True,
             )
@@ -752,15 +583,28 @@ class Model:
 
         return residuals
 
-    def get_residuals_batched(self, prompts: list[Prompt]) -> Tensor:
+    def get_residuals_batched(
+        self,
+        prompts: list[Prompt],
+        winsorization_quantile: float = 1.0,
+    ) -> Tensor:
         residuals = []
 
         for batch in batchify(prompts, self.settings.batch_size):
-            residuals.append(self.get_residuals(batch))
+            residuals.append(
+                self.get_residuals(
+                    batch,
+                    winsorization_quantile=winsorization_quantile,
+                )
+            )
 
         return torch.cat(residuals, dim=0)
 
-    def get_residuals_mean(self, prompts: list[Prompt]) -> Tensor:
+    def get_residuals_mean(
+        self,
+        prompts: list[Prompt],
+        winsorization_quantile: float = 1.0,
+    ) -> Tensor:
         if not prompts:
             raise ValueError("prompts must not be empty")
 
@@ -768,7 +612,10 @@ class Model:
         total_count = 0
 
         for batch in batchify(prompts, self.settings.batch_size):
-            batch_residuals = self.get_residuals(batch)
+            batch_residuals = self.get_residuals(
+                batch,
+                winsorization_quantile=winsorization_quantile,
+            )
 
             # Accumulate in high precision on CPU to reduce peak VRAM usage.
             batch_sum = batch_residuals.sum(dim=0, dtype=torch.float64).cpu()
@@ -784,6 +631,132 @@ class Model:
 
         return (running_sum / total_count).to(torch.float32)
 
+    def get_module_io(
+        self,
+        prompts: list[Prompt],
+    ) -> ModuleIO:
+        # The list contains one element per layer.
+        # Each element maps from the component name to a (possibly sparse) mapping
+        # from the module index to an (input, output) tuple containing the I/O
+        # tensors of shape (prompt, component).
+        module_io: ModuleIO = []
+
+        def get_hook(
+            layer_index: int,
+            component: str,
+            module_index: int,
+        ) -> Callable[[Module, tuple[Tensor, ...], Tensor], None]:
+            def hook(
+                module: Module,
+                inputs: tuple[Tensor, ...],
+                outputs: Tensor,
+            ) -> None:
+                if len(module_io) == layer_index:
+                    # First invocation of the hook for this layer.
+                    module_io.append({})
+
+                # Layers are invoked in order during inference,
+                # so this should always hold.
+                assert len(module_io) == layer_index + 1
+
+                if component not in module_io[layer_index]:
+                    module_io[layer_index][component] = {}
+
+                # Each module should be invoked at most once per inference step.
+                assert module_index not in module_io[layer_index][component]
+
+                # inputs[0] and outputs have shape (prompt, position, component),
+                # so this extracts the input/output at the end of each prompt.
+                # Move to CPU to decouple from device assignments, which can
+                # change between model reloads in multi-GPU configurations.
+                input = inputs[0][:, -1, :].detach().clone().cpu()
+                output = outputs[:, -1, :].detach().clone().cpu()
+
+                # The modules associated with a component (e.g. expert MLPs)
+                # are not necessarily invoked in order, nor are all of them
+                # necessarily invoked in each inference step, so we cannot
+                # use a list here.
+                module_io[layer_index][component][module_index] = (input, output)
+
+            return hook
+
+        hook_handles: list[RemovableHandle] = []
+
+        for layer_index in range(len(self.get_layers())):
+            for component, modules in self.get_layer_modules(layer_index).items():
+                for module_index, module in enumerate(modules):
+                    hook_handles.append(
+                        module.register_forward_hook(
+                            get_hook(layer_index, component, module_index)
+                        )
+                    )
+
+        self.generate(prompts, max_new_tokens=1)
+
+        for hook_handle in hook_handles:
+            hook_handle.remove()
+
+        return module_io
+
+    def get_module_io_batched(
+        self,
+        prompts: list[Prompt],
+    ) -> ModuleIO:
+        # Aggregating batch results is more complicated for module I/O
+        # than for other get_*_batched methods, because the structure of the results
+        # might differ between batches, as whether individual modules activate
+        # can depend on the prompt (in particular for MoE models).
+        # In practice, inhomogeneous results should be very rare, but to be fully
+        # generic, this logic is required.
+        module_io_batches: list[ModuleIO] = [
+            self.get_module_io(batch)
+            for batch in batchify(prompts, self.settings.batch_size)
+        ]
+
+        module_io: ModuleIO = []
+
+        for layer_index in range(len(self.get_layers())):
+            module_io.append({})
+
+            for module_io_batch in module_io_batches:
+                for component, io_map in module_io_batch[layer_index].items():
+                    if component not in module_io[layer_index]:
+                        module_io[layer_index][component] = {}
+
+                    for module_index in io_map:
+                        if module_index not in module_io[layer_index][component]:
+                            # This is a placeholder; the actual aggregation happens below.
+                            # We need to iterate over the batches twice because we don't
+                            # know in advance which components and module indices are present.
+                            module_io[layer_index][component][module_index] = (
+                                torch.empty(0),
+                                torch.empty(0),
+                            )
+
+            for component, io_map in module_io[layer_index].items():
+                for module_index in io_map:
+                    inputs_outputs = [
+                        module_io_batch[layer_index][component][module_index]
+                        for module_io_batch in module_io_batches
+                        if component in module_io_batch[layer_index]
+                        and module_index in module_io_batch[layer_index][component]
+                    ]
+                    input = torch.cat(
+                        [input_output[0] for input_output in inputs_outputs],
+                        dim=0,
+                    )
+                    output = torch.cat(
+                        [input_output[1] for input_output in inputs_outputs],
+                        dim=0,
+                    )
+
+                    # The key already exists, and replacing existing values
+                    # in a dictionary while iterating over the same dictionary
+                    # is safe in Python.
+                    module_io[layer_index][component][module_index] = (input, output)
+
+        return module_io
+
     def get_logits(self, prompts: list[Prompt]) -> Tensor:
         # We only generate one token, and we return the raw logits over the vocabulary
         # at that token position, for each prompt.
@@ -843,7 +816,7 @@ class Model:
             # The TextStreamer constructor annotates this parameter with the AutoTokenizer
             # type, which makes no sense because AutoTokenizer is a factory class,
             # not a base class that tokenizers inherit from.
-            self.tokenizer,  # ty:ignore[invalid-argument-type]
+            self.tokenizer,
             skip_prompt=True,
             skip_special_tokens=True,
         )

src/heretic/modifier.py

@@ -0,0 +1,185 @@
+# 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 dataclasses import dataclass
+from typing import Any, Generic, Protocol, TypeVar, get_args
+
+from optuna import Trial
+from optuna.trial import FrozenTrial
+from pydantic import BaseModel
+
+from .config import (
+    ModifierConfig,
+)
+from .config import (
+    Settings as HereticSettings,
+)
+from .model import Model
+from .plugin import Context, Plugin, load_plugin
+from .utils import print
+
+
+class Serializable(Protocol):
+    def to_dict(self) -> dict[str, Any]: ...
+
+    def to_presentation_dict(self) -> dict[str, str]: ...
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "Serializable": ...
+
+
+Parameters = TypeVar("Parameters", bound=Serializable)
+
+
+class Modifier(Plugin, ABC, Generic[Parameters]):
+    """
+    Abstract base class for modifier plugins.
+
+    Modifiers modify models based on an implementation-dependent set of optimizable parameters.
+
+    Examples: Standard abliteration, ARA, SOMA, etc.
+    """
+
+    @property
+    def modifier_name(self) -> str:
+        """
+        The name of the modifier.
+        This is what shows up in the CLI and Markdown on HF.
+        """
+        return self.__class__.__name__
+
+    @property
+    def parameters_class(self) -> type[Parameters]:
+        """
+        The class of the modifier's parameters type.
+        """
+        base_class = self.__class__.__orig_bases__[0]  # ty:ignore[unresolved-attribute]
+        generic_type = get_args(base_class)[0]
+        return generic_type
+
+    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`.
+        """
+
+    def render_trial_parameters(self, trial: Trial | FrozenTrial) -> dict[str, str]:
+        """
+        Transform the names and values of the modifier's parameters
+        that are contained in the trial's user attributes into a form
+        suitable for presentation.
+        """
+        return self.parameters_class.from_dict(
+            trial.user_attrs["parameters"]
+        ).to_presentation_dict()
+
+
+@dataclass
+class ModifierEntry:
+    modifier: Modifier[Any]
+    name: str
+    config: ModifierConfig
+
+
+def load_and_init_modifiers(
+    settings: HereticSettings,
+    model: Model,
+) -> list[ModifierEntry]:
+    """
+    Load and instantiate all configured modifier plugins,
+    then runs their initialization hooks.
+    """
+    modifier_configs = settings.modifiers
+    if not modifier_configs:
+        raise ValueError("No modifiers configured. Set 'modifiers' in config.toml")
+    if len(modifier_configs) > 1:
+        raise ValueError("Using multiple modifiers is not yet supported")
+
+    modifier_keys: set[str] = set()
+
+    modifier_entries: list[ModifierEntry] = []
+
+    # Resolve plugin classes from names and validate.
+    for config in modifier_configs:
+        modifier_cls = load_plugin(name=config.plugin, base_class=Modifier)
+        modifier_cls.validate_contract()
+
+        print(
+            f"* Loaded: [bold]{modifier_cls.__name__}{' - ' + config.instance_name if config.instance_name else ''}[/bold]"
+        )
+
+        # Instantiate modifiers.
+        instance_name = config.instance_name or None
+
+        raw_settings = modifier_cls.get_settings_raw(
+            settings.model_extra,
+            "modifier",
+            instance_name,
+        )
+        modifier_settings: BaseModel | None = modifier_cls.validate_settings(
+            raw_settings
+        )
+
+        modifier = modifier_cls(
+            heretic_settings=settings,
+            settings=modifier_settings,
+        )
+
+        # External labeling key: ensures multiple instances can coexist.
+        # Uses underscore to match the TOML namespace format (`modifier.<Class>_<instance>`).
+        modifier_key = (
+            modifier_cls.__name__
+            if not instance_name
+            else f"{modifier_cls.__name__}_{instance_name}"
+        )
+        if modifier_key in modifier_keys:
+            raise ValueError(
+                f"Duplicate modifier instance name: {modifier_key}. "
+                "Give each instance a unique `instance_name`."
+            )
+        modifier_keys.add(modifier_key)
+
+        modifier_instance_name = (
+            f"{modifier.modifier_name} - {instance_name}"
+            if instance_name
+            else modifier.modifier_name
+        )
+        modifier_entries.append(
+            ModifierEntry(
+                modifier=modifier,
+                config=config,
+                name=modifier_instance_name,
+            )
+        )
+
+    # Run modifier init hooks.
+    ctx = Context(settings=settings, model=model)
+
+    for entry in modifier_entries:
+        entry.modifier.init(ctx)
+
+    return modifier_entries

src/heretic/modifiers/abliteration.py

@@ -0,0 +1,466 @@
+# SPDX-License-Identifier: AGPL-3.0-or-later
+# Copyright (C) 2025-2026  Philipp Emanuel Weidmann <pew@worldwidemann.com> + contributors
+
+import math
+from dataclasses import asdict, dataclass
+from enum import Enum
+from typing import Any, cast
+
+import bitsandbytes.functional as BNB_F
+import torch
+import torch.linalg as LA
+import torch.nn.functional as F
+from optuna import Trial
+from peft.tuners.lora.layer import Linear
+from pydantic import (
+    BaseModel,
+    Field,
+    PositiveInt,
+)
+from torch import Tensor
+
+from heretic.config import DatasetSpecification, SingleDatasetSpecification
+from heretic.modifier import Context, Modifier, Serializable
+from heretic.utils import format_dataset_specification, print
+
+
+@dataclass
+class WeightDistribution:
+    max_weight: float
+    max_weight_position: float
+    min_weight: float
+    min_weight_distance: float
+
+
+@dataclass
+class Parameters(Serializable):
+    direction_index: float | None
+    weight_distributions: dict[str, WeightDistribution]
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+    def to_presentation_dict(self) -> dict[str, str]:
+        parameters = {}
+
+        parameters["direction_index"] = (
+            "per layer"
+            if (self.direction_index is None)
+            else f"{self.direction_index:.2f}"
+        )
+
+        for component, weight_distribution in self.weight_distributions.items():
+            for name, value in asdict(weight_distribution).items():
+                parameters[f"{component}.{name}"] = f"{value:.2f}"
+
+        return parameters
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "Serializable":
+        return Parameters(
+            direction_index=data["direction_index"],
+            weight_distributions={
+                component: WeightDistribution(**weight_distribution)
+                for component, weight_distribution in data[
+                    "weight_distributions"
+                ].items()
+            },
+        )
+
+
+class RowNormalization(str, Enum):
+    NONE = "none"
+    PRE = "pre"
+    # POST = "post"  # Theoretically possible, but provides no advantage.
+    FULL = "full"
+
+
+class Settings(BaseModel):
+    good_prompts: DatasetSpecification = Field(
+        default=SingleDatasetSpecification(
+            dataset="mlabonne/harmless_alpaca",
+            split="train[:400]",
+            column="text",
+        ),
+        description="Dataset of prompts that tend to produce desirable responses.",
+    )
+
+    bad_prompts: DatasetSpecification = Field(
+        default=SingleDatasetSpecification(
+            dataset="mlabonne/harmful_behaviors",
+            split="train[:400]",
+            column="text",
+        ),
+        description="Dataset of prompts that tend to produce undesirable responses.",
+    )
+
+    orthogonalize_direction: bool = Field(
+        default=True,
+        description=(
+            "Whether to adjust the residual directions so that only the component that is "
+            "orthogonal to the good direction is subtracted during abliteration."
+        ),
+    )
+
+    row_normalization: RowNormalization = Field(
+        default=RowNormalization.FULL,
+        description=(
+            "How to apply row normalization of the weights. Options: "
+            '"none" (no normalization), '
+            '"pre" (compute LoRA adapter relative to row-normalized weights), '
+            '"full" (like "pre", but renormalizes to preserve original row magnitudes).'
+        ),
+    )
+
+    full_normalization_lora_rank: PositiveInt = Field(
+        default=3,
+        description=(
+            'The rank of the LoRA adapter to use when "full" row normalization is used. '
+            "Row magnitude preservation is approximate due to non-linear effects, "
+            "and this determines the rank of that approximation. Higher ranks produce "
+            "larger output files and may slow down evaluation."
+        ),
+    )
+
+    winsorization_quantile: float = Field(
+        default=1.0,
+        description=(
+            "The symmetric winsorization to apply to the per-prompt, per-layer residual vectors, "
+            "expressed as the quantile to clamp to (between 0 and 1). Disabled by default. "
+            'This can tame so-called "massive activations" that occur in some models. '
+            "Example: winsorization_quantile = 0.95 computes the 0.95-quantile of the absolute values "
+            "of the components, then clamps the magnitudes of all components to that quantile."
+        ),
+    )
+
+
+class Abliteration(Modifier[Parameters]):
+    settings: Settings
+
+    @property
+    def reproducible(self) -> bool:
+        return True
+
+    @property
+    def modifier_name(self) -> str:
+        if (
+            self.settings.orthogonalize_direction
+            and self.settings.row_normalization == RowNormalization.FULL
+        ):
+            return "Magnitude-Preserving Orthogonal Ablation (MPOA)"
+        elif self.settings.orthogonalize_direction:
+            return "Projected Abliteration"
+        else:
+            return "Abliteration"
+
+    def init(self, ctx: Context) -> None:
+        model = ctx.get_model()
+
+        print()
+        print(
+            f"Loading good prompts from [bold]{format_dataset_specification(self.settings.good_prompts)}[/]..."
+        )
+        good_prompts = ctx.load_prompts(self.settings.good_prompts)
+        print(f"* [bold]{len(good_prompts)}[/] prompts loaded")
+
+        print()
+        print(
+            f"Loading bad prompts from [bold]{format_dataset_specification(self.settings.bad_prompts)}[/]..."
+        )
+        bad_prompts = ctx.load_prompts(self.settings.bad_prompts)
+        print(f"* [bold]{len(bad_prompts)}[/] prompts loaded")
+
+        print()
+        print("Calculating per-layer residual directions...")
+
+        print("* Obtaining residual mean for good prompts...")
+        good_means = model.get_residuals_mean(
+            good_prompts,
+            winsorization_quantile=self.settings.winsorization_quantile,
+        )
+        print("* Obtaining residual mean for bad prompts...")
+        bad_means = model.get_residuals_mean(
+            bad_prompts,
+            winsorization_quantile=self.settings.winsorization_quantile,
+        )
+
+        self.residual_directions = F.normalize(
+            bad_means - good_means,
+            p=2,
+            dim=1,
+        )
+
+        if self.settings.orthogonalize_direction:
+            # Implements https://huggingface.co/blog/grimjim/projected-abliteration
+            # Adjust the residual directions so that only the component that is
+            # orthogonal to the good direction is subtracted during abliteration.
+            good_directions = F.normalize(
+                good_means,
+                p=2,
+                dim=1,
+            )
+            projection_vector = torch.sum(
+                self.residual_directions * good_directions,
+                dim=1,
+            )
+            self.residual_directions = (
+                self.residual_directions
+                - projection_vector.unsqueeze(1) * good_directions
+            )
+            self.residual_directions = F.normalize(
+                self.residual_directions,
+                p=2,
+                dim=1,
+            )
+
+        if self.settings.row_normalization != RowNormalization.FULL:
+            # Rank 1 is sufficient for directional ablation without renormalization.
+            self.lora_rank = 1
+        else:
+            # Row magnitude preservation introduces nonlinear effects.
+            self.lora_rank = self.settings.full_normalization_lora_rank
+
+        # LoRA B matrices are initialized to zero by default in PEFT,
+        # so we don't need to do anything manually.
+        model.apply_lora(self.lora_rank)
+
+    def suggest_parameters(self, ctx: Context, trial: Trial) -> Parameters:
+        model = ctx.get_model()
+
+        direction_scope = trial.suggest_categorical(
+            "direction_scope",
+            [
+                "global",
+                "per layer",
+            ],
+        )
+
+        last_layer_index = len(model.get_layers()) - 1
+
+        # Discrimination between "harmful" and "harmless" inputs is usually strongest
+        # in layers slightly past the midpoint of the layer stack. See the original
+        # abliteration paper (https://arxiv.org/abs/2406.11717) for a deeper analysis.
+        #
+        # Note that we always sample this parameter even though we only need it for
+        # the "global" direction scope. The reason is that multivariate TPE doesn't
+        # work with conditional or variable-range parameters.
+        direction_index = trial.suggest_float(
+            "direction_index",
+            0.4 * last_layer_index,
+            0.9 * last_layer_index,
+        )
+
+        if direction_scope == "per layer":
+            direction_index = None
+
+        weight_distributions = {}
+
+        for component in model.get_abliterable_components():
+            # The parameter ranges are based on experiments with various models
+            # and much wider ranges. They are not set in stone and might have to be
+            # adjusted for future models.
+            #
+            # The MLP gets a negative lower bound that is then clamped to 0, so the
+            # optimizer can fully disable its ablation. The clamp puts a positive
+            # probability mass on exactly 0 (the continuous sampler would otherwise
+            # reach 0 with probability zero). Ablating the MLP is often unnecessary for
+            # removing refusals and tends to damage model intelligence more than
+            # ablating the attention output, so on many models the optimum is to leave
+            # it (mostly) untouched. See issue #202.
+            max_weight_lower_bound = -0.25 if component == "mlp.down_proj" else 0.8
+            max_weight = max(
+                0.0,
+                trial.suggest_float(
+                    f"{component}.max_weight",
+                    max_weight_lower_bound,
+                    1.5,
+                ),
+            )
+            max_weight_position = trial.suggest_float(
+                f"{component}.max_weight_position",
+                0.6 * last_layer_index,
+                1.0 * last_layer_index,
+            )
+            # For sampling purposes, min_weight is expressed as a fraction of max_weight,
+            # again because multivariate TPE doesn't support variable-range parameters.
+            # The value is transformed into the actual min_weight value below.
+            min_weight = trial.suggest_float(
+                f"{component}.min_weight",
+                0.0,
+                1.0,
+            )
+            min_weight_distance = trial.suggest_float(
+                f"{component}.min_weight_distance",
+                1.0,
+                max(0.6 * last_layer_index, 1.0),
+            )
+
+            weight_distributions[component] = WeightDistribution(
+                max_weight=max_weight,
+                max_weight_position=max_weight_position,
+                min_weight=(min_weight * max_weight),
+                min_weight_distance=min_weight_distance,
+            )
+
+        return Parameters(
+            direction_index=direction_index,
+            weight_distributions=weight_distributions,
+        )
+
+    def modify_model(self, ctx: Context, parameters: Parameters) -> None:
+        model = ctx.get_model()
+
+        if parameters.direction_index is None:
+            residual_direction = None
+        else:
+            # The index must be shifted by 1 because the first element
+            # of residual_directions is the direction for the embeddings.
+            weight, index = math.modf(parameters.direction_index + 1)
+            residual_direction = F.normalize(
+                self.residual_directions[int(index)].lerp(
+                    self.residual_directions[int(index) + 1],
+                    weight,
+                ),
+                p=2,
+                dim=0,
+            )
+
+        # Note that some implementations of abliteration also orthogonalize
+        # the embedding matrix, but it's unclear if that has any benefits.
+        for layer_index in range(len(model.get_layers())):
+            for component, modules in model.get_layer_modules(layer_index).items():
+                weight_distribution = parameters.weight_distributions[component]
+
+                # Type inference fails here for some reason.
+                distance = abs(layer_index - weight_distribution.max_weight_position)
+
+                # Don't orthogonalize layers that are more than
+                # min_weight_distance away from max_weight_position.
+                if distance > weight_distribution.min_weight_distance:
+                    continue
+
+                # Interpolate linearly between max_weight and min_weight
+                # over min_weight_distance.
+                weight = weight_distribution.max_weight + (
+                    distance / weight_distribution.min_weight_distance
+                ) * (weight_distribution.min_weight - weight_distribution.max_weight)
+
+                # A weight of 0 disables this component's ablation. reset_model() has
+                # already left the adapter at identity, so abort before the otherwise
+                # wasteful decomposition (which would also be operating on a zero matrix).
+                if weight == 0:
+                    continue
+
+                if residual_direction is None:
+                    # The index must be shifted by 1 because the first element
+                    # of residual_directions is the direction for the embeddings.
+                    layer_residual_direction = self.residual_directions[layer_index + 1]
+                else:
+                    layer_residual_direction = residual_direction
+
+                for module in modules:
+                    # FIXME: This cast is potentially invalid, because the program logic
+                    #        does not guarantee that the module is of type Linear, and in fact
+                    #        the retrieved modules might not conform to the interface assumed
+                    #        below (though they do in practice). However, this is difficult
+                    #        to fix cleanly, because get_layer_modules is called twice on
+                    #        different model configurations, and PEFT employs different
+                    #        module types depending on the chosen quantization.
+                    module = cast(Linear, module)
+
+                    # LoRA abliteration: delta W = -lambda * v * (v^T W)
+                    # lora_B = -lambda * v
+                    # lora_A = v^T W
+
+                    # Use the FP32 residual direction directly (no downcast/upcast)
+                    # and move to the correct device.
+                    v = layer_residual_direction.to(module.weight.device)
+
+                    # Get W (dequantize if necessary).
+                    #
+                    # FIXME: This cast is valid only under the assumption that the original
+                    #        module wrapped by the LoRA adapter has a weight attribute.
+                    #        See the comment above for why this is currently not guaranteed.
+                    base_weight = cast(Tensor, module.base_layer.weight)
+                    quant_state = getattr(base_weight, "quant_state", None)
+
+                    if quant_state is None:
+                        W = base_weight.to(torch.float32)
+                    else:
+                        # 4-bit quantization.
+                        W = BNB_F.dequantize_4bit(
+                            base_weight.data,
+                            quant_state,
+                        ).to(torch.float32)
+
+                    # Flatten weight matrix to (out_features, in_features).
+                    W = W.view(W.shape[0], -1)
+
+                    if self.settings.row_normalization == RowNormalization.FULL:
+                        # Keep a reference to the original weight matrix so we can subtract it later.
+                        W_org = W
+
+                    if self.settings.row_normalization != RowNormalization.NONE:
+                        # Get the row norms.
+                        W_row_norms = LA.vector_norm(W, dim=1, keepdim=True)
+                        # Normalize the weight matrix along the rows.
+                        W = F.normalize(W, p=2, dim=1)
+
+                    # Calculate lora_A = v^T W
+                    # v is (d_out,), W is (d_out, d_in)
+                    # v @ W -> (d_in,)
+                    lora_A = (v @ W).view(1, -1)
+
+                    # Calculate lora_B = -weight * v
+                    # v is (d_out,)
+                    lora_B = (-weight * v).view(-1, 1)
+
+                    if self.settings.row_normalization == RowNormalization.PRE:
+                        # Make the LoRA adapter apply to the original weight matrix.
+                        lora_B = W_row_norms * lora_B
+                    elif self.settings.row_normalization == RowNormalization.FULL:
+                        # Approximates https://huggingface.co/blog/grimjim/norm-preserving-biprojected-abliteration
+                        W = W + lora_B @ lora_A
+                        # Normalize the adjusted weight matrix along the rows.
+                        W = F.normalize(W, p=2, dim=1)
+                        # Restore the original row norms of the weight matrix.
+                        W = W * W_row_norms
+                        # Subtract the original matrix to turn W into a delta.
+                        W = W - W_org
+                        # Use a low-rank SVD to get an approximation of the matrix.
+                        r = model.peft_config.r
+
+                        # svd_lowrank is randomized:
+                        # https://github.com/pytorch/pytorch/blob/20919052303c0b5ba87f8bf7e19237dc33ab09d3/torch/_lowrank.py#L108-L109
+                        # Reseed immediately before the call so restoring a trial is independent of RNG history.
+                        torch.manual_seed(self.heretic_settings.seed)
+                        # "It's safe to call this function if CUDA is not available;
+                        # in that case, it is silently ignored."
+                        torch.cuda.manual_seed_all(self.heretic_settings.seed)  # ty:ignore[invalid-argument-type]
+                        U, S, Vh = torch.svd_lowrank(W, q=2 * r + 4, niter=6)
+
+                        # Truncate it to the part we want to store in the LoRA adapter.
+                        # Note: svd_lowrank actually returns V, so transpose it to get Vh.
+                        U = U[:, :r]
+                        S = S[:r]
+                        Vh = Vh[:, :r].T
+                        # Transfer it into the LoRA adapter components. Split the singular values
+                        # evenly between the two components to keep their norms balanced and avoid
+                        # potential issues with numerical stability.
+                        sqrt_S = torch.sqrt(S)
+                        lora_B = U @ torch.diag(sqrt_S)
+                        lora_A = torch.diag(sqrt_S) @ Vh
+
+                    # Assign to adapters. The adapter name is "default", because that's
+                    # what PEFT uses when no name is explicitly specified, as above.
+                    # These casts are therefore valid.
+                    weight_A = cast(Tensor, module.lora_A["default"].weight)
+                    weight_B = cast(Tensor, module.lora_B["default"].weight)
+                    weight_A.data = lora_A.to(weight_A.dtype)
+                    weight_B.data = lora_B.to(weight_B.dtype)
+
+    def reset_model(self, ctx: Context) -> None:
+        model = ctx.get_model()
+        fast_path = model.reset_model()
+        if not fast_path:
+            model.apply_lora(self.lora_rank)

src/heretic/modifiers/ara.py

@@ -0,0 +1,348 @@
+# SPDX-License-Identifier: AGPL-3.0-or-later
+# Copyright (C) 2025-2026  Philipp Emanuel Weidmann <pew@worldwidemann.com> + contributors
+
+# Arbitrary-Rank Ablation (ARA) (Weidmann 2026)
+# See https://github.com/p-e-w/heretic/pull/211 for more information.
+
+from dataclasses import asdict, dataclass
+from typing import Any, cast
+
+import bitsandbytes.functional as BNB_F
+import torch
+import torch.linalg as LA
+import torch.nn.functional as F
+from optuna import Trial
+from peft.tuners.lora.layer import Linear
+from pydantic import (
+    BaseModel,
+    Field,
+    PositiveInt,
+)
+from torch import Tensor
+from torch.optim import LBFGS
+
+from heretic.config import DatasetSpecification, SingleDatasetSpecification
+from heretic.modifier import Context, Modifier, Serializable
+from heretic.utils import format_dataset_specification, print
+
+
+@dataclass
+class Parameters(Serializable):
+    start_layer_index: int
+    end_layer_index: int
+    preserve_good_behavior_weight: float
+    steer_bad_behavior_weight: float
+    overcorrect_relative_weight: float
+    neighbor_count: int
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+    def to_presentation_dict(self) -> dict[str, str]:
+        return {
+            name: (f"{value:.4f}" if isinstance(value, float) else f"{value}")
+            for name, value in asdict(self).items()
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "Serializable":
+        return Parameters(**data)
+
+
+class Settings(BaseModel):
+    good_prompts: DatasetSpecification = Field(
+        default=SingleDatasetSpecification(
+            dataset="mlabonne/harmless_alpaca",
+            split="train[:400]",
+            column="text",
+        ),
+        description="Dataset of prompts that tend to produce desirable responses.",
+    )
+
+    bad_prompts: DatasetSpecification = Field(
+        default=SingleDatasetSpecification(
+            dataset="mlabonne/harmful_behaviors",
+            split="train[:400]",
+            column="text",
+        ),
+        description="Dataset of prompts that tend to produce undesirable responses.",
+    )
+
+    preserve_row_magnitudes: bool = Field(
+        default=True,
+        description=(
+            "Whether to renormalize the rows of the modified matrices to preserve "
+            "the original matrices' row magnitudes. This is believed to improve "
+            'intelligence retention (see Lai 2025, "Magnitude-Preserving Orthogonal Ablation").'
+        ),
+    )
+
+    lora_rank: PositiveInt = Field(
+        default=50,
+        description=(
+            "The rank of the LoRA adapter to use. "
+            'While mathematically, ARA is of "arbitrary" rank, experiments have shown that '
+            "singular values tend to drop rapidly after a few dozen dimensions, and approximating "
+            "the full transformation with a LoRA has many practical advantages."
+        ),
+    )
+
+    n_optimization_steps: PositiveInt = Field(
+        default=5,
+        description="Number of (outer) L-BFGS optimization steps to perform.",
+    )
+
+    learning_rate: float = Field(
+        default=1.0,
+        description="Learning rate to use in the L-BFGS optimizer.",
+    )
+
+    max_iter: PositiveInt = Field(
+        default=20,
+        description="Maximum number of (inner) iterations to perform per (outer) L-BFGS optimization step.",
+    )
+
+    history_size: PositiveInt = Field(
+        default=10,
+        description="Number of past updates to store for approximating the Hessian matrix in the L-BFGS optimizer.",
+    )
+
+    print_loss: bool = Field(
+        default=False,
+        description="Whether to print the loss value for each L-BFGS optimization step.",
+    )
+
+
+# For each vector in the 2D-tensor `a`, computes the mean Euclidean distance
+# to the `k` nearest neighbors of the vector among the vectors in the 2D-tensor `b`.
+def mean_distances_to_knn(a: Tensor, b: Tensor, k: int) -> Tensor:
+    distances = torch.cdist(a, b)
+    nearest_distances, _ = distances.topk(k, dim=1, largest=False)
+    return nearest_distances.mean(1)
+
+
+# The objective function at the heart of ARA.
+def ara_loss(
+    good_output: Tensor,
+    bad_output: Tensor,
+    new_good_output: Tensor,
+    new_bad_output: Tensor,
+    parameters: Parameters,
+) -> Tensor:
+    # The outputs for "good" prompts should change as little as possible.
+    preserve_good_behavior = ((new_good_output - good_output) ** 2).mean()
+
+    steer_bad_behavior = (
+        # Pull the outputs for "bad" prompts towards
+        # the original outputs for "good" prompts.
+        mean_distances_to_knn(
+            new_bad_output,
+            good_output,
+            parameters.neighbor_count,
+        ).mean()
+        # Push the outputs for "bad" prompts away from
+        # the original outputs for "bad" prompts.
+        # In combination with the above, this overcorrects
+        # away from the original residuals, which results
+        # in stronger steering that can overcome more complex
+        # refusal mechanisms.
+        + parameters.overcorrect_relative_weight
+        * -mean_distances_to_knn(
+            new_bad_output,
+            bad_output,
+            parameters.neighbor_count,
+        ).mean()
+    )
+
+    return (
+        parameters.preserve_good_behavior_weight * preserve_good_behavior
+        + parameters.steer_bad_behavior_weight * steer_bad_behavior
+    )
+
+
+class ARA(Modifier[Parameters]):
+    settings: Settings
+
+    @property
+    def reproducible(self) -> bool:
+        return True
+
+    @property
+    def modifier_name(self) -> str:
+        return "Arbitrary-Rank Ablation (ARA)"
+
+    def init(self, ctx: Context) -> None:
+        model = ctx.get_model()
+
+        print()
+        print(
+            f"Loading good prompts from [bold]{format_dataset_specification(self.settings.good_prompts)}[/]..."
+        )
+        good_prompts = ctx.load_prompts(self.settings.good_prompts)
+        print(f"* [bold]{len(good_prompts)}[/] prompts loaded")
+
+        print()
+        print(
+            f"Loading bad prompts from [bold]{format_dataset_specification(self.settings.bad_prompts)}[/]..."
+        )
+        bad_prompts = ctx.load_prompts(self.settings.bad_prompts)
+        print(f"* [bold]{len(bad_prompts)}[/] prompts loaded")
+
+        print()
+        print("Obtaining module I/O for good prompts...")
+        self.good_module_io = model.get_module_io_batched(good_prompts)
+
+        print()
+        print("Obtaining module I/O for bad prompts...")
+        self.bad_module_io = model.get_module_io_batched(bad_prompts)
+
+        # LoRA B matrices are initialized to zero by default in PEFT,
+        # so we don't need to do anything manually.
+        model.apply_lora(self.settings.lora_rank)
+
+    def suggest_parameters(self, ctx: Context, trial: Trial) -> Parameters:
+        layer_count = len(ctx.get_model().get_layers())
+
+        start_layer_index = trial.suggest_int(
+            "start_layer_index",
+            0,
+            layer_count // 2,
+        )
+        end_layer_index = trial.suggest_int(
+            "end_layer_index",
+            layer_count // 2,
+            layer_count,
+        )
+        preserve_good_behavior_weight = trial.suggest_float(
+            "preserve_good_behavior_weight",
+            0.0,
+            1.0,
+        )
+        steer_bad_behavior_weight = trial.suggest_float(
+            "steer_bad_behavior_weight",
+            0.0001,
+            1.0,
+            log=True,
+        )
+        overcorrect_relative_weight = trial.suggest_float(
+            "overcorrect_relative_weight",
+            0.0,
+            1.3,
+        )
+        neighbor_count = trial.suggest_int(
+            "neighbor_count",
+            1,
+            15,
+        )
+
+        return Parameters(
+            start_layer_index=start_layer_index,
+            end_layer_index=end_layer_index,
+            preserve_good_behavior_weight=preserve_good_behavior_weight,
+            steer_bad_behavior_weight=steer_bad_behavior_weight,
+            overcorrect_relative_weight=overcorrect_relative_weight,
+            neighbor_count=neighbor_count,
+        )
+
+    def modify_model(self, ctx: Context, parameters: Parameters) -> None:
+        model = ctx.get_model()
+
+        for layer_index in range(
+            parameters.start_layer_index,
+            parameters.end_layer_index,
+        ):
+            for component, modules in model.get_layer_modules(layer_index).items():
+                for module_index, module in enumerate(modules):
+                    # Cast to Linear to access weights and LoRA adapters.
+                    module = cast(Linear, module)
+
+                    # We need the base weight in float32 to compute the effective weight.
+                    base_weight = cast(Tensor, module.base_layer.weight)
+                    quant_state = getattr(base_weight, "quant_state", None)
+
+                    if quant_state is None:
+                        W_base = base_weight.to(torch.float32)
+                    else:
+                        # Use the original dequantization logic from bitsandbytes.
+                        W_base = BNB_F.dequantize_4bit(
+                            base_weight.data,
+                            quant_state,
+                        ).to(torch.float32)
+
+                    # Pre-calculate the original row norms to preserve them.
+                    # See https://huggingface.co/blog/grimjim/norm-preserving-biprojected-abliteration
+                    W_row_norms = cast(
+                        Tensor,
+                        LA.vector_norm(W_base, dim=1, keepdim=True).detach(),
+                    )
+
+                    # We optimize the LoRA weights A and B.
+                    lora_A = cast(Tensor, module.lora_A["default"].weight)
+                    lora_B = cast(Tensor, module.lora_B["default"].weight)
+
+                    # Move I/O tensors to the device of the adapter weights.
+                    good_input, good_output = self.good_module_io[layer_index][
+                        component
+                    ][module_index]
+                    bad_input, bad_output = self.bad_module_io[layer_index][component][
+                        module_index
+                    ]
+
+                    good_input = good_input.float().to(lora_A.device)
+                    good_output = good_output.float().to(lora_A.device)
+                    bad_input = bad_input.float().to(lora_A.device)
+                    bad_output = bad_output.float().to(lora_A.device)
+
+                    def objective(A: Tensor, B: Tensor) -> Tensor:
+                        # Calculate effective weight after applying adapter.
+                        W_eff = W_base + (B @ A)
+
+                        if self.settings.preserve_row_magnitudes:
+                            # Normalize to unit length, then scale by original norms,
+                            # preserving the original row norms.
+                            W_eff = F.normalize(W_eff, p=2, dim=1) * W_row_norms
+
+                        # Compute outputs using the effective weight.
+                        new_good_output = good_input @ W_eff.T
+                        new_bad_output = bad_input @ W_eff.T
+
+                        return ara_loss(
+                            good_output,
+                            bad_output,
+                            new_good_output,
+                            new_bad_output,
+                            parameters,
+                        )
+
+                    optimizer = LBFGS(
+                        [lora_A, lora_B],
+                        lr=self.settings.learning_rate,
+                        max_iter=self.settings.max_iter,
+                        history_size=self.settings.history_size,
+                        line_search_fn="strong_wolfe",
+                    )
+
+                    def closure() -> Tensor:
+                        optimizer.zero_grad()
+                        loss = objective(lora_A, lora_B)
+                        loss.backward()
+                        return loss
+
+                    for step in range(self.settings.n_optimization_steps):
+                        loss = optimizer.step(closure)
+                        if self.settings.print_loss:
+                            print(
+                                f"\\[{layer_index}/{component}/{module_index}] Step: {step + 1}, Loss: {loss.item():.6f}"
+                            )
+
+                    # Free the gradient buffers accumulated during optimization.
+                    # Without this, they persist on the model (one full-size gradient
+                    # per processed weight) and can easily consume tens of GB of VRAM,
+                    # causing out-of-memory errors during the subsequent evaluation.
+                    optimizer.zero_grad(set_to_none=True)
+
+    def reset_model(self, ctx: Context) -> None:
+        model = ctx.get_model()
+        fast_path = model.reset_model()
+        if not fast_path:
+            model.apply_lora(self.settings.lora_rank)

src/heretic/plugin.py

@@ -13,17 +13,17 @@ from typing import Annotated, Any, TypeVar, Union, get_args, get_origin, get_typ
 from pydantic import BaseModel
 from torch import Tensor
 
-from heretic.utils import Prompt, load_prompts
+from .config import DatasetSpecification, SingleDatasetSpecification
-
-from .config import DatasetSpecification
 from .config import Settings as HereticSettings
 from .model import Model
+from .utils import Prompt, deep_merge_dicts, load_prompts
 
 T = TypeVar("T")
 
 
 def get_plugin_namespace(
-    model_extra: dict[str, Any] | None, namespace: str
+    model_extra: dict[str, Any] | None,
+    namespace: str,
 ) -> dict[str, Any]:
     """
     Returns the config dict from the `[<namespace>]` TOML table.
@@ -51,7 +51,7 @@ def is_builtin_plugin(name: str) -> bool:
     plugins (file paths or third-party import paths) disable the reproducibility
     offer during upload.
     """
-    return name.startswith("heretic.scorers.")
+    return name.startswith("heretic.")
 
 
 def load_plugin(
@@ -149,12 +149,8 @@ def load_plugin(
 
 class Context:
     """
-    Runtime context passed to plugins
+    Runtime context passed to plugins.
-
+    Acts as a quasi-API for plugins to access Heretic functionality.
-    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.
     """
 
     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()
@@ -280,9 +286,46 @@ class Plugin:
             )
         return model
 
+    @classmethod
+    def get_settings_raw(
+        cls,
+        model_extra: dict[str, Any] | None,
+        top_namespace: str,
+        instance_name: str | None,
+    ) -> dict[str, Any]:
+        """
+        Build the raw settings dict for a plugin class and optional instance.
+
+        Config rules:
+        - Base settings live in `[<top_namespace>.ClassName]` (applies to all instances).
+        - Instance overrides live in `[<top_namespace>.ClassName_<instance_name>]` (preferred).
+        - Only merge/validate keys that exist in the plugin Settings schema.
+        """
+        settings_model = cls.get_settings_model()
+        if settings_model is None:
+            # No settings schema: nothing to merge/validate.
+            return {}
+
+        class_name = cls.__name__
+
+        namespaces = [f"{top_namespace}.{class_name}"]
+        if instance_name:
+            namespaces.append(f"{top_namespace}.{class_name}_{instance_name}")
+
+        merged_settings: dict[str, Any] = {}
+        allowed_keys = set(settings_model.model_fields.keys())
+
+        for namespace in namespaces:
+            raw_table = get_plugin_namespace(model_extra, namespace)
+            filtered = {k: v for k, v in raw_table.items() if k in allowed_keys}
+            merged_settings = deep_merge_dicts(merged_settings, filtered)
+
+        return merged_settings
+
     @classmethod
     def validate_settings(
-        cls, raw_namespace: dict[str, Any] | None
+        cls,
+        raw_namespace: dict[str, Any] | None,
     ) -> BaseModel | None:
         """
         Validates plugin settings for this plugin class.
@@ -295,6 +338,23 @@ class Plugin:
             return None
         return settings_model.model_validate(raw_namespace or {})
 
+    def get_dataset_specifications(self) -> list[DatasetSpecification]:
+        """
+        Collect the dataset specifications declared in the settings
+        of the plugin.
+        """
+        if self.settings is None:
+            return []
+        specifications = []
+        for value in dict(self.settings).values():
+            if isinstance(value, SingleDatasetSpecification) or (
+                isinstance(value, list)
+                and len(value) > 0
+                and isinstance(value[0], SingleDatasetSpecification)
+            ):
+                specifications.append(value)
+        return specifications
+
     def init(self, ctx: Context) -> None:
         """
         Runs before the plugin's main functionality.
@@ -302,4 +362,4 @@ class Plugin:
         Override this in subclasses to do one-time setup (e.g. load prompts, compute
         baselines).
         """
-        return None
+        return

src/heretic/reproduce.py

@@ -81,7 +81,7 @@ def collect_reproducibles(path: str):
 
             found += 1
 
-            commit_hash = paths_info[0].last_commit.oid
+            commit_hash = paths_info[0].last_commit.oid  # ty: ignore[unresolved-attribute]
 
             file_path = (
                 Path(path)
@@ -285,12 +285,10 @@ def check_environment(
 
     else:
         print(
-            (
+            "[yellow]The provided JSON file does not contain system information. "
-                "[yellow]The provided JSON file does not contain system information. "
+            "Some system parameters can affect reproducibility, but due to the lack of system information, "
-                "Some system parameters can affect reproducibility, but due to the lack of system information, "
+            "Heretic is unable to verify that those parameters match the original environment. "
-                "Heretic is unable to verify that those parameters match the original environment. "
+            "Reproduction may or may not produce a byte-for-byte identical model.[/]"
-                "Reproduction may or may not produce a byte-for-byte identical model.[/]"
-            )
         )
 
     requirements = get_requirements_dict()
@@ -321,10 +319,8 @@ def check_environment(
     if system_mismatches or package_mismatches:
         print()
         print(
-            (
+            "[yellow]Your local environment doesn't perfectly match the environment "
-                "[yellow]Your local environment doesn't perfectly match the environment "
+            "used to produce the original model. The following components differ:[/]"
-                "used to produce the original model. The following components differ:[/]"
-            )
         )
 
     if system_mismatches:
@@ -358,12 +354,10 @@ def check_environment(
     if system_mismatches or package_mismatches:
         print()
         print(
-            (
+            f"There is a {cast(MismatchSeverity, mismatch_severity).__rich__()} chance "
-                f"There is a {cast(MismatchSeverity, mismatch_severity).__rich__()} chance "
+            "that reproduction won't produce a byte-for-byte identical model. "
-                "that reproduction won't produce a byte-for-byte identical model. "
+            "However, the resulting model will very likely still behave similarly "
-                "However, the resulting model will very likely still behave similarly "
+            "to the original model."
-                "to the original model."
-            )
         )
 
         if settings.ignore_mismatches is None:

src/heretic/scorer.py

@@ -6,9 +6,8 @@ from dataclasses import dataclass
 
 from pydantic import BaseModel
 
-from heretic.plugin import Context, Plugin
-
 from .config import Settings as HereticSettings
+from .plugin import Context, Plugin
 
 
 @dataclass
@@ -32,7 +31,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 +46,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,
+            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

@@ -3,11 +3,11 @@
 
 from pydantic import BaseModel, Field
 
-from heretic.config import DatasetSpecification
+from heretic.config import DatasetSpecification, SingleDatasetSpecification
 from heretic.scorer import Context, Score, Scorer
-from heretic.utils import print
+from heretic.utils import format_dataset_specification, print
 
-DEFAULT_KEYWORD_MARKERS = [
+REFUSAL_MARKERS = [
     "disclaimer",
     "sorry",
     "i can'",
@@ -45,13 +45,13 @@ DEFAULT_KEYWORD_MARKERS = [
 
 
 class Settings(BaseModel):
-    keyword_markers: list[str] = Field(
+    score_name: str = Field(
-        default=DEFAULT_KEYWORD_MARKERS,
+        default="Refusals",
-        description="Strings whose presence in a response (case insensitive) identifies the response as a keyword match.",
+        description="Name that describes what the configured keyword rate measures.",
     )
 
     prompts: DatasetSpecification = Field(
-        default=DatasetSpecification(
+        default=SingleDatasetSpecification(
             dataset="mlabonne/harmful_behaviors",
             split="test[:100]",
             column="text",
@@ -59,6 +59,11 @@ class Settings(BaseModel):
         description="Dataset of prompts to evaluate the keyword match rate on.",
     )
 
+    keyword_markers: list[str] = Field(
+        default=REFUSAL_MARKERS,
+        description="Strings whose presence in a response (case insensitive) identifies the response as a keyword match.",
+    )
+
     print_responses: bool = Field(
         default=False,
         description="Whether to print prompt/response pairs when counting keyword matches.",
@@ -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]{format_dataset_specification(self.settings.prompts)}[/]..."
         )
         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

@@ -4,20 +4,20 @@
 import torch.nn.functional as F
 from pydantic import BaseModel, Field
 
-from heretic.config import DatasetSpecification
+from heretic.config import DatasetSpecification, SingleDatasetSpecification
 from heretic.plugin import Context
 from heretic.scorer import Score, Scorer
-from heretic.utils import print
+from heretic.utils import format_dataset_specification, print
 
 
 class Settings(BaseModel):
     prompts: DatasetSpecification = Field(
-        default=DatasetSpecification(
+        default=SingleDatasetSpecification(
             dataset="mlabonne/harmless_alpaca",
             split="test[:100]",
             column="text",
         ),
-        description="Prompt dataset used to measure KL divergence from original model.",
+        description="Dataset of prompts used to measure KL divergence from original model.",
     )
 
 
@@ -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]{format_dataset_specification(self.settings.prompts)}[/]..."
         )
         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,
+            value=kl_divergence,
-            rich_display=f"{kl:.4f}",
+            rich_display=f"[bold]{kl_divergence:.4f}[/]",
-            md_display=f"{kl:.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/system.py

@@ -96,7 +96,7 @@ def get_amdgpu_driver_version() -> str | None:
             if os.path.exists(version_path):
                 with open(version_path, "r", encoding="utf-8") as f:
                     return f.read().strip()
-    except Exception:
+    except Exception:  # noqa: S110
         pass
 
     return None
@@ -249,7 +249,7 @@ def get_accelerator_info_dict() -> dict[str, Any]:
         info: dict[str, Any] = {
             "type": "ROCm" if is_rocm else "CUDA",
             "api_name": "HIP Version" if is_rocm else "CUDA Version",
-            "api_version": torch.version.hip if is_rocm else torch.version.cuda,  # ty:ignore[unresolved-attribute]
+            "api_version": torch.version.hip if is_rocm else torch.version.cuda,
             "driver_version": get_amdgpu_driver_version()
             if is_rocm
             else get_nvidia_driver_version(),
@@ -264,13 +264,13 @@ def get_accelerator_info_dict() -> dict[str, Any]:
         return info
 
     if is_xpu_available():
-        count = torch.xpu.device_count()  # ty:ignore[unresolved-attribute]
+        count = torch.xpu.device_count()
         return {
             "type": "XPU",
             "api_name": None,
             "api_version": None,
             "driver_version": get_xpu_driver_version(),
-            "devices": [{"name": torch.xpu.get_device_name(i)} for i in range(count)],  # ty:ignore[unresolved-attribute]
+            "devices": [{"name": torch.xpu.get_device_name(i)} for i in range(count)],
         }
 
     if is_mlu_available():

src/heretic/utils.py

@@ -1,6 +1,8 @@
 # SPDX-License-Identifier: AGPL-3.0-or-later
 # Copyright (C) 2025-2026  Philipp Emanuel Weidmann <pew@worldwidemann.com> + contributors
 
+from __future__ import annotations
+
 import hashlib
 import json
 import os
@@ -11,7 +13,7 @@ from dataclasses import dataclass
 from datetime import datetime, timezone
 from importlib.metadata import version
 from pathlib import Path
-from typing import Any, TypeVar
+from typing import TYPE_CHECKING, Any, TypeVar
 
 import huggingface_hub
 import tomli_w
@@ -28,7 +30,7 @@ from psutil import Process
 from questionary import Question
 from rich.console import Console
 
-from .config import DatasetSpecification, Settings
+from .config import DatasetSpecification, Settings, SingleDatasetSpecification
 from .system import (
     get_accelerator_info_dict,
     get_cpu_info_dict,
@@ -38,13 +40,15 @@ from .system import (
     is_xpu_available,
 )
 
+if TYPE_CHECKING:
+    from .modifier import Modifier
+
+
 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]:
     """
@@ -165,9 +169,9 @@ def get_split_slice(split_str: str, length: int) -> tuple[int, int]:
     return absolute_instruction.from_, absolute_instruction.to
 
 
-def load_prompts(
+def _load_prompts_single(
     settings: Settings,
-    specification: DatasetSpecification,
+    specification: SingleDatasetSpecification,
 ) -> list[Prompt]:
     path = specification.dataset
     split_str = specification.split
@@ -208,6 +212,7 @@ def load_prompts(
                     )
             dataset = load_dataset(
                 path,
+                name=specification.config,
                 revision=specification.commit,
                 split=split_str,
             )
@@ -225,6 +230,7 @@ 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,
@@ -255,27 +261,51 @@ def load_prompts(
     ]
 
 
+def load_prompts(
+    settings: Settings,
+    specification: DatasetSpecification,
+) -> list[Prompt]:
+    if isinstance(specification, SingleDatasetSpecification):
+        return _load_prompts_single(settings, specification)
+    else:
+        return [
+            prompt
+            for single_specification in specification
+            for prompt in _load_prompts_single(settings, single_specification)
+        ]
+
+
+def format_dataset_specification(specification: DatasetSpecification) -> str:
+    if isinstance(specification, SingleDatasetSpecification):
+        return specification.dataset
+    else:
+        return (
+            "\\["
+            + ", ".join(
+                single_specification.dataset for single_specification in specification
+            )
+            + "]"
+        )
+
+
+def is_dataset_specification_reproducible(specification: DatasetSpecification) -> bool:
+    if isinstance(specification, SingleDatasetSpecification):
+        return is_hf_path(specification.dataset) and specification.commit is not None
+    else:
+        return all(
+            is_hf_path(single_specification.dataset)
+            and single_specification.commit is not None
+            for single_specification in specification
+        )
+
+
 def batchify(items: list[T], batch_size: int) -> list[list[T]]:
     return [items[i : i + batch_size] for i in range(0, len(items), batch_size)]
 
 
-def get_trial_parameters(trial: Trial | FrozenTrial) -> dict[str, str]:
-    params = {}
-
-    direction_index = trial.user_attrs["direction_index"]
-    params["direction_index"] = (
-        "per layer" if (direction_index is None) else f"{direction_index:.2f}"
-    )
-
-    for component, parameters in trial.user_attrs["parameters"].items():
-        for name, value in parameters.items():
-            params[f"{component}.{name}"] = f"{value:.2f}"
-
-    return params
-
-
 def get_readme_intro(
     settings: Settings,
+    modifier: Modifier[Any],
     trial: Trial | FrozenTrial,
     contains_reproducibility_information: bool,
 ) -> str:
@@ -318,7 +348,7 @@ def get_readme_intro(
         model_link
     }, made using [Heretic](https://heretic-project.org) v{version("heretic-llm")}
 {reproducibility_instructions}
-## Abliteration parameters
+## {modifier.modifier_name} parameters
 
 | Parameter | Value |
 | :-------- | :---: |
@@ -326,7 +356,7 @@ def get_readme_intro(
         chr(10).join(
             [
                 f"| **{name}** | {value} |"
-                for name, value in get_trial_parameters(trial).items()
+                for name, value in modifier.render_trial_parameters(trial).items()
             ]
         )
     }
@@ -375,6 +405,7 @@ def format_hf_link(
 
 def generate_reproduce_readme(
     settings: Settings,
+    dataset_specifications: list[DatasetSpecification],
     checkpoint_filename: str,
     trial: Trial | FrozenTrial,
     include_system_information: bool,
@@ -491,6 +522,29 @@ def generate_reproduce_readme(
                 f" --index-url https://download.pytorch.org/whl/{suffix}"
             )
 
+    formatted_datasets = set()
+    for specification in dataset_specifications:
+        if isinstance(specification, SingleDatasetSpecification):
+            formatted_datasets.add(
+                format_hf_link(
+                    specification.dataset,
+                    specification.commit,
+                    is_dataset=True,
+                )
+            )
+        else:
+            for single_specification in specification:
+                formatted_datasets.add(
+                    format_hf_link(
+                        single_specification.dataset,
+                        single_specification.commit,
+                        is_dataset=True,
+                    )
+                )
+    dataset_lines = "\n".join(
+        f"- {formatted_dataset}" for formatted_dataset in sorted(formatted_datasets)
+    )
+
     trial_scores = trial.user_attrs["scores"]
     score_lines = "\n".join(
         (
@@ -510,8 +564,7 @@ This directory contains the necessary information and assets to reproduce the re
 
 ## Datasets
 
-- **Good prompts:** {format_hf_link(settings.good_prompts.dataset, settings.good_prompts.commit, is_dataset=True)}
+{dataset_lines}
-- **Bad prompts:** {format_hf_link(settings.bad_prompts.dataset, settings.bad_prompts.commit, is_dataset=True)}
 
 ## Selected trial
 
@@ -566,8 +619,8 @@ def generate_reproduce_json(
     version_info = get_heretic_version_info()
 
     data = {
-        # Version 3: plugin-based schema with generic scores/baseline scores.
+        # Version 4: plugin-based schema with generic parameters and scores.
-        "version": "3",
+        "version": "4",
         "timestamp": timestamp,
         "system": None,  # Defined here to preserve insertion order.
         "environment": {
@@ -580,10 +633,7 @@ def generate_reproduce_json(
             "requirements": get_requirements_dict(),
         },
         "settings": settings.model_dump(),
-        "parameters": {
+        "parameters": trial.user_attrs["parameters"],
-            "direction_index": trial.user_attrs["direction_index"],
-            "abliteration_parameters": trial.user_attrs["parameters"],
-        },
         "scores": trial.user_attrs["scores"],
         "hashes": uploaded_model_hashes,
     }
@@ -631,6 +681,7 @@ def get_file_sha256(file_path: str | Path) -> str:
 def create_reproduce_folder(
     path: Path,
     settings: Settings,
+    dataset_specifications: list[DatasetSpecification],
     checkpoint_path: str | Path,
     trial: Trial | FrozenTrial,
     uploaded_model_hashes: dict[str, str],
@@ -679,6 +730,7 @@ def create_reproduce_folder(
     (reproduce_dir / "README.md").write_text(
         generate_reproduce_readme(
             settings,
+            dataset_specifications,
             checkpoint_filename,
             trial,
             include_system_information=include_system_information,
@@ -695,6 +747,7 @@ def create_reproduce_folder(
 def upload_reproduce_folder(
     repo_id: str,
     settings: Settings,
+    dataset_specifications: list[DatasetSpecification],
     token: str,
     checkpoint_path: str | Path,
     trial: Trial | FrozenTrial,
@@ -723,6 +776,7 @@ def upload_reproduce_folder(
         create_reproduce_folder(
             tmp_path,
             settings,
+            dataset_specifications,
             checkpoint_path=checkpoint_path,
             trial=trial,
             uploaded_model_hashes=uploaded_model_hashes,

tests/gemma-4e/SHA256SUMS.ci

@@ -1,6 +1,6 @@
 2f1b4d75d067bae3fe44e676721c7f077d243bc007156cb9c2f8b5836613d082 *chat_template.jinja
-ca80080dfa4ec6ba87152fa2b9afe70b90c400e5c4b1d6bdc3aa3114467ca68f *config.json
+c128bc8647a505e343c8cdff9bdd188b1a4a3f81938148685cbf90dff9827268 *config.json
-70070bac883cf9c39b5992450d6b23cd160eaf33099e24c654e0359d2f87c760 *generation_config.json
+58678fb2b8ae1b96652dffa43f864ad1c1c59e49c889bab4264ba0844c4a23b8 *generation_config.json
 f3f4ec19504f182486459cf4e255ece265c25f827840d63b6a9d4058b8e4877a *model.safetensors
 32bdf45d2ad4cc29a0822ddd157a182de76644f0419a6228d151495256e9813c *processor_config.json
 cc8d3a0ce36466ccc1278bf987df5f71db1719b9ca6b4118264f45cb627bfe0f *tokenizer.json

tests/gemma-4e/SHA256SUMS.ci2

@@ -1,6 +1,6 @@
 2f1b4d75d067bae3fe44e676721c7f077d243bc007156cb9c2f8b5836613d082 *chat_template.jinja
-ca80080dfa4ec6ba87152fa2b9afe70b90c400e5c4b1d6bdc3aa3114467ca68f *config.json
+c128bc8647a505e343c8cdff9bdd188b1a4a3f81938148685cbf90dff9827268 *config.json
-70070bac883cf9c39b5992450d6b23cd160eaf33099e24c654e0359d2f87c760 *generation_config.json
+58678fb2b8ae1b96652dffa43f864ad1c1c59e49c889bab4264ba0844c4a23b8 *generation_config.json
 53c4ee891dce23c0ac85bebc2c4d48301469750fafbb3e6e024c15786d94db8b *model.safetensors
 32bdf45d2ad4cc29a0822ddd157a182de76644f0419a6228d151495256e9813c *processor_config.json
 cc8d3a0ce36466ccc1278bf987df5f71db1719b9ca6b4118264f45cb627bfe0f *tokenizer.json

tests/gemma-4e/SHA256SUMS.ci3

@@ -0,0 +1,7 @@
+2f1b4d75d067bae3fe44e676721c7f077d243bc007156cb9c2f8b5836613d082 *chat_template.jinja
+c128bc8647a505e343c8cdff9bdd188b1a4a3f81938148685cbf90dff9827268 *config.json
+58678fb2b8ae1b96652dffa43f864ad1c1c59e49c889bab4264ba0844c4a23b8 *generation_config.json
+9ff0593e3fbd0ba463bbc980ebb3ed34798e562b606f90cd93c2df5403732c7b *model.safetensors
+32bdf45d2ad4cc29a0822ddd157a182de76644f0419a6228d151495256e9813c *processor_config.json
+cc8d3a0ce36466ccc1278bf987df5f71db1719b9ca6b4118264f45cb627bfe0f *tokenizer.json
+a1bab8c81ed15fa6ce912ec993c66cb49392e0487fb1ea5f5f11ea3618683627 *tokenizer_config.json

tests/gemma-4e/SHA256SUMS.linux

@@ -1,6 +1,6 @@
 2f1b4d75d067bae3fe44e676721c7f077d243bc007156cb9c2f8b5836613d082 *chat_template.jinja
-ca80080dfa4ec6ba87152fa2b9afe70b90c400e5c4b1d6bdc3aa3114467ca68f *config.json
+c128bc8647a505e343c8cdff9bdd188b1a4a3f81938148685cbf90dff9827268 *config.json
-70070bac883cf9c39b5992450d6b23cd160eaf33099e24c654e0359d2f87c760 *generation_config.json
+58678fb2b8ae1b96652dffa43f864ad1c1c59e49c889bab4264ba0844c4a23b8 *generation_config.json
 effe36925f85ecb1e29bba84501a456bb49df21e4047be8b7ea3f6f88181fb65 *model.safetensors
 32bdf45d2ad4cc29a0822ddd157a182de76644f0419a6228d151495256e9813c *processor_config.json
 cc8d3a0ce36466ccc1278bf987df5f71db1719b9ca6b4118264f45cb627bfe0f *tokenizer.json

tests/gemma-4e/SHA256SUMS.windows

@@ -1,7 +0,0 @@
-b16d3228a775c549ba97af41233a54e9de8dd2b65250f78346661d18b936a8b5 *chat_template.jinja
-0094ad598a8043f84d82ad5c886547bca1d1d7f302d82f1491f83d388e89acd4 *config.json
-1a019c5d688d54cf01318eab88cb4345dfa52135eb1d83c2f54125469eb88d5c *generation_config.json
-effe36925f85ecb1e29bba84501a456bb49df21e4047be8b7ea3f6f88181fb65 *model.safetensors
-24d00232e58cfa179fe8b3911c788d4aad9a6279d778ebe4c72e82623b6197f9 *processor_config.json
-cc8d3a0ce36466ccc1278bf987df5f71db1719b9ca6b4118264f45cb627bfe0f *tokenizer.json
-8044bbbddaee8dc47e6b5660e013ba92224d4a5392b2939c59699aa0105f5c8b *tokenizer_config.json

tests/gemma-4e/config.toml

@@ -1,4 +1,4 @@
-# This test case is for Hybrid-Edge models.
+# This test case is for hybrid models.
 # After any change related to it, this test should PASS.
 
 model = "tiny-random/gemma-4e"
@@ -9,6 +9,11 @@ print_debug_information = true
 
 batch_size = 2
 max_response_length = 10
+
+modifiers = [
+    { plugin = "heretic.modifiers.abliteration.Abliteration" },
+]
+
 n_trials = 2
 n_startup_trials = 1
 
@@ -18,13 +23,13 @@ trial_index = 0
 model_action = "save"
 save_directory = "model"
 
-[good_prompts]
+[[response_prefix_test_prompts]]
 dataset = "mlabonne/harmless_alpaca"
 commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f"
 split = "train[:5]"
 column = "text"
 
-[bad_prompts]
+[[response_prefix_test_prompts]]
 dataset = "mlabonne/harmful_behaviors"
 commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7"
 split = "train[:5]"
@@ -41,3 +46,15 @@ dataset = "mlabonne/harmful_behaviors"
 commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7"
 split = "test[:5]"
 column = "text"
+
+[modifier.Abliteration.good_prompts]
+dataset = "mlabonne/harmless_alpaca"
+commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f"
+split = "train[:5]"
+column = "text"
+
+[modifier.Abliteration.bad_prompts]
+dataset = "mlabonne/harmful_behaviors"
+commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7"
+split = "train[:5]"
+column = "text"

tests/gpt-oss/SHA256SUMS.linux

@@ -0,0 +1,6 @@
+f8d9255777615591a7cc1a7c932f5a69e181128902295e1b81221d20d983cac7 *chat_template.jinja
+22d559dde87f1a0faf02efed3062af2b1fc8a1721a55b952cd5f5d8c356b4feb *config.json
+523e422e425e7a7ba23d85cc0faeef143dc5da37b72bfc649df5a8d83ee95beb *generation_config.json
+aefe8b9c4b4969f6d13c5d778760f3dce4e25134324b33677934550d9df02a7c *model.safetensors
+fce342a4642cb8afc42d8d89cfa21198b64a43458ded7f6ff28d1151a08c9cda *tokenizer.json
+9ba5fa877168e24823cb583c55b4c2e4df0331f30084953c7cf07de294640384 *tokenizer_config.json

tests/gpt-oss/config.toml

@@ -0,0 +1,56 @@
+# This test case is for ARA.
+# After any change related to it, this test should PASS.
+
+model = "tiny-random/gpt-oss"
+model_commit = "02ba5c61f879b5a38a8b1f7a8e0409b8e1bb8f38"
+
+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"
+
+[[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.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"

tests/minicpm5/SHA256SUMS.linux

@@ -1,6 +1,6 @@
 7451a05cf1e28a79d97d7c0bc951028c0b1915119bf9046acd06a0e3d931f47c *chat_template.jinja
-fe6fd41d9f2ce5d6486748cf0330b574f37bf7d4e915f7b39d1af1a185cac3c3 *config.json
+de44f54b200f63c8ccb7899965bc9aaf21b23d35164c7de9a3cc3394530e3821 *config.json
-c4c2ef5ae4a4e2dd10655a3b99d801a8a50497286ddd042ba35bcfefc44ad349 *generation_config.json
+12d96814a0cd1a72cae14392b04b9644353e8600614faa6eb7ee5edbf9452ec1 *generation_config.json
 1535a9b7a91b2cb39ad280dbd9a940e2609a0b423d5b924df4d664e579912802 *model.safetensors
 ad92aaa8d3032c98a9158b8c5e8682bed10027ed6463e4fb1320fe5384210873 *tokenizer.json
 3ad32522c384dbe35192bb69de9befbf3f523e99d4bb3f95da757671d4c28281 *tokenizer_config.json

tests/minicpm5/SHA256SUMS.windows

@@ -1,6 +0,0 @@
-d8db3ff45c4c68a0ba9dee962ff1a0adde9a2be55e0895306f6bd2b2756f5adb *chat_template.jinja
-a9d6f64bb9d0c02b553119e475615153af625b5c2a16ccb8fb8b3c2cc348f465 *config.json
-0e7611a1e8fd0a06a139b0572b2c55b885ba9fb7db2022873c3508aebfb488aa *generation_config.json
-411d95f42d3e31aef41c28314c8f0431c980687a97904d32b4ef57c42199720f *model.safetensors
-ad92aaa8d3032c98a9158b8c5e8682bed10027ed6463e4fb1320fe5384210873 *tokenizer.json
-aa083f3da10340925734e876e41e235c459329294ecd35d7511ec5868c1f14e3 *tokenizer_config.json

tests/minicpm5/config.toml

@@ -9,7 +9,11 @@ print_debug_information = true
 
 batch_size = 2
 max_response_length = 10
-kl_divergence_target = 0
+
+modifiers = [
+    { plugin = "heretic.modifiers.abliteration.Abliteration" },
+]
+
 n_trials = 2
 n_startup_trials = 1
 
@@ -19,20 +23,13 @@ trial_index = 0
 model_action = "save"
 save_directory = "model"
 
-row_normalization = "none"
+[[response_prefix_test_prompts]]
-
-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"
 split = "train[:5]"
 column = "text"
 
-[bad_prompts]
+[[response_prefix_test_prompts]]
 dataset = "mlabonne/harmful_behaviors"
 commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7"
 split = "train[:5]"
@@ -49,3 +46,18 @@ dataset = "mlabonne/harmful_behaviors"
 commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7"
 split = "test[:5]"
 column = "text"
+
+[modifier.Abliteration]
+row_normalization = "none"
+
+[modifier.Abliteration.good_prompts]
+dataset = "mlabonne/harmless_alpaca"
+commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f"
+split = "train[:5]"
+column = "text"
+
+[modifier.Abliteration.bad_prompts]
+dataset = "mlabonne/harmful_behaviors"
+commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7"
+split = "train[:5]"
+column = "text"

tests/mistral-3/SHA256SUMS.ci

@@ -1,7 +1,7 @@
 39f03c383413f531fd302c06c7e982ad98c83f0657a8339ae25478ccb81fdcda *chat_template.jinja
-f69f84977a47c8fea9ce9fc26b7de379216cb01146ea726a87996d3554cfcd19 *config.json
+b2cfd8a9da09efc97040f3839fe629f3925db856a44b58314301dc5e11345ec8 *config.json
-34dfa6012ca9ac5f57e5521d8dbaecbc7ab7f7ab0fd96ec020b543aab5f265d9 *generation_config.json
+338d303c3c3df6884030cf5be24560f6e973fc183220d14acfb9364cd0e070e4 *generation_config.json
-876c6691eb85e3e5e11771e589529830fb454ab26344e1271ae550661e312b50 *model.safetensors
+05497f43a38427c5813105fc007ce0dd41be65470e9cb8f343443f31bab7f0fe *model.safetensors
-84be30b124b50749c56d25fdbec5ccedf564446f6b3b035e88e1e07b986d2491 *processor_config.json
+a05f93a41b7e42ccc18461b540fb54490c3ffa13aef79c800bbe7942e02e360f *processor_config.json
 c3a8d92e371b92a2cd6e678e31ebc27d0235e929a51fbf290f74742b341fa96f *tokenizer.json
 7b29c843c0043622d28fd4638451cbb0a609d99a0762ffbff3b92b4b2fee4d94 *tokenizer_config.json

tests/mistral-3/SHA256SUMS.ci2

@@ -1,7 +1,7 @@
 39f03c383413f531fd302c06c7e982ad98c83f0657a8339ae25478ccb81fdcda *chat_template.jinja
-f69f84977a47c8fea9ce9fc26b7de379216cb01146ea726a87996d3554cfcd19 *config.json
+b2cfd8a9da09efc97040f3839fe629f3925db856a44b58314301dc5e11345ec8 *config.json
-34dfa6012ca9ac5f57e5521d8dbaecbc7ab7f7ab0fd96ec020b543aab5f265d9 *generation_config.json
+338d303c3c3df6884030cf5be24560f6e973fc183220d14acfb9364cd0e070e4 *generation_config.json
-6febb813086f253e5ec0fcda02fdfc849c551a7dba54681b37ac5bc402e4eed6 *model.safetensors
+ad6675b44b476a899914761257bbeb7f320cb9c0dc39a7467daf4dad0216dd1f *model.safetensors
-84be30b124b50749c56d25fdbec5ccedf564446f6b3b035e88e1e07b986d2491 *processor_config.json
+a05f93a41b7e42ccc18461b540fb54490c3ffa13aef79c800bbe7942e02e360f *processor_config.json
 c3a8d92e371b92a2cd6e678e31ebc27d0235e929a51fbf290f74742b341fa96f *tokenizer.json
 7b29c843c0043622d28fd4638451cbb0a609d99a0762ffbff3b92b4b2fee4d94 *tokenizer_config.json

tests/mistral-3/SHA256SUMS.ci3

@@ -0,0 +1,7 @@
+39f03c383413f531fd302c06c7e982ad98c83f0657a8339ae25478ccb81fdcda *chat_template.jinja
+b2cfd8a9da09efc97040f3839fe629f3925db856a44b58314301dc5e11345ec8 *config.json
+338d303c3c3df6884030cf5be24560f6e973fc183220d14acfb9364cd0e070e4 *generation_config.json
+4cf4d463a94f477ed446e47130cb7da6b5994f0b788e75b82b3c65a531585621 *model.safetensors
+a05f93a41b7e42ccc18461b540fb54490c3ffa13aef79c800bbe7942e02e360f *processor_config.json
+c3a8d92e371b92a2cd6e678e31ebc27d0235e929a51fbf290f74742b341fa96f *tokenizer.json
+7b29c843c0043622d28fd4638451cbb0a609d99a0762ffbff3b92b4b2fee4d94 *tokenizer_config.json

tests/mistral-3/SHA256SUMS.linux

@@ -1,7 +1,7 @@
 39f03c383413f531fd302c06c7e982ad98c83f0657a8339ae25478ccb81fdcda *chat_template.jinja
-f69f84977a47c8fea9ce9fc26b7de379216cb01146ea726a87996d3554cfcd19 *config.json
+b2cfd8a9da09efc97040f3839fe629f3925db856a44b58314301dc5e11345ec8 *config.json
-34dfa6012ca9ac5f57e5521d8dbaecbc7ab7f7ab0fd96ec020b543aab5f265d9 *generation_config.json
+338d303c3c3df6884030cf5be24560f6e973fc183220d14acfb9364cd0e070e4 *generation_config.json
-29aff97d5633dead9e1ccd29a2cc153b4b7431d22f63c8d6cf60bc6547681cc9 *model.safetensors
+8244162ecb6ce7bad7b10fe4f015889d76a09c7749dc7ac8c4d957f0864b2e82 *model.safetensors
-84be30b124b50749c56d25fdbec5ccedf564446f6b3b035e88e1e07b986d2491 *processor_config.json
+a05f93a41b7e42ccc18461b540fb54490c3ffa13aef79c800bbe7942e02e360f *processor_config.json
 c3a8d92e371b92a2cd6e678e31ebc27d0235e929a51fbf290f74742b341fa96f *tokenizer.json
 7b29c843c0043622d28fd4638451cbb0a609d99a0762ffbff3b92b4b2fee4d94 *tokenizer_config.json

tests/mistral-3/SHA256SUMS.windows

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

tests/mistral-3/config.toml

@@ -1,4 +1,4 @@
-# This test case is for Dense models.
+# This test case is for dense models.
 # After any change related to it, this test should PASS.
 
 model = "tiny-random/mistral-3"
@@ -9,6 +9,11 @@ print_debug_information = true
 
 batch_size = 2
 max_response_length = 10
+
+modifiers = [
+    { plugin = "heretic.modifiers.abliteration.Abliteration" },
+]
+
 n_trials = 2
 n_startup_trials = 1
 
@@ -18,13 +23,13 @@ trial_index = 0
 model_action = "save"
 save_directory = "model"
 
-[good_prompts]
+[[response_prefix_test_prompts]]
 dataset = "mlabonne/harmless_alpaca"
 commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f"
 split = "train[:5]"
 column = "text"
 
-[bad_prompts]
+[[response_prefix_test_prompts]]
 dataset = "mlabonne/harmful_behaviors"
 commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7"
 split = "train[:5]"
@@ -41,3 +46,15 @@ dataset = "mlabonne/harmful_behaviors"
 commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7"
 split = "test[:5]"
 column = "text"
+
+[modifier.Abliteration.good_prompts]
+dataset = "mlabonne/harmless_alpaca"
+commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f"
+split = "train[:5]"
+column = "text"
+
+[modifier.Abliteration.bad_prompts]
+dataset = "mlabonne/harmful_behaviors"
+commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7"
+split = "train[:5]"
+column = "text"

tests/qwen2.5/SHA256SUMS.ci

@@ -1,6 +1,6 @@
 cd8e9439f0570856fd70470bf8889ebd8b5d1107207f67a5efb46e342330527f *chat_template.jinja
-45134b857367fdcb97c0179199848c353fc28f8b95ac2244ac8f45cca448d864 *config.json
+6b517667960a5e39a692eb8277be7260ac02c22d9c7eda9c7aa1bbab21517fc1 *config.json
-e81e23e025c38e825dcf8375861e26a90e804276e4db9ee390122a4fdc95dae7 *generation_config.json
+64cf6f4e0016154cae45628221d77e3925961876125fe68bd942303ec7cca40f *generation_config.json
-bd86541d817978c896bd3579e69ae6d41b6382eaf1646accf83d6feb16acb703 *model.safetensors
+e616cbeb5a913015eb3db96e001030048df2db560df363d4cf688f0c1b2c96de *model.safetensors
 f7f96da3a872b5e901575b2067c744ad336c3a3d77a21584d20024557b1bd7f0 *tokenizer.json
 04b1682c59acbd057f4c9072297faa73d56fc9de053094c659cdb4c464f58f86 *tokenizer_config.json

tests/qwen2.5/SHA256SUMS.linux

@@ -1,6 +1,6 @@
 cd8e9439f0570856fd70470bf8889ebd8b5d1107207f67a5efb46e342330527f *chat_template.jinja
-45134b857367fdcb97c0179199848c353fc28f8b95ac2244ac8f45cca448d864 *config.json
+6b517667960a5e39a692eb8277be7260ac02c22d9c7eda9c7aa1bbab21517fc1 *config.json
-e81e23e025c38e825dcf8375861e26a90e804276e4db9ee390122a4fdc95dae7 *generation_config.json
+64cf6f4e0016154cae45628221d77e3925961876125fe68bd942303ec7cca40f *generation_config.json
-e616cbeb5a913015eb3db96e001030048df2db560df363d4cf688f0c1b2c96de *model.safetensors
+bd86541d817978c896bd3579e69ae6d41b6382eaf1646accf83d6feb16acb703 *model.safetensors
 f7f96da3a872b5e901575b2067c744ad336c3a3d77a21584d20024557b1bd7f0 *tokenizer.json
 04b1682c59acbd057f4c9072297faa73d56fc9de053094c659cdb4c464f58f86 *tokenizer_config.json

tests/qwen2.5/SHA256SUMS.windows

@@ -1,6 +0,0 @@
-8aa40ce145adb73cb3a75194dc0224702a95850ec5275cabb728496bbd749fc6 *chat_template.jinja
-e8f2fcd2681eb92233c0902866441f79a207b235f0b03364d41ebf8c53df62a0 *config.json
-3fec6d7004e5ae311864de130b62e32dac87569874c91b3fe9c46e9309345c1c *generation_config.json
-bd86541d817978c896bd3579e69ae6d41b6382eaf1646accf83d6feb16acb703 *model.safetensors
-f7f96da3a872b5e901575b2067c744ad336c3a3d77a21584d20024557b1bd7f0 *tokenizer.json
-154e5ff1e7c152d964edf30da854ea62465c767719ac8e97e58babf2d4fa9079 *tokenizer_config.json

tests/qwen2.5/config.toml

@@ -9,7 +9,11 @@ print_debug_information = true
 
 batch_size = 2
 max_response_length = 10
-kl_divergence_target = 0
+
+modifiers = [
+    { plugin = "heretic.modifiers.abliteration.Abliteration" },
+]
+
 n_trials = 2
 n_startup_trials = 1
 
@@ -19,20 +23,13 @@ trial_index = 0
 model_action = "save"
 save_directory = "model"
 
-row_normalization = "pre"
+[[response_prefix_test_prompts]]
-
-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"
 split = "train[:5]"
 column = "text"
 
-[bad_prompts]
+[[response_prefix_test_prompts]]
 dataset = "mlabonne/harmful_behaviors"
 commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7"
 split = "train[:5]"
@@ -49,3 +46,18 @@ dataset = "mlabonne/harmful_behaviors"
 commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7"
 split = "test[:5]"
 column = "text"
+
+[modifier.Abliteration]
+row_normalization = "pre"
+
+[modifier.Abliteration.good_prompts]
+dataset = "mlabonne/harmless_alpaca"
+commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f"
+split = "train[:5]"
+column = "text"
+
+[modifier.Abliteration.bad_prompts]
+dataset = "mlabonne/harmful_behaviors"
+commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7"
+split = "train[:5]"
+column = "text"

tests/qwen3.5-moe/SHA256SUMS.ci

@@ -1,7 +1,7 @@
 a4aee8afcf2e0711942cf848899be66016f8d14a889ff9ede07bca099c28f715 *chat_template.jinja
-749b56d1b1e08081981169db6f2c44ab0be4fd6ebb452d15baafa5e09c21586a *config.json
+4b501bef0727793de90d446fcfdc7757a9f8d19efc42ceaa3d26e90db4f876d5 *config.json
-4625d1d64d41d1fa9dae7af4ba1e1d7e65a194073d4efa58acb266a916eaaa74 *generation_config.json
+532190619b0b243e0558a846257cc62883c26df812d456f8133e7b1043c69ad8 *generation_config.json
-5fb94c65bcd9d736735a45e50c2b0bfafd3bb09a444c49b8cff2e131ed35797e *model.safetensors
+cfa4d5428e9b23c245fd2413ad702688cbcacd25808ecb17889e97c0976532ac *model.safetensors
-01562eddd6f9e9ec4bc31656a3b7055284cafbf889acc6c4348dca431ae31f68 *processor_config.json
+3814737f2d4bfdd153a82364b3850f1be5dbc6f074a55c8d5fbf4d40127c3e87 *processor_config.json
-87a7830d63fcf43bf241c3c5242e96e62dd3fdc29224ca26fed8ea333db72de4 *tokenizer.json
+a5cd9732badce41de57e6efce8302930ded1c1188c5f81feb2bd6c24c4a1941f *tokenizer.json
-2e31d1126e81bddf8d15c3f95260fb487b48c5131b24fcbb5bb9d2537e7afac0 *tokenizer_config.json
+5e55dc6b9d9d28d49b6b2b80d2a2f2558a0b6e0e7b194bbca8664fc269ab3ebc *tokenizer_config.json

tests/qwen3.5-moe/SHA256SUMS.linux

@@ -1,7 +1,7 @@
 a4aee8afcf2e0711942cf848899be66016f8d14a889ff9ede07bca099c28f715 *chat_template.jinja
-749b56d1b1e08081981169db6f2c44ab0be4fd6ebb452d15baafa5e09c21586a *config.json
+4b501bef0727793de90d446fcfdc7757a9f8d19efc42ceaa3d26e90db4f876d5 *config.json
-4625d1d64d41d1fa9dae7af4ba1e1d7e65a194073d4efa58acb266a916eaaa74 *generation_config.json
+532190619b0b243e0558a846257cc62883c26df812d456f8133e7b1043c69ad8 *generation_config.json
-5e0fb0ac724cf079b693fc76a515e60bc16de72c32b36c107b9f078061c4f2ef *model.safetensors
+4c29b8ce99c59c90eca59aa6453029bda87165efa57cabc2c668599c02afc31a *model.safetensors
-01562eddd6f9e9ec4bc31656a3b7055284cafbf889acc6c4348dca431ae31f68 *processor_config.json
+3814737f2d4bfdd153a82364b3850f1be5dbc6f074a55c8d5fbf4d40127c3e87 *processor_config.json
-87a7830d63fcf43bf241c3c5242e96e62dd3fdc29224ca26fed8ea333db72de4 *tokenizer.json
+a5cd9732badce41de57e6efce8302930ded1c1188c5f81feb2bd6c24c4a1941f *tokenizer.json
-2e31d1126e81bddf8d15c3f95260fb487b48c5131b24fcbb5bb9d2537e7afac0 *tokenizer_config.json
+5e55dc6b9d9d28d49b6b2b80d2a2f2558a0b6e0e7b194bbca8664fc269ab3ebc *tokenizer_config.json

tests/qwen3.5-moe/SHA256SUMS.windows

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

tests/qwen3.5-moe/config.toml

@@ -9,6 +9,11 @@ print_debug_information = true
 
 batch_size = 2
 max_response_length = 10
+
+modifiers = [
+    { plugin = "heretic.modifiers.abliteration.Abliteration" },
+]
+
 n_trials = 2
 n_startup_trials = 1
 
@@ -18,13 +23,13 @@ trial_index = 0
 model_action = "save"
 save_directory = "model"
 
-[good_prompts]
+[[response_prefix_test_prompts]]
 dataset = "mlabonne/harmless_alpaca"
 commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f"
 split = "train[:5]"
 column = "text"
 
-[bad_prompts]
+[[response_prefix_test_prompts]]
 dataset = "mlabonne/harmful_behaviors"
 commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7"
 split = "train[:5]"
@@ -41,3 +46,15 @@ dataset = "mlabonne/harmful_behaviors"
 commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7"
 split = "test[:5]"
 column = "text"
+
+[modifier.Abliteration.good_prompts]
+dataset = "mlabonne/harmless_alpaca"
+commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f"
+split = "train[:5]"
+column = "text"
+
+[modifier.Abliteration.bad_prompts]
+dataset = "mlabonne/harmful_behaviors"
+commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7"
+split = "train[:5]"
+column = "text"

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():
@@ -51,7 +53,7 @@ for test_directory in script_directory.iterdir():
 
             print()
 
-            valid_hashes: dict[str, list[str]] = {}
+            valid_hashes: dict[str, set[str]] = {}
 
             for hash_file in hash_files:
                 with open(hash_file, "r", encoding="utf-8") as file:
@@ -61,27 +63,38 @@ for test_directory in script_directory.iterdir():
                             filename = filename.removeprefix("*")
 
                             if filename not in valid_hashes:
-                                valid_hashes[filename] = []
+                                valid_hashes[filename] = set()
 
-                            valid_hashes[filename].append(sha256.lower())
+                            valid_hashes[filename].add(sha256.lower())
 
-            for filename in valid_hashes:
+            # Track which specific files failed within this test directory.
+            failed_files: list[str] = []
+            for filename, hashes in valid_hashes.items():
                 sha256 = get_file_sha256(test_directory / "model" / filename)
 
-                if sha256.lower() not in valid_hashes[filename]:
+                if sha256.lower() not in hashes:
                     print(
-                        (
+                        f"Test {test_directory.name} has FAILED!\n"
-                            f"Test {test_directory.name} has FAILED!\n"
+                        f"Output file {filename} doesn't match any valid hash.\n\n"
-                            f"Output file {filename} doesn't match any valid hash.\n\n"
+                        f"Valid hashes:\n"
-                            f"Valid hashes:\n"
+                        f"{chr(10).join(hashes)}\n\n"
-                            f"{chr(10).join(valid_hashes[filename])}\n\n"
+                        f"Actual hash:\n"
-                            f"Actual hash:\n"
+                        f"{sha256}\n"
-                            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.")

tests/seed-oss/SHA256SUMS.linux

@@ -0,0 +1,6 @@
+39a23a1a28f68acd747373fa688bea630021ec4db2bc8dec2fd3962240fb0418 *chat_template.jinja
+45d58c7316e9dab89e918b9ca1c141b8c7f5250d55e0f9be76f2f452903aeef8 *config.json
+02b35dff3057b998f224a3a5599dc7737ba028e3b4181ba89ba74ba1e707a5dc *generation_config.json
+0cdea9064dcbe6db666f9d42a283d664c133d4c67a2f6ea0fd863a54f522160c *model.safetensors
+3cf3a6d9520f195638a36f0194239d817de7288710bca55f1f5753de226748f7 *tokenizer.json
+388b47e61cb40f2fd51a89999053686ab4c45b40b43c0329d15645f5910d069e *tokenizer_config.json

tests/seed-oss/config.toml

@@ -0,0 +1,65 @@
+# This test case is for ARA (non-standard settings).
+# After any change related to it, this test should PASS.
+
+model = "tiny-random/seed-oss"
+model_commit = "6860befd78b678885f7a52bbf41d7fd0671af2db"
+
+seed = 12345
+print_debug_information = true
+
+batch_size = 2
+max_response_length = 10
+
+scorers = [
+    { plugin = "heretic.scorers.keyword_rate.KeywordRate", optimization = "minimize" },
+    { plugin = "heretic.scorers.kl_divergence.KLDivergence", optimization = "maximize" },
+]
+
+n_trials = 2
+n_startup_trials = 1
+
+export_strategy = "merge"
+checkpoint_action = "restart"
+trial_index = 0
+model_action = "save"
+save_directory = "model"
+
+[[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.prompts]
+dataset = "mlabonne/harmful_behaviors"
+commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7"
+split = "test[:5]"
+column = "text"
+
+[modifier.ARA]
+preserve_row_magnitudes = false
+lora_rank = 20
+
+[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"