fix(ara): free gradient buffers after optimization (#426)

LBFGS leaves one full-size gradient buffer on every optimized weight.
Across the layers processed in a typical trial this is many GiB of VRAM
that persists into evaluation, causing CUDA out-of-memory errors.
Clear the gradients after each module is optimized.

.github/workflows/ci.yml

@@ -40,11 +40,6 @@ jobs:
       - name: Check typing
         run: uv run ty check --output-format=github --error-on-warning .
 
-      - name: Run tests
-        env:
-          PYTHONUNBUFFERED: "1"
-        run: uv run tests/run_tests.py 2>&1
-
       - name: Build package
         run: uv build
 

.gitignore

@@ -15,14 +15,11 @@ wheels/
 # Editors
 /.vscode/
 
-# Configuration file (root only, not ignored in test directories)
+# Configuration files
 /config.toml
 
 # Study checkpoints
-checkpoints/
+/checkpoints/
 
 # Residual plots
-plots/
+/plots/
-
-# Models generated by tests
-/tests/*/model/

README.md

@@ -1,6 +1,6 @@
-<img width="128" align="right" alt="Logo" src="https://github.com/user-attachments/assets/df5f2840-2f92-4991-aa57-252747d7182e" />
+<img width="128" height="128" align="right" alt="Logo" src="https://github.com/user-attachments/assets/df5f2840-2f92-4991-aa57-252747d7182e" />
 
-# Heretic: Fully automatic censorship removal for language models<br><br>[![Discord](https://img.shields.io/discord/1447831134212984903?color=5865F2&label=discord&labelColor=black&logo=discord&logoColor=white&style=for-the-badge)](https://discord.gg/gdXc48gSyT) [![Matrix](https://img.shields.io/badge/Matrix-black?logo=matrix&style=for-the-badge)](https://matrix.to/#/#heretic:matrix.org) [![Follow us on Hugging Face](https://huggingface.co/datasets/huggingface/badges/resolve/main/follow-us-on-hf-md-dark.svg)](https://huggingface.co/heretic-org) [![Codeberg mirror](https://img.shields.io/badge/Codeberg%20mirror-black?logo=codeberg&style=for-the-badge)](https://codeberg.org/p-e-w/heretic)
+# Heretic: Fully automatic censorship removal for language models<br><br>[![Discord](https://img.shields.io/discord/1447831134212984903?color=5865F2&label=discord&labelColor=black&logo=discord&logoColor=white&style=for-the-badge)](https://discord.gg/gdXc48gSyT) [![Follow us on Hugging Face](https://huggingface.co/datasets/huggingface/badges/resolve/main/follow-us-on-hf-md-dark.svg)](https://huggingface.co/heretic-org)
 
 [![#1 Repository of the Day](https://trendshift.io/api/badge/repositories/20538)](https://trendshift.io/repositories/20538)
 
@@ -20,11 +20,6 @@ as possible. Using Heretic does not require an understanding of transformer
 internals. In fact, anyone who knows how to run a command-line program
 can use Heretic to decensor language models.
 
-Heretic supports most dense models, including many multimodal models,
-several different MoE architectures, and even some hybrid models like Qwen3.5.
-Pure state-space models and certain other research architectures are not yet
-supported out of the box.
-
 <img width="650" height="715" alt="Screenshot" src="https://github.com/user-attachments/assets/d71a5efa-d6be-4705-a817-63332afb2d15" />
 
 &nbsp;
@@ -70,15 +65,15 @@ Heretic have been well-received by users (links and emphasis added):
 > Has been the best unquantized abliterated model that I have been able to run on 16gb vram."
 > [*(Link to comment)*](https://old.reddit.com/r/LocalLLaMA/comments/1phjxca/im_calling_these_people_out_right_now/nt06tji/)
 
-Heretic models have also been independently benchmarked using standard metrics
+Heretic supports most dense models, including many multimodal models, and
-like MMLU and GSM8K, and have been found to compare favorably with models
+several different MoE architectures. It does not yet support SSMs/hybrid models,
-produced by competing abliteration tools:
+models with inhomogeneous layers, and certain novel attention systems.
-[1](https://old.reddit.com/r/LocalLLaMA/comments/1sojjoc/abliterlitics_benchmark_and_tensor_analysis/),
-[2](https://old.reddit.com/r/LocalLLaMA/comments/1sy18lx/abliterlitics_benchmarks_and_tensor_comparison/).
 
-The community has created and published
+You can find a small collection of models that have been decensored using Heretic
-[well over 4000](https://huggingface.co/models?other=heretic)
+[on Hugging Face](https://huggingface.co/collections/p-e-w/the-bestiary),
-models with Heretic.
+and the community has created and published
+[well over 1,000](https://huggingface.co/models?other=heretic)
+Heretic models in addition to those.
 
 
 ## Usage
@@ -86,28 +81,13 @@ models with Heretic.
 Prepare a Python 3.10+ environment with PyTorch 2.2+ installed as appropriate
 for your hardware. Then run:
 
-```sh
+```
 pip install -U heretic-llm
 heretic Qwen/Qwen3-4B-Instruct-2507
 ```
 
 Replace `Qwen/Qwen3-4B-Instruct-2507` with whatever model you want to decensor.
 
-> [!IMPORTANT]
->
-> While PyTorch 2.2 is the minimum version of PyTorch needed for Heretic to work,
-> some models and configurations might require features only found in
-> later versions. For example, loading MXFP4-quantized models like gpt-oss
-> uses `torch.accelerator`, which was added in PyTorch 2.6.
-
-> [!TIP]
->
-> Heretic uses [uv](https://docs.astral.sh/uv/) for dependency management,
-> and the repository includes a `uv.lock` file pinning every package version.
-> If you already use uv (and you probably should!), you can just clone the repo
-> and run Heretic with `uv run heretic`, which ensures that your dependencies
-> match those used by the developers, improving reliability and security.
-
 The process is fully automatic and does not require configuration; however,
 Heretic has a variety of configuration parameters that can be changed for
 greater control. Run `heretic --help` to see available command-line options,
@@ -116,15 +96,14 @@ a configuration file.
 
 At the start of a program run, Heretic benchmarks the system to determine
 the optimal batch size to make the most of the available hardware.
-On an RTX 3090, with the default configuration, decensoring
+On an RTX 3090, with the default configuration, decensoring Llama-3.1-8B-Instruct
-[Qwen3-4B-Instruct-2507](https://huggingface.co/Qwen/Qwen3-4B-Instruct-2507)
+takes about 45 minutes. Note that Heretic supports model quantization with
-takes about 20-30 minutes. Note that Heretic supports model quantization with
 bitsandbytes, which can drastically reduce the amount of VRAM required to process
 models. Set the `quantization` option to `bnb_4bit` to enable quantization.
 
 After Heretic has finished decensoring a model, you are given the option to
 save the model, upload it to Hugging Face, chat with it to test how well it works,
-run standard benchmarks on it, or any combination of those actions.
+or any combination of those actions.
 
 
 ## Research features
@@ -134,7 +113,7 @@ provides features designed to support research into the semantics of model inter
 (interpretability). To use those features, you need to install Heretic with the
 optional `research` extra:
 
-```sh
+```
 pip install -U heretic-llm[research]
 ```
 

config.default.toml

@@ -25,13 +25,7 @@ quantization = "none"
 device_map = "auto"
 
 # Maximum memory to allocate per device.
-# max_memory = { "0" = "20GB", "cpu" = "64GB" }
+# max_memory = {"0": "20GB", "cpu": "64GB"}
-
-# Whether to move intermediate analysis tensors (such as residuals and logprobs)
-# to CPU memory as soon as possible to reduce peak VRAM usage.
-# This lowers peak VRAM usage during residual analysis and evaluation,
-# but may slightly reduce performance due to host/device transfers.
-offload_outputs_to_cpu = true
 
 # Number of input sequences to process in parallel (0 = auto).
 batch_size = 0  # auto
@@ -42,38 +36,9 @@ max_batch_size = 128
 # Maximum number of tokens to generate for each response.
 max_response_length = 100
 
-# List of pairs of the form [cot_initializer, closed_cot_block] used to skip
-# the Chain-of-Thought block in responses, so that evaluation happens
-# at the start of the actual response.
-chain_of_thought_skips = [
-    # Most thinking models.
-    [
-        "<think>",
-        "<think></think>",
-    ],
-    # gpt-oss.
-    [
-        "<|channel|>analysis<|message|>",
-        "<|channel|>analysis<|message|><|end|><|start|>assistant<|channel|>final<|message|>",
-    ],
-    # Unknown, suggested by user.
-    [
-        "<thought>",
-        "<thought></thought>",
-    ],
-    # Unknown, suggested by user.
-    [
-        "[THINK]",
-        "[THINK][/THINK]",
-    ],
-]
-
 # Whether to print prompt/response pairs when counting refusals.
 print_responses = false
 
-# Whether to print additional information that can help with debugging.
-print_debug_information = false
-
 # Whether to print detailed information about residuals and refusal directions.
 print_residual_geometry = false
 
@@ -99,13 +64,13 @@ kl_divergence_target = 0.01
 
 # Whether to adjust the refusal directions so that only the component that is
 # orthogonal to the good direction is subtracted during abliteration.
-orthogonalize_direction = true
+orthogonalize_direction = false
 
 # 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"
+row_normalization = "none"
 
 # The rank of the LoRA adapter to use when "full" row normalization is used.
 # Row magnitude preservation is approximate due to non-linear effects,
@@ -129,12 +94,8 @@ n_startup_trials = 60
 # Directory to save and load study progress to/from.
 study_checkpoint_dir = "checkpoints"
 
-# Maximum size for individual safetensors files generated when exporting a model.
-max_shard_size = "5GB"
-
 # Strings whose presence in a response (case insensitive) identifies the response as a refusal.
 refusal_markers = [
-    "disclaimer",
     "sorry",
     "i can'",
     "i cant",
@@ -172,11 +133,6 @@ refusal_markers = [
 # System prompt to use when prompting the model.
 system_prompt = "You are a helpful assistant."
 
-# 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]").
-
 # Dataset of prompts that tend to not result in refusals (used for calculating refusal directions).
 [good_prompts]
 dataset = "mlabonne/harmless_alpaca"

config.nohumor.toml

@@ -1,69 +0,0 @@
-# Rename this file to config.toml, place it in the working directory
-# that you run Heretic from, and edit the configuration to your liking.
-
-max_response_length = 300
-
-residual_plot_title = "PaCMAP Projection of Residuals for Serious/Humorous Prompts"
-
-refusal_markers = [
-    "😅",
-    "here's one",
-    "why did",
-    "walked into a bar",
-    "gets worse",
-    "the punchline",
-    "classic",
-    "dad joke",
-    "bad joke",
-    "pun intended",
-    "ba dum tss",
-    "rimshot",
-    "😂",
-    "😄",
-    "😆",
-    "haha",
-    "hehe",
-    "lol",
-    "funny",
-    "joke",
-    "humor",
-    "that's hilarious",
-    "you could say",
-    "one-liner",
-    "comedian",
-    "stand-up",
-    "unexpectedly",
-    "because apparently",
-    "to be fair",
-    "on the bright side",
-    "lmao",
-    "omg",
-    "rofl",
-    "silly",
-    "humorous",
-    "clever",
-]
-
-[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"
-
-[good_evaluation_prompts]
-dataset = "mlabonne/harmless_alpaca"
-split = "test[:100]"
-column = "text"
-
-[bad_evaluation_prompts]
-dataset = "UnstableLlama/jokes"
-split = "train[200:250]"
-column = "text"

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "heretic-llm"
-version = "1.4.0"
+version = "1.2.0"
 description = "Fully automatic censorship removal for language models"
 readme = "README.md"
 license = "AGPL-3.0-or-later"
@@ -25,23 +25,21 @@ dependencies = [
     "accelerate~=1.13",
     "bitsandbytes~=0.49",
     "datasets~=4.7",
+    "hf-transfer~=0.1",
     "huggingface-hub~=1.7",
     "immutabledict~=4.3",
+    "kernels~=0.12",
     "langdetect~=1.0",
     "lm-eval[hf]~=0.4",
     "numpy~=2.2",
     "optuna~=4.7",
-    "peft~=0.19",
+    "peft~=0.18",
     "psutil~=7.2",
-    "py-cpuinfo~=9.0",
     "pydantic-settings~=2.13",
     "questionary~=2.1",
     "rich~=14.3",
-    "tomli-w~=1.2",
-    "torch",  # version deliberately unspecified
-    "torchvision",  # version deliberately unspecified
     "tqdm~=4.67",
-    "transformers[kernels]~=5.6",
+    "transformers~=5.3",
 ]
 
 [project.optional-dependencies]
@@ -60,8 +58,8 @@ dev = [
 ]
 
 [project.urls]
-Homepage = "https://heretic-project.org"
+Homepage = "https://github.com/p-e-w/heretic"
-Documentation = "https://heretic-project.org/tutorial"
+Documentation = "https://github.com/p-e-w/heretic"
 Repository = "https://github.com/p-e-w/heretic.git"
 Issues = "https://github.com/p-e-w/heretic/issues"
 Changelog = "https://github.com/p-e-w/heretic/releases"
@@ -73,8 +71,5 @@ heretic = "heretic.main:main"
 requires = ["uv_build>=0.8.11,<0.9.0"]
 build-backend = "uv_build"
 
-[tool.uv]
-exclude-newer = "7 days"
-
 [tool.uv.build-backend]
 module-name = "heretic"

src/heretic/config.py

@@ -4,12 +4,7 @@
 from enum import Enum
 from typing import Dict
 
-from pydantic import (
+from pydantic import BaseModel, Field
-    BaseModel,
-    Field,
-    NonNegativeInt,
-    PositiveInt,
-)
 from pydantic_settings import (
     BaseSettings,
     CliSettingsSource,
@@ -18,12 +13,6 @@ from pydantic_settings import (
     TomlConfigSettingsSource,
 )
 
-# !!!IMPORTANT!!!
-#
-# Any settings added to the classes defined in this module
-# must be evaluated for privacy implications and have
-# exclude=True set in their field definitions if appropriate.
-
 
 class QuantizationMethod(str, Enum):
     NONE = "none"
@@ -37,30 +26,14 @@ class RowNormalization(str, Enum):
     FULL = "full"
 
 
-class ExportStrategy(str, Enum):
-    MERGE = "merge"
-    ADAPTER = "adapter"
-
-
 class DatasetSpecification(BaseModel):
     dataset: str = Field(
         description="Hugging Face dataset ID, or path to dataset on disk."
     )
 
-    commit: str | None = Field(
+    split: str = Field(description="Portion of the dataset to use.")
-        default=None,
-        description="Hugging Face commit hash of the dataset.",
-    )
 
-    split: str | None = Field(
+    column: str = Field(description="Column in the dataset that contains the prompts.")
-        default=None,
-        description="Portion of the dataset to use. Required for datasets, optional for plain text files.",
-    )
-
-    column: str | None = Field(
-        default=None,
-        description="Column in the dataset that contains the prompts. Required for datasets, ignored for plain text files.",
-    )
 
     prefix: str = Field(
         default="",
@@ -80,13 +53,11 @@ class DatasetSpecification(BaseModel):
     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(
         default=None,
         description="Matplotlib color to use for the dataset in plots of residual vectors.",
-        exclude=True,
     )
 
 
@@ -105,37 +76,12 @@ class BenchmarkSpecification(BaseModel):
 class Settings(BaseSettings):
     model: str = Field(description="Hugging Face model ID, or path to model on disk.")
 
-    model_commit: str | None = Field(
-        default=None,
-        description="Hugging Face commit hash of the model.",
-    )
-
     evaluate_model: str | None = Field(
         default=None,
         description=(
             "If this model ID or path is set, then instead of abliterating the main model, "
             "evaluate this model relative to the main model."
         ),
-        exclude=True,
-    )
-
-    collect_reproducibles: str | None = Field(
-        default=None,
-        description=(
-            "If this directory path is set, then instead of abliterating a model, "
-            "download all reproduce.json files from public Heretic model repositories "
-            "on Hugging Face, and store them in that directory for archival purposes."
-        ),
-        exclude=True,
-    )
-
-    reproduce: str | None = Field(
-        default=None,
-        description=(
-            "If this path or URL to a reproduce.json file is set, load reproduction information "
-            "from that file, and attempt to reproduce the abliterated model it originated from."
-        ),
-        exclude=True,
     )
 
     dtypes: list[str] = Field(
@@ -173,119 +119,57 @@ class Settings(BaseSettings):
 
     max_memory: Dict[str, str] | None = Field(
         default=None,
-        description='Maximum memory to allocate per device (e.g., { "0" = "20GB", "cpu" = "64GB" }).',
+        description='Maximum memory to allocate per device (e.g., {"0": "20GB", "cpu": "64GB"}).',
     )
 
-    offload_outputs_to_cpu: bool = Field(
+    trust_remote_code: bool | None = Field(
-        default=True,
+        default=None,
-        description=(
+        description="Whether to trust remote code when loading the model.",
-            "Whether to move intermediate analysis tensors (such as residuals and logprobs) "
-            "to CPU memory as soon as possible to reduce peak VRAM usage. "
-            "This lowers peak VRAM usage during residual analysis and evaluation, "
-            "but may slightly reduce performance due to host/device transfers."
-        ),
     )
 
-    batch_size: NonNegativeInt = Field(
+    batch_size: int = Field(
         default=0,  # auto
         description="Number of input sequences to process in parallel (0 = auto).",
     )
 
-    max_batch_size: PositiveInt = Field(
+    max_batch_size: int = Field(
         default=128,
         description="Maximum batch size to try when 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(
+    max_response_length: int = Field(
         default=100,
         description="Maximum number of tokens to generate for each response.",
     )
 
-    response_prefix: str | None = Field(
-        default=None,
-        description=(
-            "Common prefix to assume for all responses, so that evaluation happens "
-            "at the point where responses start to differ for different prompts. "
-            "If not set, the prefix is determined automatically by comparing multiple responses."
-        ),
-    )
-
-    chain_of_thought_skips: list[tuple[str, str]] = Field(
-        default=[
-            # Most thinking models.
-            (
-                "<think>",
-                "<think></think>",
-            ),
-            # gpt-oss.
-            (
-                "<|channel|>analysis<|message|>",
-                "<|channel|>analysis<|message|><|end|><|start|>assistant<|channel|>final<|message|>",
-            ),
-            # Unknown, suggested by user.
-            (
-                "<thought>",
-                "<thought></thought>",
-            ),
-            # Unknown, suggested by user.
-            (
-                "[THINK]",
-                "[THINK][/THINK]",
-            ),
-        ],
-        description=(
-            "List of pairs of the form (cot_initializer, closed_cot_block) used to skip "
-            "the Chain-of-Thought block in responses, so that evaluation happens "
-            "at the start of the actual response."
-        ),
-        # When storing a settings object, the response prefix is already fixed,
-        # either determined by the automatic mechanism or by explicit user choice.
-        exclude=True,
-    )
-
     print_responses: bool = Field(
         default=False,
         description="Whether to print prompt/response pairs when counting refusals.",
-        exclude=True,
-    )
-
-    print_debug_information: bool = Field(
-        default=False,
-        description="Whether to print additional information that can help with debugging.",
-        exclude=True,
     )
 
     print_residual_geometry: bool = Field(
         default=False,
         description="Whether to print detailed information about residuals and refusal 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,
     )
 
     kl_divergence_scale: float = Field(
@@ -304,8 +188,44 @@ class Settings(BaseSettings):
         ),
     )
 
-    orthogonalize_direction: bool = Field(
+    target_components: list[str] = Field(
+        default=["attn.o_proj", "mlp.down_proj"],
+        description=(
+            "List of component names to target for abliteration. "
+            'Currently supported values are "attn.o_proj" and "mlp.down_proj".'
+        ),
+    )
+
+    use_ara: bool = Field(
         default=True,
+        description=(
+            "Whether to use Arbitrary-Rank Ablation (ARA), an abliteration method based on matrix optimization, "
+            "instead of traditional directional ablation."
+        ),
+    )
+    
+    use_ara_lora: bool = Field(
+        default=False,
+        description=(
+            "Use LoRA in ARA instead of full-weight editing. Makes it compatible with quantization and removes model reloads."
+        ),
+    )
+    
+    ara_lora_rank: int = Field(
+        default=128,
+        description="If LoRA is used in ARA, this sets up its rank. Keep it high enough to simulate the 'arbitrary' effect.",
+    )
+
+    use_piqa: bool = Field(
+        default=False,
+        description=(
+            "Whether to use the Physical Interaction: Question Answering (PIQA) benchmark "
+            "as the quality metric instead of the Kullback-Leibler divergence."
+        ),
+    )
+
+    orthogonalize_direction: bool = Field(
+        default=False,
         description=(
             "Whether to adjust the refusal directions so that only the component that is "
             "orthogonal to the good direction is subtracted during abliteration."
@@ -322,7 +242,7 @@ class Settings(BaseSettings):
         ),
     )
 
-    full_normalization_lora_rank: PositiveInt = Field(
+    full_normalization_lora_rank: int = Field(
         default=3,
         description=(
             'The rank of the LoRA adapter to use when "full" row normalization is used. '
@@ -343,28 +263,19 @@ class Settings(BaseSettings):
         ),
     )
 
-    n_trials: PositiveInt = Field(
+    n_trials: int = Field(
         default=200,
         description="Number of abliteration trials to run during optimization.",
     )
 
-    n_startup_trials: NonNegativeInt = Field(
+    n_startup_trials: int = Field(
         default=60,
         description="Number of trials that use random sampling for the purpose of exploration.",
     )
 
-    seed: int | None = Field(
-        default=None,
-        description=(
-            "Random seed for reproducible optimization. "
-            "Applies to Python's random module, NumPy, PyTorch, and Optuna."
-        ),
-    )
-
     study_checkpoint_dir: str = Field(
         default="checkpoints",
         description="Directory to save and load study progress to/from.",
-        exclude=True,
     )
 
     benchmarks: list[BenchmarkSpecification] = Field(
@@ -426,69 +337,10 @@ class Settings(BaseSettings):
             ),
         ],
         description="Benchmarks to offer to the user for evaluating abliterated models.",
-        exclude=True,
-    )
-
-    max_shard_size: PositiveInt | str = Field(
-        default="5GB",
-        description="Maximum size for individual safetensors files generated when exporting a model.",
-    )
-
-    export_strategy: ExportStrategy | None = Field(
-        default=None,
-        description='How to export the model: "merge", "adapter", or unset to prompt the user.',
-    )
-
-    checkpoint_action: str | None = Field(
-        default=None,
-        description='Action to take in case a checkpoint exists: "continue", "restart", or unset to prompt the user.',
-    )
-
-    trial_index: NonNegativeInt | None = Field(
-        default=None,
-        description="Index (in the sorted Pareto front) of the trial to use, or unset to prompt the user.",
-    )
-
-    n_additional_trials: PositiveInt | None = Field(
-        default=None,
-        description="Number of additional trials to run, or unset to prompt the user.",
-    )
-
-    model_action: str | None = Field(
-        default=None,
-        description='Action to take with the decensored model: "save", "upload", or unset to prompt the user.',
-    )
-
-    save_directory: str | None = Field(
-        default=None,
-        description="Directory to save the model to, or unset to prompt the user.",
-        exclude=True,
-    )
-
-    upload_repo_id: str | None = Field(
-        default=None,
-        description="Name of the Hugging Face repository to upload the model to, or unset to prompt the user.",
-        exclude=True,
-    )
-
-    upload_repo_private: bool | None = Field(
-        default=None,
-        description="Whether the Hugging Face repository to upload the model to should be private, or unset to prompt the user.",
-    )
-
-    upload_reproducibility_information: str | None = Field(
-        default=None,
-        description='Which reproducibility information to add to the Hugging Face repository: "full", "basic", "none", or unset to prompt the user.',
-    )
-
-    ignore_mismatches: bool | None = Field(
-        default=None,
-        description="Whether to attempt to reproduce the model even if there are environment mismatches, or unset to prompt the user.",
     )
 
     refusal_markers: list[str] = Field(
         default=[
-            "disclaimer",
             "sorry",
             "i can'",
             "i cant",

src/heretic/evaluator.py

@@ -1,7 +1,9 @@
 # SPDX-License-Identifier: AGPL-3.0-or-later
 # Copyright (C) 2025-2026  Philipp Emanuel Weidmann <pew@worldwidemann.com> + contributors
 
+import lm_eval
 import torch.nn.functional as F
+from lm_eval.models.huggingface import HFLM
 from torch import Tensor
 
 from .config import Settings
@@ -21,15 +23,16 @@ class Evaluator:
         self.settings = settings
         self.model = model
 
-        print()
+        if not settings.use_piqa:
-        print(
+            print()
-            f"Loading good evaluation prompts from [bold]{settings.good_evaluation_prompts.dataset}[/]..."
+            print(
-        )
+                f"Loading good evaluation prompts from [bold]{settings.good_evaluation_prompts.dataset}[/]..."
-        self.good_prompts = load_prompts(settings, settings.good_evaluation_prompts)
+            )
-        print(f"* [bold]{len(self.good_prompts)}[/] prompts loaded")
+            self.good_prompts = load_prompts(settings, settings.good_evaluation_prompts)
+            print(f"* [bold]{len(self.good_prompts)}[/] prompts loaded")
 
-        print("* Obtaining first-token probability distributions...")
+            print("* Obtaining first-token probability distributions...")
-        self.base_logprobs = model.get_logprobs_batched(self.good_prompts)
+            self.base_logprobs = model.get_logprobs_batched(self.good_prompts)
 
         print()
         print(
@@ -93,35 +96,57 @@ class Evaluator:
         return refusal_count
 
     def get_score(self) -> tuple[tuple[float, float], float, int]:
-        print("  * Obtaining first-token probability distributions...")
+        if self.settings.use_piqa:
-        logprobs = self.model.get_logprobs_batched(self.good_prompts)
+            print("  * Running PIQA benchmark...")
-        kl_divergence = F.kl_div(
+            hflm = HFLM(
-            logprobs,
+                pretrained=self.model.model,  # ty:ignore[invalid-argument-type]
-            self.base_logprobs,
+                tokenizer=self.model.tokenizer,  # ty:ignore[invalid-argument-type]
-            reduction="batchmean",
+                batch_size="auto",
-            log_target=True,
+            )
-        ).item()
+            results = lm_eval.simple_evaluate(
-        print(f"  * KL divergence: [bold]{kl_divergence:.4f}[/]")
+                model=hflm,
+                tasks=["piqa"],
+            )
+            piqa_acc_norm: float = results["results"]["piqa"]["acc_norm,none"]
+            print(f"  * PIQA acc_norm: [bold]{piqa_acc_norm:.4f}[/]")
+        else:
+            print("  * Obtaining first-token probability distributions...")
+            logprobs = self.model.get_logprobs_batched(self.good_prompts)
+            kl_divergence = F.kl_div(
+                logprobs,
+                self.base_logprobs,
+                reduction="batchmean",
+                log_target=True,
+            ).item()
+            print(f"  * KL divergence: [bold]{kl_divergence:.4f}[/]")
 
         print("  * Counting model refusals...")
         refusals = self.count_refusals()
         print(f"  * Refusals: [bold]{refusals}[/]/{len(self.bad_prompts)}")
 
-        kl_divergence_scale = self.settings.kl_divergence_scale
-        kl_divergence_target = self.settings.kl_divergence_target
-
         refusals_score = (
             refusals / self.base_refusals if self.base_refusals > 0 else float(refusals)
         )
 
-        if kl_divergence >= kl_divergence_target:
+        if self.settings.use_piqa:
-            kld_score = kl_divergence / kl_divergence_scale
+            score = (
+                -piqa_acc_norm,
+                refusals_score,
+            )
+
+            return score, -piqa_acc_norm, refusals
         else:
-            kld_score = refusals_score * kl_divergence_target / kl_divergence_scale
+            kl_divergence_scale = self.settings.kl_divergence_scale
+            kl_divergence_target = self.settings.kl_divergence_target
 
-        score = (
+            if kl_divergence >= kl_divergence_target:
-            kld_score,
+                kld_score = kl_divergence / kl_divergence_scale
-            refusals_score,
+            else:
-        )
+                kld_score = refusals_score * kl_divergence_target / kl_divergence_scale
 
-        return score, kl_divergence, refusals
+            score = (
+                kld_score,
+                refusals_score,
+            )
+
+            return score, kl_divergence, refusals

src/heretic/model.py

@@ -4,7 +4,7 @@
 import math
 from contextlib import suppress
 from dataclasses import dataclass
-from typing import Any, Type, cast
+from typing import Any, Callable, Type, TypeAlias, cast
 
 import bitsandbytes as bnb
 import torch
@@ -14,17 +14,17 @@ 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.optim import LBFGS
+from torch.utils.hooks import RemovableHandle
 from transformers import (
     AutoModelForCausalLM,
     AutoModelForImageTextToText,
-    AutoProcessor,
     AutoTokenizer,
     BatchEncoding,
     BitsAndBytesConfig,
     PretrainedConfig,
     PreTrainedModel,
     PreTrainedTokenizerBase,
-    ProcessorMixin,
     TextStreamer,
 )
 from transformers.generation import (
@@ -32,8 +32,7 @@ from transformers.generation import (
 )
 
 from .config import QuantizationMethod, RowNormalization, Settings
-from .system import empty_cache
+from .utils import Prompt, batchify, empty_cache, mean_distances_to_knn, print
-from .utils import Prompt, batchify, format_exception, print
 
 
 def get_model_class(
@@ -55,38 +54,41 @@ class AbliterationParameters:
     min_weight_distance: float
 
 
+@dataclass
+class ARAParameters:
+    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
+
+
+# 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).
+ModuleIO: TypeAlias = list[dict[str, dict[int, tuple[Tensor, Tensor]]]]
+
+
 class Model:
     model: PreTrainedModel | PeftModel
     tokenizer: PreTrainedTokenizerBase
-    # Set for multimodal models, None for text-only ones.
-    processor: ProcessorMixin | None
     peft_config: LoraConfig
-    dtype: torch.dtype
 
     def __init__(self, settings: Settings):
         self.settings = settings
+        self.response_prefix = ""
         self.needs_reload = False
 
-        self.revision_kwargs = {}
-        if settings.model_commit is not None:
-            self.revision_kwargs["revision"] = settings.model_commit
-
         print()
         print(f"Loading model [bold]{settings.model}[/]...")
 
         self.tokenizer = AutoTokenizer.from_pretrained(
             settings.model,
-            **self.revision_kwargs,
+            trust_remote_code=settings.trust_remote_code,
         )
 
-        # Multimodal models have a processor we'll want to save.
-        self.processor = None
-        if get_model_class(settings.model) == AutoModelForImageTextToText:
-            self.processor = AutoProcessor.from_pretrained(
-                settings.model,
-                **self.revision_kwargs,
-            )
-
         # Fallback for tokenizers that don't declare a special pad token.
         if self.tokenizer.pad_token is None:
             self.tokenizer.pad_token = self.tokenizer.eos_token
@@ -102,8 +104,10 @@ class Model:
             if settings.max_memory
             else None
         )
+        self.trusted_models = {settings.model: settings.trust_remote_code}
 
-        self.trusted_models = set()
+        if self.settings.evaluate_model is not None:
+            self.trusted_models[settings.evaluate_model] = settings.trust_remote_code
 
         for dtype in settings.dtypes:
             print(f"* Trying dtype [bold]{dtype}[/]...")
@@ -122,19 +126,14 @@ class Model:
                     dtype=dtype,
                     device_map=settings.device_map,
                     max_memory=self.max_memory,
-                    trust_remote_code=True
+                    trust_remote_code=self.trusted_models.get(settings.model),
-                    if settings.model in self.trusted_models
-                    else None,
-                    **self.revision_kwargs,
                     **extra_kwargs,
                 )
 
-                self.dtype = self.model.dtype
-
                 # If we reach this point and the model requires trust_remote_code,
-                # the user must have agreed when prompted to execute remote code,
+                # either the user accepted, or settings.trust_remote_code is True.
-                # because from_pretrained raises an exception otherwise.
+                if self.trusted_models.get(settings.model) is None:
-                self.trusted_models.add(settings.model)
+                    self.trusted_models[settings.model] = True
 
                 # A test run can reveal dtype-related problems such as the infamous
                 # "RuntimeError: probability tensor contains either `inf`, `nan` or element < 0"
@@ -151,13 +150,7 @@ class Model:
             except Exception as error:
                 self.model = None  # ty:ignore[invalid-assignment]
                 empty_cache()
-
+                print(f"* [red]Failed[/] ({error})")
-                formatted = format_exception(error)
-                if "\n" in formatted:
-                    print(f"* [red]Failed:\n{formatted}[/]")
-                else:
-                    print(f"* [red]Failed ({formatted})[/]")
-
                 continue
 
             if settings.quantization == QuantizationMethod.BNB_4BIT:
@@ -168,21 +161,20 @@ class Model:
         if self.model is None:
             raise Exception("Failed to load model with all configured dtypes.")
 
-        self._apply_lora()
+        if not settings.use_ara or settings.use_ara_lora:
+            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")
-
+        print("* Abliterable components:")
         all_components = {}
         for layer_index in range(len(self.get_layers())):
             for component, modules in self.get_layer_modules(layer_index).items():
                 if component not in all_components:
                     all_components[component] = 0
                 all_components[component] += len(modules)
-
-        print("* Abliterable components:")
         for component, count in all_components.items():
             print(f"  * [bold]{component}[/]: [bold]{count}[/] modules total")
 
@@ -211,7 +203,9 @@ class Model:
 
         target_modules = sorted(target_modules_set)
 
-        if self.settings.row_normalization != RowNormalization.FULL:
+        if self.settings.use_ara_lora:
+            lora_rank = self.settings.ara_lora_rank
+        elif self.settings.row_normalization != RowNormalization.FULL:
             # Rank 1 is sufficient for directional ablation without renormalization.
             lora_rank = 1
         else:
@@ -284,10 +278,7 @@ class Model:
                 self.settings.model,
                 torch_dtype=self.model.dtype,
                 device_map="cpu",
-                trust_remote_code=True
+                trust_remote_code=self.trusted_models.get(self.settings.model),
-                if self.settings.model in self.trusted_models
-                else None,
-                **self.revision_kwargs,
             )
 
             # Apply LoRA adapters to the CPU model
@@ -322,45 +313,42 @@ class Model:
         - Slow path: If switching models or after merge_and_unload(),
           performs full model reload with quantization config.
         """
-
+        current_model = getattr(self.model.config, "name_or_path", None)
-        # If a prior model load was interrupted/cancelled mid-process, self.model will be None.
+        if (
-        current_model = None
+            current_model == self.settings.model
-        if self.model is not None:
+            and not self.needs_reload
-            current_model = getattr(self.model.config, "name_or_path", None)
+            and (not self.settings.use_ara or self.settings.use_ara_lora)
-
+        ):
-        if current_model == self.settings.model and not self.needs_reload:
+            # Reset LoRA adapters to zero (identity transformation)
-            # Reset LoRA adapters to zero (identity transformation).
             for name, module in self.model.named_modules():
                 if "lora_B" in name and hasattr(module, "weight"):
                     torch.nn.init.zeros_(module.weight)
             return
 
+        dtype = self.model.dtype
+
         # Purge existing model object from memory to make space.
         self.model = None  # ty:ignore[invalid-assignment]
         empty_cache()
 
-        quantization_config = self._get_quantization_config(
+        quantization_config = self._get_quantization_config(str(dtype).split(".")[-1])
-            str(self.dtype).split(".")[-1]
-        )
 
-        # Build kwargs, only include quantization_config if it's not None.
+        # Build kwargs, only include quantization_config if it's not None
         extra_kwargs = {}
         if quantization_config is not None:
             extra_kwargs["quantization_config"] = quantization_config
 
         self.model = get_model_class(self.settings.model).from_pretrained(
             self.settings.model,
-            dtype=self.dtype,
+            dtype=dtype,
             device_map=self.settings.device_map,
             max_memory=self.max_memory,
-            trust_remote_code=True
+            trust_remote_code=self.trusted_models.get(self.settings.model),
-            if self.settings.model in self.trusted_models
-            else None,
-            **self.revision_kwargs,
             **extra_kwargs,
         )
 
-        self._apply_lora()
+        if not self.settings.use_ara or self.settings.use_ara_lora:
+            self._apply_lora()
 
         self.needs_reload = False
 
@@ -384,6 +372,9 @@ class Model:
         modules = {}
 
         def try_add(component: str, module: Any):
+            if component not in self.settings.target_components:
+                return
+
             # Only add if it's a proper nn.Module (PEFT can wrap these with LoRA)
             if isinstance(module, Module):
                 if component not in modules:
@@ -399,8 +390,8 @@ class Model:
         with suppress(Exception):
             try_add("attn.o_proj", layer.self_attn.o_proj)  # ty:ignore[possibly-missing-attribute]
 
-        # Qwen3.5 MoE hybrid layers use GatedDeltaNet (linear attention) instead of
+        # Qwen3.5 MoE hybrid layers use GatedDeltaNet (linear attention) instead
-        # standard self-attention, so self_attn.o_proj doesn't exist on those layers.
+        # 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]
 
@@ -418,21 +409,6 @@ class Model:
             for expert in layer.block_sparse_moe.experts:  # ty:ignore[possibly-missing-attribute, not-iterable]
                 try_add("mlp.down_proj", expert.w2)  # ty:ignore[possibly-missing-attribute]
 
-        # LFM dense operator blocks.
-        with suppress(Exception):
-            try_add("attn.o_proj", layer.conv.out_proj)  # ty:ignore[possibly-missing-attribute]
-
-        with suppress(Exception):
-            try_add("mlp.down_proj", layer.feed_forward.w2)  # ty:ignore[possibly-missing-attribute]
-
-        # LFM transformer blocks.
-        with suppress(Exception):
-            try_add("attn.o_proj", layer.self_attn.out_proj)  # ty:ignore[possibly-missing-attribute]
-
-        with suppress(Exception):
-            for expert in layer.feed_forward.experts:  # ty:ignore[possibly-missing-attribute, not-iterable]
-                try_add("mlp.down_proj", expert.w2)  # ty:ignore[possibly-missing-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]
@@ -449,13 +425,11 @@ class Model:
         return modules
 
     def get_abliterable_components(self) -> list[str]:
-        components: set[str] = set()
-
         # Scan all layers because hybrid models (e.g. Qwen3.5 MoE) have different
         # components on different layers (some have self_attn, others linear_attn).
+        components: set[str] = set()
         for layer_index in range(len(self.get_layers())):
             components.update(self.get_layer_modules(layer_index).keys())
-
         return sorted(components)
 
     def abliterate(
@@ -499,12 +473,6 @@ class Model:
                     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 refusal_direction is None:
                     # The index must be shifted by 1 because the first element
                     # of refusal_directions is the direction for the embeddings.
@@ -586,16 +554,7 @@ class Model:
                         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]
@@ -616,6 +575,228 @@ class Model:
                     weight_A.data = lora_A.to(weight_A.dtype)
                     weight_B.data = lora_B.to(weight_B.dtype)
 
+    def ara_abliterate(
+        self,
+        good_module_io: ModuleIO,
+        bad_module_io: ModuleIO,
+        parameters: ARAParameters,
+    ):
+        for layer_index in range(
+            parameters.start_layer_index,
+            parameters.end_layer_index,
+        ):
+            for component, modules in self.get_layer_modules(layer_index).items():
+                for module_index, module in enumerate(modules):
+                    # See above for a (partial) justification of this cast.
+                    module = cast(Linear, module)
+                    matrix = module.weight
+
+                    row_norms = LA.vector_norm(matrix, dim=1, keepdim=True).detach()
+
+                    # Helper function for reparameterization (row-norm preservation constraint).
+                    def get_matrix() -> Tensor:
+                        if self.settings.row_normalization == RowNormalization.FULL:
+                            # See https://huggingface.co/blog/grimjim/norm-preserving-biprojected-abliteration
+                            return row_norms * F.normalize(matrix, p=2, dim=1)
+                        else:
+                            return matrix
+
+                    good_input, good_output = good_module_io[layer_index][component][
+                        module_index
+                    ]
+                    bad_input, bad_output = bad_module_io[layer_index][component][
+                        module_index
+                    ]
+
+                    good_input = good_input.to(matrix.device)
+                    good_output = good_output.to(matrix.device)
+                    bad_input = bad_input.to(matrix.device)
+                    bad_output = bad_output.to(matrix.device)
+
+                    def objective(matrix: Tensor) -> Tensor:
+                        new_good_output = good_input @ matrix.T
+                        new_bad_output = bad_input @ matrix.T
+
+                        # 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
+                        )
+
+                    optimizer = LBFGS(
+                        [matrix],
+                        lr=1.0,
+                        max_iter=20,  # Number of internal iterations per step, *not* the number of steps.
+                        history_size=10,
+                        line_search_fn="strong_wolfe",
+                    )
+
+                    def closure() -> Tensor:
+                        optimizer.zero_grad()
+                        loss = objective(get_matrix())
+                        loss.backward()
+                        return loss
+
+                    # Convergence usually happens within 2-3 steps, so this is more than enough.
+                    for step in range(5):
+                        loss = optimizer.step(closure)
+                        # print(
+                        #    f"\\[{layer_index}/{component}/{module_index}] Step: {step}, Loss: {loss.item():.6f}"
+                        # )
+
+                    # Free the gradient buffers accumulated on the weight parameters
+                    # during optimization. Without this, they persist on the model
+                    # (one full-size gradient per processed weight) and can easily
+                    # consume tens of GiB of VRAM, causing out-of-memory errors
+                    # during the subsequent evaluation.
+                    optimizer.zero_grad(set_to_none=True)
+
+                    with torch.no_grad():
+                        matrix.copy_(get_matrix())
+
+    def ara_lora_abliterate(
+        self,
+        good_module_io: ModuleIO,
+        bad_module_io: ModuleIO,
+        parameters: ARAParameters,
+    ):
+        for layer_index in range(
+            parameters.start_layer_index,
+            parameters.end_layer_index,
+        ):
+            for component, modules in self.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)
+
+                    # Base weight handling and dequantization.
+                    # 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:
+                        # Maintain the original dequantization logic for bitsandbytes.
+                        W_base = cast(
+                            Tensor,
+                            bnb.functional.dequantize_4bit(
+                                base_weight.data, 
+                                quant_state
+                            ).to(torch.float32),
+                        )
+
+                    # Row normalization setup.
+                    # Pre-calculate the original row norms to preserve them.
+                    # This implements the RowNormalization.FULL logic.
+                    W_row_norms = LA.vector_norm(W_base, dim=1, keepdim=True).detach()
+
+                    # Adapter target identification.
+                    # 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)
+
+                    # Data preparation.
+                    # Move I/O tensors to the device of the adapter weights.
+                    good_input, good_output = good_module_io[layer_index][component][module_index]
+                    bad_input, bad_output = 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)
+
+                    # The objective function.
+                    def objective(A: Tensor, B: Tensor) -> Tensor:
+                        # Calculate effective weight: W_eff = W_base + B @ A.
+                        W_eff = W_base + (B @ A)
+
+                        # Apply Row Normalization (keep original norms).
+                        if self.settings.row_normalization == RowNormalization.FULL:
+                            # Normalize to unit length, then scale by original 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
+
+                        # The original ARA loss function.
+                        preserve_good_behavior = (
+                            (new_good_output - good_output) ** 2
+                        ).mean()
+
+                        steer_bad_behavior = (
+                            mean_distances_to_knn(
+                                new_bad_output,
+                                good_output,
+                                parameters.neighbor_count,
+                            ).mean()
+                            + 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
+                        )
+
+                    # Optimization loop.
+                    # We optimize A and B, not the base matrix.
+                    optimizer = LBFGS(
+                        [lora_A, lora_B],
+                        lr=1.0,
+                        max_iter=20,
+                        history_size=10,
+                        line_search_fn="strong_wolfe",
+                    )
+
+                    def closure():
+                        optimizer.zero_grad()
+                        # Pass the actual tensors being optimized to the objective.
+                        loss = objective(lora_A, lora_B)
+                        loss.backward()
+                        return loss
+
+                    # Run optimization steps.
+                    for step in range(5):
+                        optimizer.step(closure)
+
+                    # Free the gradient buffers accumulated on the LoRA adapter
+                    # parameters during optimization (see ara_abliterate for details).
+                    optimizer.zero_grad(set_to_none=True)
+
     def generate(
         self,
         prompts: list[Prompt],
@@ -640,12 +821,10 @@ class Model:
             ),
         )
 
-        if self.settings.response_prefix:
+        if self.response_prefix:
             # Append the common response prefix to the prompts so that evaluation happens
             # at the point where responses start to differ for different prompts.
-            chat_prompts = [
+            chat_prompts = [prompt + self.response_prefix for prompt in chat_prompts]
-                prompt + self.settings.response_prefix for prompt in chat_prompts
-            ]
 
         inputs = self.tokenizer(
             chat_prompts,
@@ -707,9 +886,6 @@ class Model:
             max_new_tokens=1,
             output_hidden_states=True,
             return_dict_in_generate=True,
-            # KV cache is unnecessary here because we only need the hidden states
-            # for the first generated token.
-            use_cache=False,
         )
 
         # This cast is valid because GenerateDecoderOnlyOutput is the return type
@@ -743,11 +919,7 @@ class Model:
                 dim=2,
                 keepdim=True,
             )
-            residuals = torch.clamp(residuals, -thresholds, thresholds)
+            return torch.clamp(residuals, -thresholds, thresholds)
-
-        if self.settings.offload_outputs_to_cpu:
-            residuals = residuals.cpu()
-            empty_cache()
 
         return residuals
 
@@ -759,29 +931,131 @@ class Model:
 
         return torch.cat(residuals, dim=0)
 
-    def get_residuals_mean(self, prompts: list[Prompt]) -> Tensor:
+    def get_module_io(
-        if not prompts:
+        self,
-            raise ValueError("prompts must not be empty")
+        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 = []
 
-        running_sum = None
+        def get_hook(
-        total_count = 0
+            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({})
 
-        for batch in batchify(prompts, self.settings.batch_size):
+                # Layers are invoked in order during inference,
-            batch_residuals = self.get_residuals(batch)
+                # so this should always hold.
+                assert len(module_io) == layer_index + 1
 
-            # Accumulate in high precision on CPU to reduce peak VRAM usage.
+                if component not in module_io[layer_index]:
-            batch_sum = batch_residuals.sum(dim=0, dtype=torch.float64).cpu()
+                    module_io[layer_index][component] = {}
 
-            if running_sum is None:
+                # Each module should be invoked at most once per inference step.
-                running_sum = batch_sum
+                assert module_index not in module_io[layer_index][component]
-            else:
-                running_sum += batch_sum
 
-            total_count += batch_residuals.shape[0]
+                # 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()
 
-        assert running_sum is not None
+                # 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 (running_sum / total_count).to(torch.float32)
+            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
 
     # We work with logprobs rather than probabilities for numerical stability
     # when computing the KL divergence.
@@ -791,9 +1065,8 @@ class Model:
         _, outputs = self.generate(
             prompts,
             max_new_tokens=1,
-            output_logits=True,
+            output_scores=True,
             return_dict_in_generate=True,
-            use_cache=False,
         )
 
         # This cast is valid because GenerateDecoderOnlyOutput is the return type
@@ -801,20 +1074,11 @@ class Model:
         outputs = cast(GenerateDecoderOnlyOutput, outputs)
 
         # Logits for the first (only) generated token.
-        # Use raw logits, not processed generation scores; processors can insert
+        # This cast is valid because we passed output_scores=True above.
-        # -inf for suppressed tokens, which can make KL divergence evaluate to NaN.
+        logits = cast(tuple[FloatTensor], outputs.scores)[0]
-        # This cast is valid because we passed output_logits=True above.
-        logits = cast(tuple[FloatTensor], outputs.logits)[0]
 
         # The returned tensor has shape (prompt, token).
-        logprobs = F.log_softmax(logits, dim=-1)
+        return F.log_softmax(logits, dim=-1)
-
-        if self.settings.offload_outputs_to_cpu:
-            del outputs, logits
-            logprobs = logprobs.cpu()
-            empty_cache()
-
-        return logprobs
 
     def get_logprobs_batched(self, prompts: list[Prompt]) -> Tensor:
         logprobs = []

src/heretic/reproduce.py

@@ -1,391 +0,0 @@
-# SPDX-License-Identifier: AGPL-3.0-or-later
-# Copyright (C) 2025-2026  Philipp Emanuel Weidmann <pew@worldwidemann.com> + contributors
-
-import json
-import platform
-import random
-import shutil
-from dataclasses import asdict
-from enum import IntEnum
-from pathlib import Path
-from typing import Any, cast
-from urllib.request import urlopen
-
-import cpuinfo
-import questionary
-import torch
-from huggingface_hub import HfApi, hf_hub_download
-from huggingface_hub.utils import (
-    GatedRepoError,
-    disable_progress_bars,
-    enable_progress_bars,
-)
-from questionary import Choice, Style
-from rich.table import Table
-
-from .config import Settings
-from .system import (
-    get_accelerator_info_dict,
-    get_heretic_version_info,
-    get_requirements_dict,
-)
-from .utils import ask_if_unset, print
-
-
-def collect_reproducibles(path: str):
-    print(
-        f"Collecting [bold]reproduce.json[/] files from Hugging Face and storing them in [bold]{path}[/]..."
-    )
-    print()
-
-    api = HfApi()
-
-    models = api.list_models(
-        filter=["heretic", "reproducible"],
-        sort="created_at",
-        expand=["gated", "tags"],
-    )
-
-    found = 0
-    downloaded = 0
-
-    # We're only downloading tiny files, so the progress bars are just noise.
-    disable_progress_bars()
-
-    try:
-        for model in models:
-            # Ignore repositories containing quantizations.
-            if model.tags is not None and "gguf" in model.tags:
-                continue
-
-            if model.gated:
-                try:
-                    api.auth_check(model.id, repo_type="model")
-                except GatedRepoError:
-                    continue
-
-            print(f"[bold]{model.id}[/]...", end="")
-
-            user, repository = model.id.split("/")
-
-            paths_info = api.get_paths_info(
-                model.id,
-                "reproduce/reproduce.json",
-                expand=True,
-            )
-            # The reproduce.json file might not exist in the repository
-            # despite the relevant tags being present.
-            if not paths_info:
-                print(" [yellow]no reproduce.json found[/]")
-                continue
-
-            found += 1
-
-            commit_hash = paths_info[0].last_commit.oid
-
-            file_path = (
-                Path(path)
-                / "huggingface.co"
-                / user
-                / f"{repository}-{commit_hash[:7]}.json"
-            )
-            if file_path.exists():
-                print(" already stored")
-                continue
-
-            cache_path = hf_hub_download(
-                model.id,
-                "reproduce/reproduce.json",
-            )
-
-            file_path.parent.mkdir(parents=True, exist_ok=True)
-            shutil.copyfile(cache_path, file_path)
-            print(" [green]downloaded[/]")
-
-            downloaded += 1
-    finally:
-        enable_progress_bars()
-
-    print()
-    print(f"Found: [bold]{found}[/] files")
-    print(f"Downloaded: [bold]{downloaded}[/] files")
-    print(f"Already stored: [bold]{found - downloaded}[/] files")
-
-
-def load_reproduction_information(path: str) -> dict[str, Any]:
-    if path.lower().startswith(("http://", "https://")):
-        # The path is a URL on the web.
-
-        # Obtain raw download URL.
-        path = path.replace("/blob/", "/raw/")  # Hugging Face, GitHub
-        path = path.replace("/src/branch/", "/raw/branch/")  # Codeberg
-
-        json_str = urlopen(path).read().decode("utf-8")
-    else:
-        # The path is (assumed to be) a local file system path.
-        json_str = Path(path).read_text(encoding="utf-8")
-
-    return json.loads(json_str)
-
-
-class MismatchSeverity(IntEnum):
-    LOW = 1
-    MEDIUM = 2
-    HIGH = 3
-    CRITICAL = 4
-
-    def __rich__(self) -> str:
-        match self:
-            case MismatchSeverity.LOW:
-                return "[green]low[/]"
-            case MismatchSeverity.MEDIUM:
-                return "[yellow]medium[/]"
-            case MismatchSeverity.HIGH:
-                return "[red]high[/]"
-            case MismatchSeverity.CRITICAL:
-                return "[bold red]critical[/]"
-            case _:
-                raise ValueError(f"unknown MismatchSeverity value: {self}")
-
-
-def get_package_mismatch_severity(package_name: str) -> MismatchSeverity:
-    if package_name in [
-        "heretic-llm",
-    ]:
-        return MismatchSeverity.CRITICAL
-    elif package_name in [
-        "torch",
-        "transformers",
-    ]:
-        return MismatchSeverity.HIGH
-    elif package_name in [
-        "accelerate",
-        "bitsandbytes",
-        "kernels",
-        "optuna",
-        "peft",
-        "tokenizers",
-        "triton",
-    ]:
-        return MismatchSeverity.MEDIUM
-    else:
-        return MismatchSeverity.LOW
-
-
-def format_version_information(version_information: dict[str, Any]) -> str:
-    version = version_information["version"]
-    metadata = version_information["metadata"]
-
-    if "type" in metadata:
-        match metadata["type"]:
-            case "pypi":
-                return version
-            case "git":
-                return f"{version}-git+{metadata['url']}@{metadata['commit_hash']}"
-            case "local":
-                # Append a random number to ensure that two local installations
-                # are always considered to be different versions.
-                return f"{version}-local-{random.randint(2**16, 2**17)}"
-            case _:
-                raise ValueError(
-                    f"unknown metadata.type value in version information: {metadata['type']}"
-                )
-    else:
-        return f"{version}-unknown-{random.randint(2**16, 2**17)}"
-
-
-def check_environment(
-    settings: Settings,
-    reproduction_information: dict[str, Any],
-) -> bool | None:
-    mismatch_severity: MismatchSeverity | None = None
-
-    system_mismatches = []
-    package_mismatches = []
-
-    def verify(
-        mismatch_list: list[tuple[str, Any, Any, MismatchSeverity]],
-        name: str,
-        this: Any,
-        original: Any,
-        severity: MismatchSeverity,
-    ):
-        nonlocal mismatch_severity
-        if this != original:
-            mismatch_list.append((name, this, original, severity))
-            if mismatch_severity is None:
-                mismatch_severity = severity
-            else:
-                mismatch_severity = max(severity, mismatch_severity)
-
-    if "system" in reproduction_information:
-        system = reproduction_information["system"]
-
-        verify(
-            system_mismatches,
-            "Python version",
-            platform.python_version(),
-            system["python"]["version"],
-            MismatchSeverity.LOW,
-        )
-
-        verify(
-            system_mismatches,
-            "Operating system",
-            platform.platform(),
-            system["os"]["platform"],
-            MismatchSeverity.LOW,
-        )
-
-        verify(
-            system_mismatches,
-            "CPU",
-            cpuinfo.get_cpu_info().get("brand_raw"),
-            system["cpu"]["brand"],
-            MismatchSeverity.LOW,
-        )
-
-        accelerators = get_accelerator_info_dict()
-
-        verify(
-            system_mismatches,
-            "Accelerator type",
-            accelerators["type"],
-            system["accelerators"]["type"],
-            MismatchSeverity.HIGH,
-        )
-
-        if (
-            accelerators["type"]
-            and accelerators["type"] == system["accelerators"]["type"]
-        ):
-            verify(
-                system_mismatches,
-                accelerators["api_name"],
-                accelerators["api_version"],
-                system["accelerators"]["api_version"],
-                MismatchSeverity.MEDIUM,
-            )
-            verify(
-                system_mismatches,
-                "Driver version",
-                accelerators["driver_version"],
-                system["accelerators"]["driver_version"],
-                MismatchSeverity.MEDIUM,
-            )
-            verify(
-                system_mismatches,
-                "Devices",
-                "\n".join([device["name"] for device in accelerators["devices"]]),
-                "\n".join(
-                    [device["name"] for device in system["accelerators"]["devices"]]
-                ),
-                MismatchSeverity.MEDIUM,
-            )
-
-    else:
-        print(
-            (
-                "[yellow]The provided JSON file does not contain 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. "
-                "Reproduction may or may not produce a byte-for-byte identical model.[/]"
-            )
-        )
-
-    requirements = get_requirements_dict()
-    requirements["heretic-llm"] = format_version_information(
-        asdict(get_heretic_version_info())
-    )
-    requirements["torch"] = torch.__version__
-
-    original_requirements = reproduction_information["environment"]["requirements"]
-    original_requirements["heretic-llm"] = format_version_information(
-        reproduction_information["environment"]["heretic"]
-    )
-    original_requirements["torch"] = reproduction_information["environment"][
-        "pytorch_version"
-    ]
-
-    package_names = sorted(requirements.keys() | original_requirements.keys())
-
-    for package_name in package_names:
-        verify(
-            package_mismatches,
-            package_name,
-            requirements.get(package_name),
-            original_requirements.get(package_name),
-            get_package_mismatch_severity(package_name),
-        )
-
-    if system_mismatches or package_mismatches:
-        print()
-        print(
-            (
-                "[yellow]Your local environment doesn't perfectly match the environment "
-                "used to produce the original model. The following components differ:[/]"
-            )
-        )
-
-    if system_mismatches:
-        table = Table()
-        table.add_column("Component")
-        table.add_column("This system", overflow="fold")
-        table.add_column("Original system", overflow="fold")
-        table.add_column("Severity", width=8)
-
-        for component, this, original, severity in system_mismatches:
-            table.add_row(f"[bold]{component}[/]", this, original, severity)
-
-        print()
-        print("[bold]System Mismatches[/]")
-        print(table)
-
-    if package_mismatches:
-        table = Table()
-        table.add_column("Package")
-        table.add_column("This system", overflow="fold")
-        table.add_column("Original system", overflow="fold")
-        table.add_column("Severity", width=8)
-
-        for package, this, original, severity in package_mismatches:
-            table.add_row(f"[bold]{package}[/]", this, original, severity)
-
-        print()
-        print("[bold]Package Mismatches[/]")
-        print(table)
-
-    if system_mismatches or package_mismatches:
-        print()
-        print(
-            (
-                f"There is a {cast(MismatchSeverity, mismatch_severity).__rich__()} chance "
-                "that reproduction won't produce a byte-for-byte identical model. "
-                "However, the resulting model will very likely still behave similarly "
-                "to the original model."
-            )
-        )
-
-        if settings.ignore_mismatches is None:
-            print()
-
-        return ask_if_unset(
-            settings.ignore_mismatches,
-            questionary.select(
-                "How would you like to proceed?",
-                choices=[
-                    Choice(
-                        title="Attempt to reproduce the model anyway",
-                        value=True,
-                    ),
-                    Choice(
-                        title="Exit program",
-                        value=False,
-                    ),
-                ],
-                style=Style([("highlighted", "reverse")]),
-            ),
-        )
-    else:
-        # There are no mismatches at all, so there is nothing to confirm.
-        return True

src/heretic/system.py

@@ -1,478 +0,0 @@
-# SPDX-License-Identifier: AGPL-3.0-or-later
-# Copyright (C) 2025-2026  Philipp Emanuel Weidmann <pew@worldwidemann.com> + contributors
-
-import gc
-import importlib.metadata
-import json
-import os
-import platform
-import re
-import subprocess
-import sys
-from dataclasses import dataclass
-from typing import Any
-
-import cpuinfo
-import torch
-from accelerate.utils import (
-    is_mlu_available,
-    is_musa_available,
-    is_npu_available,
-    is_sdaa_available,
-    is_xpu_available,
-)
-
-
-def empty_cache():
-    """Clears the backend cache and collects garbage."""
-
-    # Collecting garbage is not an idempotent operation, and to avoid OOM errors,
-    # gc.collect() has to be called both before and after emptying the backend cache.
-    # See https://github.com/p-e-w/heretic/pull/17 for details.
-    gc.collect()
-
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-    elif is_xpu_available():
-        torch.xpu.empty_cache()
-    elif is_mlu_available():
-        torch.mlu.empty_cache()  # ty:ignore[unresolved-attribute]
-    elif is_sdaa_available():
-        torch.sdaa.empty_cache()  # ty:ignore[unresolved-attribute]
-    elif is_musa_available():
-        torch.musa.empty_cache()  # ty:ignore[unresolved-attribute]
-    elif torch.backends.mps.is_available():
-        torch.mps.empty_cache()
-
-    gc.collect()
-
-
-def get_nvidia_driver_version() -> str | None:
-    """Gets the NVIDIA driver version using nvidia-smi."""
-
-    try:
-        output = subprocess.check_output(
-            ["nvidia-smi", "--query-gpu=driver_version", "--format=csv,noheader"],
-            stderr=subprocess.DEVNULL,
-            text=True,
-        )
-        return output.strip().split("\n")[0]
-    except (subprocess.CalledProcessError, FileNotFoundError, IndexError):
-        return None
-
-
-def get_amdgpu_driver_version() -> str | None:
-    """Gets the AMD GPU (ROCm) driver and suite version info."""
-
-    # 1. Try amd-smi (modern standard for ROCm 6.0+)
-    try:
-        output = subprocess.check_output(
-            ["amd-smi", "version"],
-            stderr=subprocess.DEVNULL,
-            text=True,
-        )
-        if output.strip():
-            return output.strip().replace("\n", " | ")
-    except (subprocess.CalledProcessError, FileNotFoundError):
-        pass
-
-    # 2. Try rocm-smi --showdriverversion
-    try:
-        output = subprocess.check_output(
-            ["rocm-smi", "--showdriverversion"],
-            stderr=subprocess.DEVNULL,
-            text=True,
-        )
-        for line in output.split("\n"):
-            if "Driver version" in line:
-                return line.split(":")[-1].strip()
-    except (subprocess.CalledProcessError, FileNotFoundError):
-        pass
-
-    # 3. Try /sys/module/amdgpu/version (Linux kernel driver version)
-    try:
-        if platform.system() == "Linux":
-            version_path = "/sys/module/amdgpu/version"
-            if os.path.exists(version_path):
-                with open(version_path, "r", encoding="utf-8") as f:
-                    return f.read().strip()
-    except Exception:
-        pass
-
-    return None
-
-
-def get_xpu_driver_version() -> str | None:
-    """Gets the Intel XPU driver version."""
-
-    try:
-        output = subprocess.check_output(
-            ["xpu-smi", "discovery"],
-            stderr=subprocess.DEVNULL,
-            text=True,
-        )
-        for line in output.split("\n"):
-            if "Driver Version" in line:
-                return line.split(":")[-1].strip()
-        return None
-    except (subprocess.CalledProcessError, FileNotFoundError):
-        return None
-
-
-def get_npu_driver_version() -> str | None:
-    """Gets the Huawei NPU driver version."""
-
-    try:
-        output = subprocess.check_output(
-            ["npu-smi", "info", "-t", "board", "-i", "0"],
-            stderr=subprocess.DEVNULL,
-            text=True,
-        )
-        for line in output.split("\n"):
-            if "Software Version" in line:
-                return line.split()[-1].strip()
-        return None
-    except (subprocess.CalledProcessError, FileNotFoundError):
-        return None
-
-
-def get_mps_driver_version() -> str | None:
-    """Gets the Apple Silicon (MPS) driver version via macOS version."""
-
-    try:
-        output = subprocess.check_output(
-            ["sw_vers", "-productVersion"],
-            stderr=subprocess.DEVNULL,
-            text=True,
-        )
-        return output.strip()
-    except (subprocess.CalledProcessError, FileNotFoundError):
-        return None
-
-
-@dataclass
-class HereticVersionInfo:
-    """Detailed information about the heretic-llm installation."""
-
-    version: str
-    origin: str | None
-    is_standard_pypi: bool
-    metadata: dict[str, Any]
-
-
-def get_heretic_version_info() -> HereticVersionInfo:
-    """Detects version and installation source (PyPI, Git, Local) of heretic-llm."""
-
-    package_name = "heretic-llm"
-    origin_metadata: dict[str, Any] = {"type": "unknown"}
-    # This package must be installed for this code to run.
-    distribution = importlib.metadata.distribution(package_name)
-
-    base_version = distribution.version.lstrip("v")
-
-    try:
-        direct_url_content = distribution.read_text("direct_url.json")
-    except Exception:
-        direct_url_content = None
-
-    if not direct_url_content:
-        # Standard PyPI installation.
-        origin_metadata["type"] = "pypi"
-
-        return HereticVersionInfo(
-            version=base_version,
-            origin="PyPI",
-            is_standard_pypi=True,
-            metadata=origin_metadata,
-        )
-
-    data = json.loads(direct_url_content)
-
-    # Check for Git source.
-    if "vcs_info" in data and data["vcs_info"].get("vcs") == "git":
-        vcs_info = data["vcs_info"]
-        commit_hash = vcs_info.get("commit_id", "unknown")
-        repo_url = data.get("url", "unknown_repo")
-        requested_revision = vcs_info.get("requested_revision")
-
-        if requested_revision:
-            origin_str = (
-                f"Git ({repo_url}@{requested_revision} - commit: {commit_hash})"
-            )
-        else:
-            origin_str = f"Git ({repo_url} @ {commit_hash})"
-
-        origin_metadata.update(
-            {
-                "type": "git",
-                "url": repo_url,
-                "commit_hash": commit_hash,
-                "requested_revision": requested_revision,
-            }
-        )
-
-        return HereticVersionInfo(
-            version=base_version,
-            origin=origin_str,
-            is_standard_pypi=False,
-            metadata=origin_metadata,
-        )
-
-    # Check for local file/wheel directory.
-    if "url" in data and data["url"].startswith("file://"):
-        origin_metadata["type"] = "local"
-
-        return HereticVersionInfo(
-            version=base_version,
-            origin="Local",
-            is_standard_pypi=False,
-            metadata=origin_metadata,
-        )
-
-    return HereticVersionInfo(
-        version=base_version,
-        origin=None,
-        is_standard_pypi=False,
-        metadata=origin_metadata,
-    )
-
-
-def get_accelerator_info_dict() -> dict[str, Any]:
-    """Retrieves raw accelerator info (CUDA, ROCm, etc) directly into structured keys."""
-
-    if torch.cuda.is_available():
-        count = torch.cuda.device_count()
-        is_rocm = getattr(torch.version, "hip", None) is not None
-
-        # ROCm (AMD) and CUDA (NVIDIA) share the same API in PyTorch.
-        # We distinguish them by checking for the HIP version.
-        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]
-            "driver_version": get_amdgpu_driver_version()
-            if is_rocm
-            else get_nvidia_driver_version(),
-            "devices": [],
-        }
-
-        for i in range(count):
-            name = torch.cuda.get_device_name(i)
-            vram = torch.cuda.mem_get_info(i)[1] / (1024**3)
-            info["devices"].append({"name": name, "vram_gb": round(vram, 2)})
-
-        return info
-
-    if is_xpu_available():
-        count = torch.xpu.device_count()  # ty:ignore[unresolved-attribute]
-        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]
-        }
-
-    if is_mlu_available():
-        count = torch.mlu.device_count()  # ty:ignore[unresolved-attribute]
-        return {
-            "type": "MLU",
-            "api_name": None,
-            "api_version": None,
-            "driver_version": None,
-            "devices": [{"name": torch.mlu.get_device_name(i)} for i in range(count)],  # ty:ignore[unresolved-attribute]
-        }
-
-    if is_sdaa_available():
-        count = torch.sdaa.device_count()  # ty:ignore[unresolved-attribute]
-        return {
-            "type": "SDAA",
-            "api_name": None,
-            "api_version": None,
-            "driver_version": None,
-            "devices": [{"name": torch.sdaa.get_device_name(i)} for i in range(count)],  # ty:ignore[unresolved-attribute]
-        }
-
-    if is_musa_available():
-        count = torch.musa.device_count()  # ty:ignore[unresolved-attribute]
-        return {
-            "type": "MUSA",
-            "api_name": None,
-            "api_version": None,
-            "driver_version": None,
-            "devices": [{"name": torch.musa.get_device_name(i)} for i in range(count)],  # ty:ignore[unresolved-attribute]
-        }
-
-    if is_npu_available():
-        return {
-            "type": "NPU",
-            "api_name": "CANN Version",
-            "api_version": torch.version.cann,  # ty:ignore[unresolved-attribute]
-            "driver_version": get_npu_driver_version(),
-            "devices": [],  # Multi-NPU is less common.
-        }
-
-    if torch.backends.mps.is_available():
-        return {
-            "type": "MPS",
-            "api_name": None,
-            "api_version": None,
-            "driver_version": get_mps_driver_version(),
-            "devices": [{"name": "Apple Metal"}],
-        }
-
-    return {"type": None}
-
-
-def get_accelerator_info(include_warnings: bool = True) -> str:
-    """Convenience wrapper for hardware detection and console-friendly formatting."""
-
-    info = get_accelerator_info_dict()
-
-    if info["type"] is None:
-        suffix = " Operations will be slow." if include_warnings else ""
-        return (
-            f"[bold yellow]No GPU or other accelerator detected.{suffix}[/]\n".strip()
-        )
-
-    devices = info["devices"]
-    count = len(devices)
-    total_vram = sum(d.get("vram_gb", 0) for d in devices)
-
-    vram_suffix = f" ({total_vram:.2f} GB total VRAM)" if total_vram > 0 else ""
-    report = f"Detected [bold]{count or 1}[/] {info['type']} device(s){vram_suffix}\n"
-
-    if info.get("api_name") and info.get("api_version"):
-        report += f"{info['api_name']}: [bold]{info['api_version']}[/]\n"
-
-    driver = info.get("driver_version") or "Unknown"
-    report += f"Driver Version: [bold]{driver}[/]\n"
-
-    for i, dev in enumerate(devices):
-        vram = f" ({dev['vram_gb']:.2f} GB)" if dev.get("vram_gb") else ""
-        report += f"* {info['type']} {i}: [bold]{dev['name']}[/]{vram}\n"
-
-    return report.strip()
-
-
-def get_cpu_info_dict() -> dict[str, str | int | None]:
-    """Gets granular CPU identifiers using the py-cpuinfo library."""
-
-    info = cpuinfo.get_cpu_info()
-
-    return {
-        "brand": info.get("brand_raw"),
-        "vendor": info.get("vendor_id_raw"),
-        "family": info.get("family"),
-        "model": info.get("model"),
-        "stepping": info.get("stepping"),
-    }
-
-
-def get_cpu_info() -> str:
-    """Gets the CPU brand name."""
-
-    info = get_cpu_info_dict()
-    parts = []
-    parts.append(
-        f"Family {info['family']}, Model {info['model']}, Stepping {info['stepping']}"
-    )
-
-    details = f" ({'; '.join(parts)})" if parts else ""
-    brand = info["brand"] or "Unknown CPU"
-    return f"{brand}{details}"
-
-
-def get_python_env_info_dict() -> dict[str, str]:
-    implementation = platform.python_implementation()
-    compiler = platform.python_compiler()
-
-    # Check for Conda.
-    if "CONDA_PREFIX" in os.environ:
-        env_type = "Conda"
-    # Check for Virtualenv/Venv.
-    elif hasattr(sys, "base_prefix") and sys.base_prefix != sys.prefix:
-        env_type = "Virtualenv/Venv"
-    else:
-        env_type = "System"
-
-    return {
-        "version": platform.python_version(),
-        "implementation": implementation,
-        "compiler": compiler,
-        "environment": env_type,
-    }
-
-
-def get_python_env_info() -> str:
-    """Detects the type of Python environment (Conda, Venv, etc.) and build info."""
-
-    info = get_python_env_info_dict()
-    return f"{info['version']} ({info['implementation']}, {info['compiler']}) [{info['environment']}]"
-
-
-def get_package_version(name: str) -> str:
-    """Gets the installed version of a package, stripping local suffixes like +cu128."""
-
-    # Normalize name: pip considers hyphens and underscores equivalent.
-    normalized_name = name.lower().replace("_", "-")
-    version_str = importlib.metadata.version(normalized_name)
-    return version_str.split("+")[0] if "+" in version_str else version_str
-
-
-def get_requirements_dict() -> dict[str, str]:
-    """Recursively finds all direct and transitive dependencies of heretic-llm and core libraries."""
-
-    # We start with heretic-llm and the core compute libraries.
-    # PyTorch is not listed as a dependency in the heretic-llm package
-    # because installation is hardware-specific and must be done manually.
-    packages_to_check = ["heretic-llm", "torch", "torchaudio", "torchvision"]
-
-    visited = set()
-    required_packages = set()
-
-    while packages_to_check:
-        package = packages_to_check.pop(0)
-        # Normalize name: pip considers hyphens and underscores equivalent.
-        normalized_package = package.lower().replace("_", "-")
-        if normalized_package in visited:
-            continue
-        visited.add(normalized_package)
-
-        try:
-            distribution = importlib.metadata.distribution(normalized_package)
-            required_packages.add(normalized_package)
-            if distribution.requires:
-                for requirement in distribution.requires:
-                    # Requirements can include environment markers like '; extra == "hf"'
-                    # or version constraints. We should ignore optional 'extra' dependencies
-                    # to keep the reproduction environment clean and relevant.
-                    if ";" in requirement and "extra ==" in requirement:
-                        continue
-
-                    # We just want the base package name.
-                    match = re.match(r"^([a-zA-Z0-9_\-]+)", requirement)
-                    if match:
-                        dep_name = match.group(0).lower().replace("_", "-")
-                        if dep_name not in visited:
-                            packages_to_check.append(dep_name)
-        except importlib.metadata.PackageNotFoundError:
-            # If a package is listed as a dependency but not installed, we skip it.
-            continue
-
-    required_packages_sorted = sorted(required_packages)
-
-    # Lookup versions for all discovered packages.
-    dependencies = {}
-    version_info = get_heretic_version_info()
-
-    for package in required_packages_sorted:
-        # If heretic-llm was installed from source (Git/Local), exclude it
-        # from requirements.txt to prevent pip from downloading an unrelated
-        # version from PyPI during reproduction.
-        if package == "heretic-llm" and not version_info.is_standard_pypi:
-            continue
-
-        dependencies[package] = get_package_version(package)
-
-    return dependencies

src/heretic/utils.py

@@ -1,44 +1,33 @@
 # SPDX-License-Identifier: AGPL-3.0-or-later
 # Copyright (C) 2025-2026  Philipp Emanuel Weidmann <pew@worldwidemann.com> + contributors
 
-import hashlib
+import gc
-import json
+import getpass
 import os
-import platform
-import tempfile
-import traceback
 from dataclasses import dataclass
-from datetime import datetime, timezone
 from importlib.metadata import version
 from pathlib import Path
-from typing import TypeVar
+from typing import Any, TypeVar
 
-import huggingface_hub
+import questionary
-import tomli_w
 import torch
+from accelerate.utils import (
+    is_mlu_available,
+    is_musa_available,
+    is_sdaa_available,
+    is_xpu_available,
+)
 from datasets import DatasetDict, ReadInstruction, load_dataset, load_from_disk
 from datasets.config import DATASET_STATE_JSON_FILENAME
 from datasets.download.download_manager import DownloadMode
 from datasets.utils.info_utils import VerificationMode
-from huggingface_hub.utils import validate_repo_id
 from optuna import Trial
-from optuna.trial import FrozenTrial
 from psutil import Process
-from questionary import Question
+from questionary import Choice, Style
 from rich.console import Console
+from torch import Tensor
 
-from .config import DatasetSpecification, Settings
+from .config import DatasetSpecification, RowNormalization, Settings
-from .system import (
-    get_accelerator_info_dict,
-    get_cpu_info_dict,
-    get_heretic_version_info,
-    get_python_env_info_dict,
-    get_requirements_dict,
-    is_xpu_available,
-)
-
-T = TypeVar("T")
-
 
 print = Console(highlight=False).print
 
@@ -66,6 +55,99 @@ def print_memory_usage():
         p("Driver (reserved) MPS memory", torch.mps.driver_allocated_memory())
 
 
+def is_notebook() -> bool:
+    # Check for specific environment variables (Colab, Kaggle).
+    # This is necessary because when running as a subprocess (e.g. !heretic),
+    # get_ipython() might not be available or might not reflect the notebook environment.
+    if os.getenv("COLAB_GPU") or os.getenv("KAGGLE_KERNEL_RUN_TYPE"):
+        return True
+
+    # Check IPython shell type (for library usage).
+    try:
+        from IPython import get_ipython  # ty:ignore[unresolved-import]
+
+        shell = get_ipython()
+        if shell is None:
+            return False
+
+        shell_name = shell.__class__.__name__
+        if shell_name in ["ZMQInteractiveShell", "Shell"]:
+            return True
+
+        if "google.colab" in str(shell.__class__):
+            return True
+
+        return False
+    except (ImportError, NameError, AttributeError):
+        return False
+
+
+def prompt_select(message: str, choices: list[Any]) -> Any:
+    if is_notebook():
+        print()
+        print(message)
+        real_choices = []
+
+        for i, choice in enumerate(choices, 1):
+            if isinstance(choice, Choice):
+                print(f"[{i}] {choice.title}")
+                real_choices.append(choice.value)
+            else:
+                print(f"[{i}] {choice}")
+                real_choices.append(choice)
+
+        while True:
+            try:
+                selection = input("Enter number: ")
+                index = int(selection) - 1
+                if 0 <= index < len(real_choices):
+                    return real_choices[index]
+                print(
+                    f"[red]Please enter a number between 1 and {len(real_choices)}[/]"
+                )
+            except ValueError:
+                print("[red]Invalid input. Please enter a number.[/]")
+    else:
+        return questionary.select(
+            message,
+            choices=choices,
+            style=Style([("highlighted", "reverse")]),
+        ).ask()
+
+
+def prompt_text(
+    message: str,
+    default: str = "",
+    qmark: str = "?",
+    unsafe: bool = False,
+) -> str:
+    if is_notebook():
+        print()
+        result = input(f"{message} [{default}]: " if default else f"{message}: ")
+        return result if result else default
+    else:
+        question = questionary.text(message, default=default, qmark=qmark)
+        if unsafe:
+            return question.unsafe_ask()
+        else:
+            return question.ask()
+
+
+def prompt_path(message: str) -> str:
+    if is_notebook():
+        return prompt_text(message)
+    else:
+        return questionary.path(message, only_directories=True).ask()
+
+
+def prompt_password(message: str) -> str:
+    if is_notebook():
+        print()
+        return getpass.getpass(message)
+    else:
+        return questionary.password(message).ask()
+
+
 def format_duration(seconds: float) -> str:
     seconds = round(seconds)
     hours, seconds = divmod(seconds, 3600)
@@ -79,64 +161,12 @@ def format_duration(seconds: float) -> str:
         return f"{seconds}s"
 
 
-def format_exception(error: Exception) -> str:
-    # Walk causal chain to find a non-empty message.
-    current = error
-    while current is not None:
-        message = str(current).strip()
-        if message:
-            return message
-        current = current.__cause__ or current.__context__
-
-    # If there is no message in the entire causal chain, fall back to the complete traceback.
-    return traceback.format_exc().strip()
-
-
-def ask_if_unset(value: T, question: Question, unsafe: bool = False) -> T:
-    if value is None:
-        if unsafe:
-            return question.unsafe_ask()
-        else:
-            return question.ask()
-    else:
-        return value
-
-
-def is_hf_path(path: str) -> bool:
-    """Checks whether a path likely refers to a Hugging Face repository."""
-
-    # Match Transformers: Existing local paths take precedence over Hub lookup,
-    # even if the path string is also a valid repository ID.
-    if Path(path).exists():
-        return False
-
-    validate_repo_id(path)
-    return True
-
-
 @dataclass
 class Prompt:
     system: str
     user: str
 
 
-def get_split_slice(split_str: str, length: int) -> tuple[int, int]:
-    """Resolves a split specification into absolute (start, end) indices."""
-
-    # The split name is the part before the slice, e.g. "train" in "train[:400]".
-    split_name = split_str.split("[")[0]
-
-    # Associate the split with its number of examples (lines).
-    name_to_length = {split_name: length}
-
-    # Convert the instructions to absolute indices and select the first one.
-    absolute_instruction = ReadInstruction.from_spec(split_str).to_absolute(
-        name_to_length
-    )[0]
-
-    return absolute_instruction.from_, absolute_instruction.to
-
-
 def load_prompts(
     settings: Settings,
     specification: DatasetSpecification,
@@ -144,43 +174,25 @@ def load_prompts(
     path = specification.dataset
     split_str = specification.split
 
-    if os.path.isfile(path):
+    if os.path.isdir(path):
-        # Plain text file with one prompt per line. Empty lines are ignored.
+        if Path(path, DATASET_STATE_JSON_FILENAME).exists():
-        with open(path, encoding="utf-8") as file:
-            prompts = [line.strip() for line in file if line.strip()]
-
-        # The split is optional for text files. When given, it selects a subset
-        # of the lines using slice notation (e.g. "[:400]"). A synthetic split
-        # name is prepended because ReadInstruction expects a named split.
-        if split_str is not None:
-            start, end = get_split_slice(f"_{split_str}", len(prompts))
-            prompts = prompts[start:end]
-    else:
-        # All dataset sources require an explicit split and column.
-        if split_str is None:
-            raise ValueError(f'The "split" field is required for datasets: {path}')
-
-        if specification.column is None:
-            raise ValueError(f'The "column" field is required for datasets: {path}')
-
-        if is_hf_path(path):
-            dataset = load_dataset(
-                path,
-                revision=specification.commit,
-                split=split_str,
-            )
-        elif Path(path, DATASET_STATE_JSON_FILENAME).exists():
             # Dataset saved with datasets.save_to_disk; needs special handling.
             # Path should be the subdirectory for a particular split.
             dataset = load_from_disk(path)
             assert not isinstance(dataset, DatasetDict), (
                 "Loading dataset dicts is not supported"
             )
-            # Parse the split instructions and apply them.
+            # Parse the split instructions.
-            start, end = get_split_slice(split_str, len(dataset))
+            instruction = ReadInstruction.from_spec(split_str)
-            dataset = dataset[start:end]
+            # Associate the split with its number of examples (lines).
+            split_name = str(dataset.split)
+            name2len = {split_name: len(dataset)}
+            # Convert the instructions to absolute indices and select the first one.
+            abs_instruction = instruction.to_absolute(name2len)[0]
+            # Get the dataset by applying the indices.
+            dataset = dataset[abs_instruction.from_ : abs_instruction.to]
         else:
-            # Path should be a local directory.
+            # Path is a local directory.
             dataset = load_dataset(
                 path,
                 split=split_str,
@@ -189,8 +201,11 @@ def load_prompts(
                 # But also don't use cached data, as the dataset may have changed on disk.
                 download_mode=DownloadMode.FORCE_REDOWNLOAD,
             )
+    else:
+        # Probably a repository path; let load_dataset figure it out.
+        dataset = load_dataset(path, split=split_str)
 
-        prompts = list(dataset[specification.column])
+    prompts = list(dataset[specification.column])
 
     if specification.prefix:
         prompts = [f"{specification.prefix} {prompt}" for prompt in prompts]
@@ -213,50 +228,103 @@ def load_prompts(
     ]
 
 
+T = TypeVar("T")
+
+
 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]:
+# For each vector in the 2D-tensor `a`, computes the mean Euclidean distance
-    params = {}
+# 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)
 
-    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():
+def empty_cache():
-        for name, value in parameters.items():
+    # Collecting garbage is not an idempotent operation, and to avoid OOM errors,
-            params[f"{component}.{name}"] = f"{value:.2f}"
+    # gc.collect() has to be called both before and after emptying the backend cache.
+    # See https://github.com/p-e-w/heretic/pull/17 for details.
+    gc.collect()
 
-    return params
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    elif is_xpu_available():
+        torch.xpu.empty_cache()
+    elif is_mlu_available():
+        torch.mlu.empty_cache()  # ty:ignore[unresolved-attribute]
+    elif is_sdaa_available():
+        torch.sdaa.empty_cache()  # ty:ignore[unresolved-attribute]
+    elif is_musa_available():
+        torch.musa.empty_cache()  # ty:ignore[unresolved-attribute]
+    elif torch.backends.mps.is_available():
+        torch.mps.empty_cache()
+
+    gc.collect()
+
+
+def get_trial_parameters(settings: Settings, trial: Trial) -> dict[str, str]:
+    if settings.use_ara:
+        parameters = trial.user_attrs["ara_parameters"]
+
+        return {
+            name: (f"{value:.4f}" if isinstance(value, float) else f"{value}")
+            for name, value in parameters.items()
+        }
+    else:
+        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_method_description(settings: Settings) -> str:
+    if settings.use_ara:
+        return (
+            " with the [Arbitrary-Rank Ablation (ARA)](https://github.com/p-e-w/heretic/pull/211) method"
+            + (
+                " (with row-norm preservation)"
+                if settings.row_normalization == RowNormalization.FULL
+                else ""
+            )
+        )
+    elif (
+        settings.orthogonalize_direction
+        and settings.row_normalization == RowNormalization.FULL
+    ):
+        return " with a variant of the [Magnitude-Preserving Orthogonal Ablation (MPOA)](https://huggingface.co/blog/grimjim/norm-preserving-biprojected-abliteration) method"
+    else:
+        return ""
 
 
 def get_readme_intro(
     settings: Settings,
-    trial: Trial | FrozenTrial,
+    trial: Trial,
-    contains_reproducibility_information: bool,
+    base_refusals: int,
+    bad_prompts: list[Prompt],
 ) -> str:
-    if is_hf_path(settings.model):
+    if Path(settings.model).exists():
-        model_link = f"[{settings.model}](https://huggingface.co/{settings.model})"
-    else:
         # Hide the path, which may contain private information.
         model_link = "a model"
-
-    if contains_reproducibility_information:
-        reproducibility_instructions = """
-> [!TIP]
-> **This model is reproducible!**
->
-> See the [README](reproduce/README.md) in the `reproduce` directory for more information.
-"""
     else:
-        reproducibility_instructions = ""
+        model_link = f"[{settings.model}](https://huggingface.co/{settings.model})"
 
     return f"""# This is a decensored version of {
         model_link
-    }, made using [Heretic](https://heretic-project.org) v{version("heretic-llm")}
+    }, made using [Heretic](https://github.com/p-e-w/heretic) v{version("heretic-llm")}{
-{reproducibility_instructions}
+        get_method_description(settings)
+    }
+
 ## Abliteration parameters
 
 | Parameter | Value |
@@ -265,7 +333,7 @@ def get_readme_intro(
         chr(10).join(
             [
                 f"| **{name}** | {value} |"
-                for name, value in get_trial_parameters(trial).items()
+                for name, value in get_trial_parameters(settings, trial).items()
             ]
         )
     }
@@ -274,416 +342,14 @@ def get_readme_intro(
 
 | Metric | This model | Original model ({model_link}) |
 | :----- | :--------: | :---------------------------: |
-| **KL divergence** | {trial.user_attrs["kl_divergence"]:.4f} | 0 *(by definition)* |
+| **{"PIQA acc_norm" if settings.use_piqa else "KL divergence"}** | {
-| **Refusals** | {trial.user_attrs["refusals"]}/{trial.user_attrs["n_bad_prompts"]} | {
+        (-1 if settings.use_piqa else 1) * trial.user_attrs["kl_divergence"]:.4f} | {
-        trial.user_attrs["base_refusals"]
+        "*Unknown*" if settings.use_piqa else "0 *(by definition)*"
-    }/{trial.user_attrs["n_bad_prompts"]} |
+    } |
+| **Refusals** | {trial.user_attrs["refusals"]}/{len(bad_prompts)} | {base_refusals}/{
+        len(bad_prompts)
+    } |
 
 -----
 
 """
-
-
-def generate_config_toml(settings: Settings) -> str:
-    """Serializes the full Settings object to TOML."""
-
-    return tomli_w.dumps(settings.model_dump(exclude_none=True))
-
-
-def generate_requirements_txt() -> str:
-    """Collects direct project dependencies as a formatted string."""
-
-    requirements = [
-        f"{package}=={version}" for package, version in get_requirements_dict().items()
-    ]
-    return "\n".join(requirements) + "\n"
-
-
-def format_hf_link(
-    path: str,
-    commit: str | None = None,
-    is_dataset: bool = False,
-) -> str:
-    prefix = "datasets/" if is_dataset else ""
-    base_url = f"https://huggingface.co/{prefix}{path}"
-    link = f"[{path}]({base_url})"
-
-    if commit:
-        commit_url = f"{base_url}/commit/{commit}"
-        link += f" (Commit: [`{commit[:7]}`]({commit_url}))"
-
-    return link
-
-
-def generate_reproduce_readme(
-    settings: Settings,
-    checkpoint_filename: str,
-    trial: Trial | FrozenTrial,
-    include_system_information: bool,
-) -> str:
-    """Generates the contents of a README.md for the reproduce/ folder."""
-
-    heterogeneous_warning = ""
-
-    if include_system_information:
-        if torch.cuda.is_available():
-            count = torch.cuda.device_count()
-            if count > 1:
-                device_names = {torch.cuda.get_device_name(i) for i in range(count)}
-                if len(device_names) > 1:
-                    heterogeneous_warning = """
-> [!WARNING]
-> **Heterogeneous GPUs**
->
-> This model was generated using multiple non-identical GPUs. When operations are distributed across different GPUs
-> (e.g. via `device_map='auto'`), non-deterministic behavior can occur.
->
-> Reproducibility *cannot* be guaranteed in this environment.
-"""
-
-        cpu = get_cpu_info_dict()
-        python_env = get_python_env_info_dict()
-
-        accelerators = get_accelerator_info_dict()
-        if accelerators["type"] is None:
-            accelerator_report = "**No GPU or other accelerator detected.**"
-        else:
-            devices = accelerators["devices"]
-            total_vram = sum(device.get("vram_gb", 0) for device in devices)
-            vram_suffix = f" ({total_vram:.2f} GB total VRAM)" if total_vram > 0 else ""
-            accelerator_lines = [
-                f"- **{accelerators['type']}:** Detected {len(devices)} device(s){vram_suffix}"
-            ]
-
-            if accelerators.get("api_name") and accelerators.get("api_version"):
-                accelerator_lines.append(
-                    f"  - **{accelerators['api_name']}:** {accelerators['api_version']}"
-                )
-
-            if accelerators.get("driver_version"):
-                accelerator_lines.append(
-                    f"  - **Driver Version:** {accelerators['driver_version']}"
-                )
-
-            accelerator_lines.append("- **Devices:**")
-            for i, device in enumerate(devices):
-                vram = f" ({device['vram_gb']:.2f} GB)" if device.get("vram_gb") else ""
-                accelerator_lines.append(
-                    f"  - **{accelerators['type']} {i}:** {device['name']}{vram}"
-                )
-            accelerator_report = "\n".join(accelerator_lines)
-
-        system_report = f"""## System
-
-- **Python:** {python_env["version"]} ({python_env["implementation"]}, {python_env["compiler"]}) [{python_env["environment"]}]
-- **Operating system:** {platform.platform()} ({platform.machine()})
-- **CPU:** {cpu["brand"] or "Unknown"}
-
-### Accelerators
-
-{accelerator_report}
-
-"""
-        system_instructions = (
-            "1. Ensure your system matches the specifications in the **System** section above. "
-            "Exact reproducibility is only guaranteed if all aspects of your system are identical to the one the model was originally generated on.\n"
-        )
-    else:
-        system_report = ""
-        system_instructions = ""
-
-    version_info = get_heretic_version_info()
-    origin_warning = ""
-    if not version_info.is_standard_pypi:
-        if version_info.origin and version_info.origin.startswith("Git"):
-            repo_info = version_info.origin.split("Git (")[1].rstrip(")")
-            origin_warning = f"""
-> [!IMPORTANT]
-> **Git installation**
->
-> This system installed Heretic from a Git repository: {repo_info}
->
-> To reproduce the model, you must install Heretic from this exact repository and commit.
-"""
-        elif version_info.origin == "Local":
-            origin_warning = """
-> [!WARNING]
-> **Local code**
->
-> This system installed Heretic from a local directory or wheel. Uncommitted or experimental code may have been executed.
->
-> Reproducibility *cannot* be guaranteed in this environment.
-"""
-        else:
-            origin_warning = """
-> [!WARNING]
-> **Non-standard installation**
->
-> This system installed Heretic from an unknown non-standard source.
->
-> Reproducibility *cannot* be guaranteed in this environment.
-"""
-
-    pytorch_version = torch.__version__
-    pytorch_install_command = f"pip install torch=={pytorch_version}"
-    if "+" in pytorch_version:
-        suffix = pytorch_version.split("+")[1]
-        if suffix:
-            pytorch_install_command += (
-                f" --index-url https://download.pytorch.org/whl/{suffix}"
-            )
-
-    return f"""# Reproduction guide
-
-This directory contains the necessary information and assets to reproduce the results obtained during this Heretic run.{heterogeneous_warning}{origin_warning}
-
-## Models
-
-- **Base model:** {format_hf_link(settings.model, settings.model_commit)}
-
-## Datasets
-
-- **Good prompts:** {format_hf_link(settings.good_prompts.dataset, settings.good_prompts.commit, is_dataset=True)}
-- **Bad prompts:** {format_hf_link(settings.bad_prompts.dataset, settings.bad_prompts.commit, is_dataset=True)}
-- **Good evaluation prompts:** {format_hf_link(settings.good_evaluation_prompts.dataset, settings.good_evaluation_prompts.commit, is_dataset=True)}
-- **Bad evaluation prompts:** {format_hf_link(settings.bad_evaluation_prompts.dataset, settings.bad_evaluation_prompts.commit, is_dataset=True)}
-
-## Selected trial
-
-- **Trial number:** {trial.user_attrs["index"]}
-- **KL divergence:** {trial.user_attrs["kl_divergence"]:.6f}
-- **Refusals:** {trial.user_attrs["refusals"]}/{trial.user_attrs["n_bad_prompts"]}
-
-{system_report}## Environment
-
-- **Heretic:** v{version_info.version}{f" (Origin: {version_info.origin})" if version_info.origin else ""}
-- **PyTorch:** {pytorch_version}
-- **Other dependencies:** See [`requirements.txt`](requirements.txt).
-
-## Contents of this directory
-
-- [`requirements.txt`](requirements.txt): The exact versions of all Python packages.
-- [`config.toml`](config.toml): The exact configuration used, including the RNG seed.
-- [`{checkpoint_filename}`]({checkpoint_filename}): The Optuna study journal containing the history of all trials.
-- [`SHA256SUMS`](SHA256SUMS): Cryptographic hashes for all weight files.
-- [`reproduce.json`](reproduce.json): A machine-readable file containing all reproducibility information.
-
-## How to reproduce
-
-> [!TIP]
-> You can automate this process, including all verification steps, by downloading the `reproduce.json` file and running
-> `heretic --reproduce reproduce.json`.
-
-{system_instructions}1. Install the exact version of Heretic indicated in the **Environment** section above, from its original source.
-1. Install the packages listed in `requirements.txt`: `pip install -r requirements.txt`
-1. Install the correct version of PyTorch: `{pytorch_install_command}`
-1. Place the provided `config.toml` in your working directory.
-1. Run Heretic without any additional arguments: `heretic`
-1. Wait for the run to finish, then select trial **{trial.user_attrs["index"]}** and export the model.
-1. Verify that the weight files have been exactly reproduced by comparing their SHA-256 hashes against those in `SHA256SUMS`:
-   `sha256sum -c SHA256SUMS` (or look at the hashes online if you uploaded to Hugging Face)
-
-> [!TIP]
-> To use the included Optuna study journal `{checkpoint_filename}`, place it in the checkpoints directory (usually `checkpoints/`) before running Heretic.
->
-> This allows you to export other models from the Pareto front, or to run additional trials without having to re-run the stored trials.
-"""
-
-
-def generate_reproduce_json(
-    settings: Settings,
-    trial: Trial | FrozenTrial,
-    timestamp: str,
-    uploaded_model_hashes: dict[str, str],
-    include_system_information: bool,
-) -> str:
-    """Generates the contents of a reproduce.json file for the reproduce/ folder."""
-
-    version_info = get_heretic_version_info()
-
-    data = {
-        "version": "2",  # Version number of the reproduce.json file format, to allow for future changes.
-        "timestamp": timestamp,
-        "system": None,  # Defined here to preserve insertion order.
-        "environment": {
-            "heretic": {
-                "version": version_info.version,
-                "is_standard_pypi": version_info.is_standard_pypi,
-                "metadata": version_info.metadata,
-            },
-            "pytorch_version": torch.__version__,
-            "requirements": get_requirements_dict(),
-        },
-        "settings": settings.model_dump(),
-        "parameters": {
-            "direction_index": trial.user_attrs["direction_index"],
-            "abliteration_parameters": trial.user_attrs["parameters"],
-        },
-        "metrics": {
-            "kl_divergence": trial.user_attrs["kl_divergence"],
-            "refusals": trial.user_attrs["refusals"],
-            "base_refusals": trial.user_attrs["base_refusals"],
-            "n_bad_prompts": trial.user_attrs["n_bad_prompts"],
-        },
-        "hashes": uploaded_model_hashes,
-    }
-
-    if include_system_information:
-        data["system"] = {
-            "python": get_python_env_info_dict(),
-            "os": {
-                "platform": platform.platform(),
-                "machine": platform.machine(),
-            },
-            "cpu": get_cpu_info_dict(),
-            "accelerators": get_accelerator_info_dict(),
-        }
-    else:
-        del data["system"]
-
-    return json.dumps(data, indent=4)
-
-
-def generate_sha256sums(hashes: dict[str, str]) -> str:
-    """Generates GNU Coreutils compatible SHA256SUMS file content."""
-
-    lines = []
-
-    for filename, sha256 in sorted(hashes.items()):
-        # Use '*' to indicate binary mode for model weights.
-        lines.append(f"{sha256} *{filename}")
-
-    return "\n".join(lines) + "\n"
-
-
-# TODO: Replace this with hashlib.file_digest when we drop support for Python 3.10.
-def get_file_sha256(file_path: str | Path) -> str:
-    hash = hashlib.sha256()
-
-    with open(file_path, "rb") as file:
-        # Read the file in 64 kB blocks.
-        for block in iter(lambda: file.read(65536), b""):
-            hash.update(block)
-
-    return hash.hexdigest()
-
-
-def create_reproduce_folder(
-    path: Path,
-    settings: Settings,
-    checkpoint_path: str | Path,
-    trial: Trial | FrozenTrial,
-    uploaded_model_hashes: dict[str, str],
-    include_system_information: bool,
-):
-    reproduce_dir = path / "reproduce"
-    reproduce_dir.mkdir(parents=True, exist_ok=True)
-
-    checkpoint_filename = Path(checkpoint_path).name
-
-    # Fetch commit hash for the base model.
-    settings.model_commit = huggingface_hub.model_info(settings.model).sha
-
-    # Fetch commit hashes for all HF datasets to ensure reproducibility.
-    for spec in [
-        settings.good_prompts,
-        settings.bad_prompts,
-        settings.good_evaluation_prompts,
-        settings.bad_evaluation_prompts,
-    ]:
-        spec.commit = huggingface_hub.dataset_info(spec.dataset).sha
-
-    # Strip microseconds and timezone for a clean format.
-    timestamp = (
-        datetime.now(timezone.utc).replace(microsecond=0, tzinfo=None).isoformat()
-    )
-
-    (reproduce_dir / "requirements.txt").write_text(
-        generate_requirements_txt(),
-        encoding="utf-8",
-    )
-
-    (reproduce_dir / "config.toml").write_text(
-        generate_config_toml(settings),
-        encoding="utf-8",
-    )
-
-    if uploaded_model_hashes:
-        (reproduce_dir / "SHA256SUMS").write_text(
-            generate_sha256sums(uploaded_model_hashes),
-            encoding="utf-8",
-        )
-
-    (reproduce_dir / "reproduce.json").write_text(
-        generate_reproduce_json(
-            settings,
-            trial,
-            timestamp=timestamp,
-            uploaded_model_hashes=uploaded_model_hashes,
-            include_system_information=include_system_information,
-        ),
-        encoding="utf-8",
-    )
-
-    (reproduce_dir / "README.md").write_text(
-        generate_reproduce_readme(
-            settings,
-            checkpoint_filename,
-            trial,
-            include_system_information=include_system_information,
-        ),
-        encoding="utf-8",
-    )
-
-    # Copy Optuna study journal.
-    checkpoint_file = Path(checkpoint_path)
-    if checkpoint_file.exists():
-        (reproduce_dir / checkpoint_file.name).write_bytes(checkpoint_file.read_bytes())
-
-
-def upload_reproduce_folder(
-    repo_id: str,
-    settings: Settings,
-    token: str,
-    checkpoint_path: str | Path,
-    trial: Trial | FrozenTrial,
-    include_system_information: bool,
-):
-    api = huggingface_hub.HfApi()
-    info = api.model_info(repo_id=repo_id, files_metadata=True, token=token)
-
-    if not info.siblings:
-        raise RuntimeError("Could not fetch uploaded model hashes.")
-
-    # For weights, we only care about safetensors.
-    weight_extensions = (".safetensors",)
-
-    uploaded_model_hashes = {}
-
-    for file in info.siblings:
-        if file.rfilename.endswith(weight_extensions):
-            sha256 = getattr(file, "lfs", {}).get("sha256")
-            if not sha256:
-                raise RuntimeError("Could not fetch uploaded model hashes.")
-            uploaded_model_hashes[file.rfilename] = sha256
-
-    with tempfile.TemporaryDirectory() as tmpdir:
-        tmp_path = Path(tmpdir)
-        create_reproduce_folder(
-            tmp_path,
-            settings,
-            checkpoint_path=checkpoint_path,
-            trial=trial,
-            uploaded_model_hashes=uploaded_model_hashes,
-            include_system_information=include_system_information,
-        )
-
-        reproduce_dir = tmp_path / "reproduce"
-        for file_path in reproduce_dir.iterdir():
-            if file_path.is_file():
-                huggingface_hub.upload_file(
-                    path_or_fileobj=str(file_path),
-                    path_in_repo=f"reproduce/{file_path.name}",
-                    repo_id=repo_id,
-                    token=token,
-                )

tests/README.md

@@ -1,17 +0,0 @@
-Run the tests with
-
-```sh
-uv run run_tests.py
-```
-
-To update the hashes after a logic change, run the tests, then execute
-
-```sh
-cd TEST_DIR/model
-sha256sum -b * > ../SHA256SUMS.LABEL
-```
-
-where `LABEL` describes the type of system you are running the tests on.
-Since PyTorch does not guarantee exact cross-system reproducibility regardless of configuration,
-multiple valid hashes can be provided for each output file. The above update must be performed
-for each `TEST_DIR` and on each type of system.

tests/gemma-4e/SHA256SUMS.ci

@@ -1,7 +0,0 @@
-2f1b4d75d067bae3fe44e676721c7f077d243bc007156cb9c2f8b5836613d082 *chat_template.jinja
-ca80080dfa4ec6ba87152fa2b9afe70b90c400e5c4b1d6bdc3aa3114467ca68f *config.json
-70070bac883cf9c39b5992450d6b23cd160eaf33099e24c654e0359d2f87c760 *generation_config.json
-f3f4ec19504f182486459cf4e255ece265c25f827840d63b6a9d4058b8e4877a *model.safetensors
-32bdf45d2ad4cc29a0822ddd157a182de76644f0419a6228d151495256e9813c *processor_config.json
-cc8d3a0ce36466ccc1278bf987df5f71db1719b9ca6b4118264f45cb627bfe0f *tokenizer.json
-a1bab8c81ed15fa6ce912ec993c66cb49392e0487fb1ea5f5f11ea3618683627 *tokenizer_config.json

tests/gemma-4e/SHA256SUMS.ci2

@@ -1,7 +0,0 @@
-2f1b4d75d067bae3fe44e676721c7f077d243bc007156cb9c2f8b5836613d082 *chat_template.jinja
-ca80080dfa4ec6ba87152fa2b9afe70b90c400e5c4b1d6bdc3aa3114467ca68f *config.json
-70070bac883cf9c39b5992450d6b23cd160eaf33099e24c654e0359d2f87c760 *generation_config.json
-53c4ee891dce23c0ac85bebc2c4d48301469750fafbb3e6e024c15786d94db8b *model.safetensors
-32bdf45d2ad4cc29a0822ddd157a182de76644f0419a6228d151495256e9813c *processor_config.json
-cc8d3a0ce36466ccc1278bf987df5f71db1719b9ca6b4118264f45cb627bfe0f *tokenizer.json
-a1bab8c81ed15fa6ce912ec993c66cb49392e0487fb1ea5f5f11ea3618683627 *tokenizer_config.json

tests/gemma-4e/SHA256SUMS.linux

@@ -1,7 +0,0 @@
-2f1b4d75d067bae3fe44e676721c7f077d243bc007156cb9c2f8b5836613d082 *chat_template.jinja
-ca80080dfa4ec6ba87152fa2b9afe70b90c400e5c4b1d6bdc3aa3114467ca68f *config.json
-70070bac883cf9c39b5992450d6b23cd160eaf33099e24c654e0359d2f87c760 *generation_config.json
-effe36925f85ecb1e29bba84501a456bb49df21e4047be8b7ea3f6f88181fb65 *model.safetensors
-32bdf45d2ad4cc29a0822ddd157a182de76644f0419a6228d151495256e9813c *processor_config.json
-cc8d3a0ce36466ccc1278bf987df5f71db1719b9ca6b4118264f45cb627bfe0f *tokenizer.json
-a1bab8c81ed15fa6ce912ec993c66cb49392e0487fb1ea5f5f11ea3618683627 *tokenizer_config.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,41 +0,0 @@
-model = "tiny-random/gemma-4e"
-model_commit = "3a207ada2c2cd95e9671942e84cf47ea58f0f6af"
-
-seed = 12345
-print_debug_information = true
-
-batch_size = 2
-max_response_length = 10
-kl_divergence_target = 0
-n_trials = 2
-n_startup_trials = 1
-
-export_strategy = "merge"
-checkpoint_action = "restart"
-trial_index = 0
-model_action = "save"
-save_directory = "model"
-
-[good_prompts]
-dataset = "mlabonne/harmless_alpaca"
-commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f"
-split = "train[:5]"
-column = "text"
-
-[bad_prompts]
-dataset = "mlabonne/harmful_behaviors"
-commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7"
-split = "train[:5]"
-column = "text"
-
-[good_evaluation_prompts]
-dataset = "mlabonne/harmless_alpaca"
-commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f"
-split = "test[:5]"
-column = "text"
-
-[bad_evaluation_prompts]
-dataset = "mlabonne/harmful_behaviors"
-commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7"
-split = "test[:5]"
-column = "text"

tests/mistral-3/SHA256SUMS.ci

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

tests/mistral-3/SHA256SUMS.ci2

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

tests/mistral-3/SHA256SUMS.linux

@@ -1,7 +0,0 @@
-39f03c383413f531fd302c06c7e982ad98c83f0657a8339ae25478ccb81fdcda *chat_template.jinja
-f69f84977a47c8fea9ce9fc26b7de379216cb01146ea726a87996d3554cfcd19 *config.json
-34dfa6012ca9ac5f57e5521d8dbaecbc7ab7f7ab0fd96ec020b543aab5f265d9 *generation_config.json
-29aff97d5633dead9e1ccd29a2cc153b4b7431d22f63c8d6cf60bc6547681cc9 *model.safetensors
-84be30b124b50749c56d25fdbec5ccedf564446f6b3b035e88e1e07b986d2491 *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,41 +0,0 @@
-model = "tiny-random/mistral-3"
-model_commit = "931aa2e5c9668fc3679e56aa44972fe18597d55d"
-
-seed = 12345
-print_debug_information = true
-
-batch_size = 2
-max_response_length = 10
-kl_divergence_target = 0
-n_trials = 2
-n_startup_trials = 1
-
-export_strategy = "merge"
-checkpoint_action = "restart"
-trial_index = 0
-model_action = "save"
-save_directory = "model"
-
-[good_prompts]
-dataset = "mlabonne/harmless_alpaca"
-commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f"
-split = "train[:5]"
-column = "text"
-
-[bad_prompts]
-dataset = "mlabonne/harmful_behaviors"
-commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7"
-split = "train[:5]"
-column = "text"
-
-[good_evaluation_prompts]
-dataset = "mlabonne/harmless_alpaca"
-commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f"
-split = "test[:5]"
-column = "text"
-
-[bad_evaluation_prompts]
-dataset = "mlabonne/harmful_behaviors"
-commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7"
-split = "test[:5]"
-column = "text"

tests/qwen3.5-moe/SHA256SUMS.ci

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

tests/qwen3.5-moe/SHA256SUMS.linux

@@ -1,7 +0,0 @@
-a4aee8afcf2e0711942cf848899be66016f8d14a889ff9ede07bca099c28f715 *chat_template.jinja
-749b56d1b1e08081981169db6f2c44ab0be4fd6ebb452d15baafa5e09c21586a *config.json
-4625d1d64d41d1fa9dae7af4ba1e1d7e65a194073d4efa58acb266a916eaaa74 *generation_config.json
-5e0fb0ac724cf079b693fc76a515e60bc16de72c32b36c107b9f078061c4f2ef *model.safetensors
-01562eddd6f9e9ec4bc31656a3b7055284cafbf889acc6c4348dca431ae31f68 *processor_config.json
-87a7830d63fcf43bf241c3c5242e96e62dd3fdc29224ca26fed8ea333db72de4 *tokenizer.json
-2e31d1126e81bddf8d15c3f95260fb487b48c5131b24fcbb5bb9d2537e7afac0 *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

@@ -1,41 +0,0 @@
-model = "tiny-random/qwen3.5-moe"
-model_commit = "2ebfa8d9717238c5dda927008104fa172a149050"
-
-seed = 12345
-print_debug_information = true
-
-batch_size = 2
-max_response_length = 10
-kl_divergence_target = 0
-n_trials = 2
-n_startup_trials = 1
-
-export_strategy = "merge"
-checkpoint_action = "restart"
-trial_index = 0
-model_action = "save"
-save_directory = "model"
-
-[good_prompts]
-dataset = "mlabonne/harmless_alpaca"
-commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f"
-split = "train[:5]"
-column = "text"
-
-[bad_prompts]
-dataset = "mlabonne/harmful_behaviors"
-commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7"
-split = "train[:5]"
-column = "text"
-
-[good_evaluation_prompts]
-dataset = "mlabonne/harmless_alpaca"
-commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f"
-split = "test[:5]"
-column = "text"
-
-[bad_evaluation_prompts]
-dataset = "mlabonne/harmful_behaviors"
-commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7"
-split = "test[:5]"
-column = "text"

tests/run_tests.py

@@ -1,87 +0,0 @@
-# SPDX-License-Identifier: AGPL-3.0-or-later
-# Copyright (C) 2025-2026  Philipp Emanuel Weidmann <pew@worldwidemann.com> + contributors
-
-import hashlib
-import subprocess
-import sys
-from pathlib import Path
-
-
-# TODO: Replace this with hashlib.file_digest when we drop support for Python 3.10.
-def get_file_sha256(file_path: str | Path) -> str:
-    hash = hashlib.sha256()
-
-    with open(file_path, "rb") as file:
-        # Read the file in 64 kB blocks.
-        for block in iter(lambda: file.read(65536), b""):
-            hash.update(block)
-
-    return hash.hexdigest()
-
-
-script_directory = Path(__file__).resolve().parent
-
-project_directory = script_directory.parent
-
-tests_failed = False
-
-for test_directory in script_directory.iterdir():
-    if test_directory.is_dir():
-        config_file = test_directory / "config.toml"
-        hash_files = list(test_directory.glob("SHA256SUMS.*"))
-
-        if config_file.is_file() and hash_files:
-            print("#" * 50)
-            print(f"Running test {test_directory.name}")
-            print("#" * 50)
-            print()
-
-            subprocess.run(
-                [
-                    "uv",
-                    "run",
-                    "--project",
-                    project_directory,
-                    "--directory",
-                    test_directory,
-                    "heretic",
-                ],
-                check=True,
-            )
-
-            print()
-
-            valid_hashes: dict[str, list[str]] = {}
-
-            for hash_file in hash_files:
-                with open(hash_file, "r", encoding="utf-8") as file:
-                    for line in file:
-                        if line.strip():
-                            sha256, filename = line.split()
-                            filename = filename.removeprefix("*")
-
-                            if filename not in valid_hashes:
-                                valid_hashes[filename] = []
-
-                            valid_hashes[filename].append(sha256.lower())
-
-            for filename in valid_hashes:
-                sha256 = get_file_sha256(test_directory / "model" / filename)
-
-                if sha256.lower() not in valid_hashes[filename]:
-                    print(
-                        (
-                            f"Test {test_directory.name} has FAILED!\n"
-                            f"Output file {filename} doesn't match any valid hash.\n\n"
-                            f"Valid hashes:\n"
-                            f"{chr(10).join(valid_hashes[filename])}\n\n"
-                            f"Actual hash:\n"
-                            f"{sha256}\n"
-                        )
-                    )
-                    tests_failed = True
-
-if tests_failed:
-    sys.exit("Tests failed.")
-else:
-    print("All tests passed.")