build: bump version to 1.1.0

Ref #75

README.md

@@ -1,5 +1,7 @@
 # Heretic: Fully automatic censorship removal for language models
 
+[![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)
+
 Heretic is a tool that removes censorship (aka "safety alignment") from
 transformer-based language models without expensive post-training.
 It combines an advanced implementation of directional ablation, also known
@@ -37,6 +39,28 @@ e.g. `heretic --model google/gemma-3-12b-it --evaluate-model p-e-w/gemma-3-12b-i
 Note that the exact values might be platform- and hardware-dependent.
 The table above was compiled using PyTorch 2.8 on an RTX 5090.)*
 
+Of course, mathematical metrics and automated benchmarks never tell the whole
+story, and are no substitute for human evaluation. Models generated with
+Heretic have been well-received by users (links and emphasis added):
+
+> "I was skeptical before, but I just downloaded
+> [**GPT-OSS 20B Heretic**](https://huggingface.co/p-e-w/gpt-oss-20b-heretic)
+> model and holy shit. It gives properly formatted long responses to sensitive topics,
+> using the exact uncensored words that you would expect from an uncensored model,
+> produces markdown format tables with details and whatnot. Looks like this is
+> the best abliterated version of this model so far..."
+> [*(Link to comment)*](https://old.reddit.com/r/LocalLLaMA/comments/1oymku1/heretic_fully_automatic_censorship_removal_for/np6tba6/)
+
+> "[**Heretic GPT 20b**](https://huggingface.co/p-e-w/gpt-oss-20b-heretic)
+> seems to be the best uncensored model I have tried yet. It doesn't destroy a
+> the model's intelligence and it is answering prompts normally would be
+> rejected by the base model."
+> [*(Link to comment)*](https://old.reddit.com/r/LocalLLaMA/comments/1oymku1/heretic_fully_automatic_censorship_removal_for/npe9jng/)
+
+> "[[**Qwen3-4B-Instruct-2507-heretic**](https://huggingface.co/p-e-w/Qwen3-4B-Instruct-2507-heretic)]
+> 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 supports most dense models, including many multimodal models, and
 several different MoE architectures. It does not yet support SSMs/hybrid models,
 models with inhomogeneous layers, and certain novel attention systems.
@@ -51,7 +75,7 @@ Prepare a Python 3.10+ environment with PyTorch 2.2+ installed as appropriate
 for your hardware. Then run:
 
 ```
-pip install heretic-llm
+pip install -U heretic-llm
 heretic Qwen/Qwen3-4B-Instruct-2507
 ```
 
@@ -73,7 +97,88 @@ save the model, upload it to Hugging Face, chat with it to test how well it work
 or any combination of those actions.
 
 
-## How it works
+## Research features
+
+In addition to its primary function of removing model censorship, Heretic also
+provides features designed to support research into the semantics of model internals
+(interpretability). To use those features, you need to install Heretic with the
+optional `research` extra:
+
+```
+pip install -U heretic-llm[research]
+```
+
+This gives you access to the following functionality:
+
+### Generate plots of residual vectors by passing `--plot-residuals`
+
+When run with this flag, Heretic will:
+
+1. Compute residual vectors (hidden states) for the first output token,
+   for each transformer layer, for both "harmful" and "harmless" prompts.
+2. Perform a [PaCMAP projection](https://github.com/YingfanWang/PaCMAP)
+   from residual space to 2D-space.
+3. Left-right align the projections of "harmful"/"harmless" residuals
+   by their geometric medians to make projections for consecutive layers
+   more similar. Additionally, PaCMAP is initialized with the previous
+   layer's projections for each new layer, minimizing disruptive transitions.
+4. Scatter-plot the projections, generating a PNG image for each layer.
+5. Generate an animation showing how residuals transform between layers,
+   as an animated GIF.
+
+<img width="800" height="600" alt="Plot of residual vectors" src="https://github.com/user-attachments/assets/981aa6ed-5ab9-48f0-9abf-2b1a2c430295" />
+
+See [the configuration file](config.default.toml) for options that allow you
+to control various aspects of the generated plots.
+
+Note that PaCMAP is an expensive operation that is performed on the CPU.
+For larger models, it can take an hour or more to compute projections
+for all layers.
+
+### Print details about residual geometry by passing `--print-residual-geometry`
+
+If you are interested in a quantitative analysis of how residual vectors
+for "harmful" and "harmless" prompts relate to each other, this flag gives you
+the following table, packed with metrics that can facilitate understanding
+the same (for [gemma-3-270m-it](https://huggingface.co/google/gemma-3-270m-it)
+in this case):
+
+```
+┏━━━━━━━┳━━━━━━━━┳━━━━━━━━━━┳━━━━━━━━━┳━━━━━━━━━━┳━━━━━━━━━┳━━━━━━━━━━┳━━━━━━━━━━┳━━━━━━━━━━┳━━━━━━━━━━┳━━━━━━━━━━┳━━━━━━━━━┳━━━━━━━━━┳━━━━━━━━┓
+┃ Layer ┃ S(g,b) ┃ S(g*,b*) ┃  S(g,r) ┃ S(g*,r*) ┃  S(b,r) ┃ S(b*,r*) ┃      |g| ┃     |g*| ┃      |b| ┃     |b*| ┃     |r| ┃    |r*| ┃   Silh ┃
+┡━━━━━━━╇━━━━━━━━╇━━━━━━━━━━╇━━━━━━━━━╇━━━━━━━━━━╇━━━━━━━━━╇━━━━━━━━━━╇━━━━━━━━━━╇━━━━━━━━━━╇━━━━━━━━━━╇━━━━━━━━━━╇━━━━━━━━━╇━━━━━━━━━╇━━━━━━━━┩
+│     1 │ 1.0000 │   1.0000 │ -0.4311 │  -0.4906 │ -0.4254 │  -0.4847 │   170.29 │   170.49 │   169.78 │   169.85 │    1.19 │    1.31 │ 0.0480 │
+│     2 │ 1.0000 │   1.0000 │  0.4297 │   0.4465 │  0.4365 │   0.4524 │   768.55 │   768.77 │   771.32 │   771.36 │    6.39 │    5.76 │ 0.0745 │
+│     3 │ 0.9999 │   1.0000 │ -0.5699 │  -0.5577 │ -0.5614 │  -0.5498 │  1020.98 │  1021.13 │  1013.80 │  1014.71 │   12.70 │   11.60 │ 0.0920 │
+│     4 │ 0.9999 │   1.0000 │  0.6582 │   0.6553 │  0.6659 │   0.6627 │  1356.39 │  1356.20 │  1368.71 │  1367.95 │   18.62 │   17.84 │ 0.0957 │
+│     5 │ 0.9987 │   0.9990 │ -0.6880 │  -0.6761 │ -0.6497 │  -0.6418 │   766.54 │   762.25 │   731.75 │   732.42 │   51.97 │   45.24 │ 0.1018 │
+│     6 │ 0.9998 │   0.9998 │ -0.1983 │  -0.2312 │ -0.1811 │  -0.2141 │  2417.35 │  2421.08 │  2409.18 │  2411.40 │   43.06 │   43.47 │ 0.0900 │
+│     7 │ 0.9998 │   0.9997 │ -0.5258 │  -0.5746 │ -0.5072 │  -0.5560 │  3444.92 │  3474.99 │  3400.01 │  3421.63 │   86.94 │   94.38 │ 0.0492 │
+│     8 │ 0.9990 │   0.9991 │  0.8235 │   0.8312 │  0.8479 │   0.8542 │  4596.54 │  4615.62 │  4918.32 │  4934.20 │  384.87 │  377.87 │ 0.2278 │
+│     9 │ 0.9992 │   0.9992 │  0.5335 │   0.5441 │  0.5678 │   0.5780 │  5322.30 │  5316.96 │  5468.65 │  5466.98 │  265.68 │  267.28 │ 0.1318 │
+│    10 │ 0.9974 │   0.9973 │  0.8189 │   0.8250 │  0.8579 │   0.8644 │  5328.81 │  5325.63 │  5953.35 │  5985.15 │  743.95 │  779.74 │ 0.2863 │
+│    11 │ 0.9977 │   0.9978 │  0.4262 │   0.4045 │  0.4862 │   0.4645 │  9644.02 │  9674.06 │  9983.47 │  9990.28 │  743.28 │  726.99 │ 0.1576 │
+│    12 │ 0.9904 │   0.9907 │  0.4384 │   0.4077 │  0.5586 │   0.5283 │ 10257.40 │ 10368.50 │ 11114.51 │ 11151.21 │ 1711.18 │ 1664.69 │ 0.1890 │
+│    13 │ 0.9867 │   0.9874 │  0.4007 │   0.3680 │  0.5444 │   0.5103 │ 12305.12 │ 12423.75 │ 13440.31 │ 13432.47 │ 2386.43 │ 2282.47 │ 0.1293 │
+│    14 │ 0.9921 │   0.9922 │  0.3198 │   0.2682 │  0.4364 │   0.3859 │ 16929.16 │ 17080.37 │ 17826.97 │ 17836.03 │ 2365.23 │ 2301.87 │ 0.1282 │
+│    15 │ 0.9846 │   0.9850 │  0.1198 │   0.0963 │  0.2913 │   0.2663 │ 16858.58 │ 16949.44 │ 17496.00 │ 17502.88 │ 3077.08 │ 3029.60 │ 0.1611 │
+│    16 │ 0.9686 │   0.9689 │ -0.0029 │  -0.0254 │  0.2457 │   0.2226 │ 18912.77 │ 19074.86 │ 19510.56 │ 19559.62 │ 4848.35 │ 4839.75 │ 0.1516 │
+│    17 │ 0.9782 │   0.9784 │ -0.0174 │  -0.0381 │  0.1908 │   0.1694 │ 27098.09 │ 27273.00 │ 27601.12 │ 27653.12 │ 5738.19 │ 5724.21 │ 0.1641 │
+│    18 │ 0.9184 │   0.9196 │  0.1343 │   0.1430 │  0.5155 │   0.5204 │   190.16 │   190.35 │   219.91 │   220.62 │   87.82 │   87.59 │ 0.1855 │
+└───────┴────────┴──────────┴─────────┴──────────┴─────────┴──────────┴──────────┴──────────┴──────────┴──────────┴─────────┴─────────┴────────┘
+g = mean of residual vectors for good prompts
+g* = geometric median of residual vectors for good prompts
+b = mean of residual vectors for bad prompts
+b* = geometric median of residual vectors for bad prompts
+r = refusal direction for means (i.e., b - g)
+r* = refusal direction for geometric medians (i.e., b* - g*)
+S(x,y) = cosine similarity of x and y
+|x| = L2 norm of x
+Silh = Mean silhouette coefficient of residuals for good/bad clusters
+```
+
+
+## How Heretic works
 
 Heretic implements a parametrized variant of directional ablation. For each
 supported transformer component (currently, attention out-projection and

config.default.toml

@@ -7,8 +7,11 @@ dtypes = [
     "auto",
     # If that doesn't work (e.g. on pre-Ampere hardware), fall back to float16.
     "float16",
-    # If that still doesn't work (e.g. due to https://github.com/meta-llama/llama/issues/380),
-    # fall back to float32.
+    # If "auto" resolves to float32, and that fails because it is too large,
+    # and float16 fails due to range issues, try bfloat16.
+    "bfloat16",
+    # If neither of those work, fall back to float32 (which will of course fail
+    # if that was the dtype "auto" resolved to).
     "float32",
 ]
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "heretic-llm"
-version = "1.0.1"
+version = "1.1.0"
 description = "Fully automatic censorship removal for language models"
 readme = "README.md"
 license = "AGPL-3.0-or-later"
@@ -40,6 +40,7 @@ research = [
     "matplotlib>=3.10.7",
     "numpy>=2.2.6",
     "pacmap>=0.8.0",
+    "scikit-learn>=1.7.2",
 ]
 
 [dependency-groups]
@@ -52,7 +53,7 @@ Homepage = "https://github.com/p-e-w/heretic"
 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/commits/master/"
+Changelog = "https://github.com/p-e-w/heretic/releases"
 
 [project.scripts]
 heretic = "heretic.main:main"

src/heretic/analyzer.py

@@ -6,6 +6,7 @@ from pathlib import Path
 import torch
 import torch.linalg as LA
 import torch.nn.functional as F
+from rich.progress import track
 from rich.table import Table
 from torch import Tensor
 
@@ -30,13 +31,14 @@ class Analyzer:
     def print_residual_geometry(self):
         try:
             from geom_median.torch import compute_geometric_median
+            from sklearn.metrics import silhouette_score
         except ImportError:
             print()
             print(
                 (
                     "[red]Research dependencies not found. Printing residual geometry requires "
                     "installing Heretic with the optional research feature, i.e., "
-                    'using "pip install heretic-llm\\[research]".[/]'
+                    'using "pip install -U heretic-llm\\[research]".[/]'
                 )
             )
             return
@@ -58,6 +60,7 @@ class Analyzer:
         table.add_column("|b*|", justify="right")
         table.add_column("|r|", justify="right")
         table.add_column("|r*|", justify="right")
+        table.add_column("Silh", justify="right")
 
         g = self.good_residuals.mean(dim=0)
         g_star = torch.stack(
@@ -94,6 +97,24 @@ class Analyzer:
         r_norms = LA.vector_norm(r, dim=-1)
         r_star_norms = LA.vector_norm(r_star, dim=-1)
 
+        residuals = (
+            torch.cat(
+                [
+                    self.good_residuals,
+                    self.bad_residuals,
+                ],
+                dim=0,
+            )
+            .detach()
+            .cpu()
+            .numpy()
+        )
+        labels = [0] * len(self.good_residuals) + [1] * len(self.bad_residuals)
+        silhouettes = [
+            silhouette_score(residuals[:, layer_index, :], labels)
+            for layer_index in range(len(self.model.get_layers()) + 1)
+        ]
+
         for layer_index in range(1, len(self.model.get_layers()) + 1):
             table.add_row(
                 f"{layer_index}",
@@ -109,6 +130,7 @@ class Analyzer:
                 f"{b_star_norms[layer_index].item():.2f}",
                 f"{r_norms[layer_index].item():.2f}",
                 f"{r_star_norms[layer_index].item():.2f}",
+                f"{silhouettes[layer_index]:.4f}",
             )
 
         print()
@@ -124,6 +146,9 @@ class Analyzer:
         )
         print("[bold]S(x,y)[/] = cosine similarity of [bold]x[/] and [bold]y[/]")
         print("[bold]|x|[/] = L2 norm of [bold]x[/]")
+        print(
+            "[bold]Silh[/] = Mean silhouette coefficient of residuals for good/bad clusters"
+        )
 
     def plot_residuals(self):
         try:
@@ -139,7 +164,7 @@ class Analyzer:
                 (
                     "[red]Research dependencies not found. Plotting residuals requires "
                     "installing Heretic with the optional research feature, i.e., "
-                    'using "pip install heretic-llm\\[research]".[/]'
+                    'using "pip install -U heretic-llm\\[research]".[/]'
                 )
             )
             return
@@ -150,12 +175,14 @@ class Analyzer:
 
         print()
         print("Plotting residual vectors...")
-        print("* Computing PaCMAP projections...")
 
         layer_residuals_2d = []
         pacmap_init = None
 
-        for layer_index in range(1, len(self.model.get_layers()) + 1):
+        for layer_index in track(
+            range(1, len(self.model.get_layers()) + 1),
+            description="* Computing PaCMAP projections...",
+        ):
             good_residuals = (
                 self.good_residuals[:, layer_index, :].detach().cpu().numpy()
             )
@@ -192,8 +219,6 @@ class Analyzer:
 
             layer_residuals_2d.append((good_residuals_2d, bad_residuals_2d))
 
-        print("* Generating plots...")
-
         plt.style.use(self.settings.residual_plot_style)
 
         def plot(
@@ -268,7 +293,13 @@ class Analyzer:
         for layer_index, (
             good_residuals_2d,
             bad_residuals_2d,
-        ) in enumerate(layer_residuals_2d, 1):
+        ) in enumerate(
+            track(
+                layer_residuals_2d,
+                description="* Generating plots...",
+            ),
+            1,
+        ):
             image_path = base_path / f"layer_{layer_index:03}.png"
 
             plot(image_path, layer_index, good_residuals_2d, bad_residuals_2d)
@@ -310,6 +341,8 @@ class Analyzer:
                     # other than building the animation.
                     image_path.unlink()
 
+        print("* Generating animation...")
+
         iio.imwrite(
             base_path / "animation.gif",
             images,

src/heretic/config.py

@@ -46,10 +46,11 @@ class Settings(BaseSettings):
             "auto",
             # If that doesn't work (e.g. on pre-Ampere hardware), fall back to float16.
             "float16",
-            # If float16 fails (e.g. due to range issues) and float32 is too large, try bfloat16.
+            # If "auto" resolves to float32, and that fails because it is too large,
+            # and float16 fails due to range issues, try bfloat16.
             "bfloat16",
-            # If that still doesn't work (e.g. due to https://github.com/meta-llama/llama/issues/380),
-            # fall back to float32.
+            # If neither of those work, fall back to float32 (which will of course fail
+            # if that was the dtype "auto" resolved to).
             "float32",
         ],
         description="List of PyTorch dtypes to try when loading model tensors. If loading with a dtype fails, the next dtype in the list will be tried.",
@@ -212,6 +213,7 @@ class Settings(BaseSettings):
         toml_file="config.toml",
         env_prefix="HERETIC_",
         cli_parse_args=True,
+        cli_implicit_flags=True,
         cli_kebab_case=True,
     )
 

src/heretic/evaluator.py

@@ -70,7 +70,7 @@ class Evaluator:
             reduction="batchmean",
             log_target=True,
         ).item()
-        print(f"  * KL divergence: [bold]{kl_divergence:.2f}[/]")
+        print(f"  * KL divergence: [bold]{kl_divergence:.4f}[/]")
 
         print("  * Counting model refusals...")
         refusals = self.count_refusals()

src/heretic/main.py

@@ -7,6 +7,7 @@ import sys
 import time
 import warnings
 from importlib.metadata import version
+from os.path import commonprefix
 from pathlib import Path
 
 import huggingface_hub
@@ -27,7 +28,7 @@ from optuna.exceptions import ExperimentalWarning
 from optuna.samplers import TPESampler
 from optuna.study import StudyDirection
 from pydantic import ValidationError
-from questionary import Choice, Style
+from questionary import Choice
 from rich.traceback import install
 
 from .analyzer import Analyzer
@@ -65,10 +66,11 @@ def run():
     print()
 
     if (
-        # An odd number of arguments have been passed (argv[0] is the program name),
-        # so that after accounting for "--param VALUE" pairs, there is one left over.
-        len(sys.argv) % 2 == 0
-        # The leftover argument is a parameter value rather than a flag (such as "--help").
+        # There is at least one argument (argv[0] is the program name).
+        len(sys.argv) > 1
+        # No model has been explicitly provided.
+        and "--model" not in sys.argv
+        # The last argument is a parameter value rather than a flag (such as "--help").
         and not sys.argv[-1].startswith("-")
     ):
         # Assume the last argument is the model.
@@ -186,6 +188,31 @@ def run():
         settings.batch_size = best_batch_size
         print(f"* Chosen batch size: [bold]{settings.batch_size}[/]")
 
+    print()
+    print("Checking for common response prefix...")
+    responses = model.get_responses_batched(good_prompts[:100] + bad_prompts[:100])
+
+    # Despite being located in os.path, commonprefix actually performs
+    # a naive string operation without any path-specific logic,
+    # which is exactly what we need here. Trailing spaces are removed
+    # to avoid issues where multiple different tokens that all start
+    # with a space character lead to the common prefix ending with
+    # a space, which would result in an uncommon tokenization.
+    model.response_prefix = commonprefix(responses).rstrip(" ")
+
+    # Suppress CoT output.
+    if model.response_prefix.startswith("<think>"):
+        # Most thinking models.
+        model.response_prefix = "<think></think>"
+    elif model.response_prefix.startswith("<|channel|>analysis<|message|>"):
+        # gpt-oss.
+        model.response_prefix = "<|channel|>analysis<|message|><|end|><|start|>assistant<|channel|>final<|message|>"
+
+    if model.response_prefix:
+        print(f"* Prefix found: [bold]{model.response_prefix!r}[/]")
+    else:
+        print("* None found")
+
     evaluator = Evaluator(settings, model)
 
     if settings.evaluate_model is not None:
@@ -364,7 +391,7 @@ def run():
             title=(
                 f"[Trial {trial.user_attrs['index']:>3}] "
                 f"Refusals: {trial.user_attrs['refusals']:>2}/{len(evaluator.bad_prompts)}, "
-                f"KL divergence: {trial.user_attrs['kl_divergence']:.2f}"
+                f"KL divergence: {trial.user_attrs['kl_divergence']:.4f}"
             ),
             value=trial,
         )
@@ -392,11 +419,7 @@ def run():
 
     while True:
         print()
-        trial = prompt_select(
-            "Which trial do you want to use?",
-            choices=choices,
-            style=Style([("highlighted", "reverse")]),
-        )
+        trial = prompt_select("Which trial do you want to use?", choices)
 
         if trial is None or trial == "":
             break
@@ -416,13 +439,12 @@ def run():
             print()
             action = prompt_select(
                 "What do you want to do with the decensored model?",
-                choices=[
+                [
                     "Save the model to a local folder",
                     "Upload the model to Hugging Face",
                     "Chat with the model",
                     "Nothing (return to trial selection menu)",
                 ],
-                style=Style([("highlighted", "reverse")]),
             )
 
             if action is None or action == "Nothing (return to trial selection menu)":
@@ -434,9 +456,7 @@ def run():
             try:
                 match action:
                     case "Save the model to a local folder":
-                        save_directory = prompt_path(
-                            "Path to the folder:", only_directories=True
-                        )
+                        save_directory = prompt_path("Path to the folder:")
                         if not save_directory:
                             continue
 
@@ -470,11 +490,10 @@ def run():
 
                         visibility = prompt_select(
                             "Should the repository be public or private?",
-                            choices=[
+                            [
                                 "Public",
                                 "Private",
                             ],
-                            style=Style([("highlighted", "reverse")]),
                         )
                         private = visibility == "Private"
 

src/heretic/model.py

@@ -34,6 +34,7 @@ class AbliterationParameters:
 class Model:
     def __init__(self, settings: Settings):
         self.settings = settings
+        self.response_prefix = ""
 
         print()
         print(f"Loading model [bold]{settings.model}[/]...")
@@ -70,7 +71,7 @@ class Model:
                 )
 
                 # If we reach this point and the model requires trust_remote_code,
-                # the user must have confirmed it.
+                # either the user accepted, or settings.trust_remote_code is True.
                 if self.trusted_models.get(settings.model) is None:
                     self.trusted_models[settings.model] = True
 
@@ -231,27 +232,16 @@ class Model:
 
                 # Projects any right-multiplied vector(s) onto the subspace
                 # spanned by the refusal direction.
-                # We use the property (r r^T) W = r (r^T W) to avoid computing
-                # the O(d^2) projector matrix and the O(d^2 k) matrix multiplication.
-                # (α is the weight)
-                # W_new = W - α(r (r^T W))
-                r = layer_refusal_direction.to(self.model.dtype)
+                projector = torch.outer(
+                    layer_refusal_direction,
+                    layer_refusal_direction,
+                ).to(self.model.dtype)
 
                 for matrix in matrices:
-                    # Ensure r is on the same device as the matrix for multi-GPU support.
-                    r_device = r.to(matrix.device)
-
-                    # Calculate the projection scalars: (r^T W)
-                    # r is (d,), matrix is (d, k) -> result is (k,)
-                    r_transpose_W = torch.matmul(r_device, matrix)
-
-                    # Compute the rank-1 update r (r^T W) using the outer product form
-                    # r_device: (d,)          — projection direction
-                    # r_transpose_W: (k,)     — r^T W result for this matrix
-                    # torch.outer(r_device, r_times_W) constructs the (d, k) matrix with
-                    # entries r[i] * (r^T W)[j], equivalent to the outer product of two
-                    # vectors, avoiding materializing the full (d x d) projector.
-                    matrix.sub_(weight * torch.outer(r_device, r_transpose_W))
+                    # Ensure projector is on the same device as the matrix for multi-GPU support.
+                    device_projector = projector.to(matrix.device)
+                    # In-place subtraction is safe as we're not using Autograd.
+                    matrix.sub_(weight * (device_projector @ matrix))
 
     def get_chat(self, prompt: str) -> list[dict[str, str]]:
         return [
@@ -272,6 +262,11 @@ class Model:
             tokenize=False,
         )
 
+        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 = [prompt + self.response_prefix for prompt in chat_prompts]
+
         inputs = self.tokenizer(
             chat_prompts,
             return_tensors="pt",
@@ -282,7 +277,7 @@ class Model:
         return inputs, self.model.generate(
             **inputs,
             **kwargs,
-            pad_token_id=self.tokenizer.eos_token_id,
+            pad_token_id=self.tokenizer.pad_token_id,
             do_sample=False,  # Use greedy decoding to ensure deterministic outputs.
         )
 
@@ -293,10 +288,7 @@ class Model:
         )
 
         # Return only the newly generated part.
-        return self.tokenizer.batch_decode(
-            outputs[:, inputs["input_ids"].shape[1] :],
-            skip_special_tokens=True,
-        )
+        return self.tokenizer.batch_decode(outputs[:, inputs["input_ids"].shape[1] :])
 
     def get_responses_batched(self, prompts: list[str]) -> list[str]:
         responses = []

src/heretic/utils.py

@@ -22,7 +22,7 @@ from datasets.config import DATASET_STATE_JSON_FILENAME
 from datasets.download.download_manager import DownloadMode
 from datasets.utils.info_utils import VerificationMode
 from optuna import Trial
-from questionary import Choice
+from questionary import Choice, Style
 from rich.console import Console
 
 from .config import DatasetSpecification, Settings
@@ -31,15 +31,15 @@ print = Console(highlight=False).print
 
 
 def is_notebook() -> bool:
-    # Check for specific environment variables (Colab, Kaggle)
+    # 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)
+    # Check IPython shell type (for library usage).
     try:
-        from IPython import get_ipython
+        from IPython import get_ipython  # pyright: ignore[reportMissingModuleSource]
 
         shell = get_ipython()
         if shell is None:
@@ -57,11 +57,12 @@ def is_notebook() -> bool:
         return False
 
 
-def prompt_select(message: str, choices: list[Any], style=None) -> Any:
+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}")
@@ -73,47 +74,45 @@ def prompt_select(message: str, choices: list[Any], style=None) -> Any:
         while True:
             try:
                 selection = input("Enter number: ")
-                idx = int(selection) - 1
-                if 0 <= idx < len(real_choices):
-                    return real_choices[idx]
+                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).ask()
+        return questionary.select(
+            message,
+            choices=choices,
+            style=Style([("highlighted", "reverse")]),
+        ).ask()
 
 
 def prompt_text(
     message: str,
     default: str = "",
-    unsafe: bool = False,
     qmark: str = "?",
+    unsafe: bool = False,
 ) -> str:
     if is_notebook():
         print()
-        prompt_msg = f"{message} [{default}]: " if default else f"{message}: "
-        result = input(prompt_msg)
+        result = input(f"{message} [{default}]: " if default else f"{message}: ")
         return result if result else default
     else:
-        # For text input, we might need unsafe_ask if requested
-        q = questionary.text(message, default=default, qmark=qmark)
+        question = questionary.text(message, default=default, qmark=qmark)
         if unsafe:
-            return q.unsafe_ask()
-        return q.ask()
+            return question.unsafe_ask()
+        else:
+            return question.ask()
 
 
-def prompt_path(message: str, default: str = "", only_directories: bool = False) -> str:
+def prompt_path(message: str) -> str:
     if is_notebook():
-        print()
-        prompt_msg = f"{message} [{default}]: " if default else f"{message}: "
-        result = input(prompt_msg)
-        return result if result else default
+        return prompt_text(message)
     else:
-        return questionary.path(
-            message, default=default, only_directories=only_directories
-        ).ask()
+        return questionary.path(message, only_directories=True).ask()
 
 
 def prompt_password(message: str) -> str:
@@ -140,20 +139,21 @@ def format_duration(seconds: float) -> str:
 def load_prompts(specification: DatasetSpecification) -> list[str]:
     path = specification.dataset
     split_str = specification.split
+
     if os.path.isdir(path):
         if 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)
             # Parse the split instructions.
-            ri = ReadInstruction.from_spec(split_str)
+            instruction = ReadInstruction.from_spec(split_str)
             # 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_i = ri.to_absolute(name2len)[0]
+            abs_instruction = instruction.to_absolute(name2len)[0]
             # Get the dataset by applying the indices.
-            dataset = dataset[abs_i.from_ : abs_i.to]
+            dataset = dataset[abs_instruction.from_ : abs_instruction.to]
         else:
             # Path is a local directory.
             dataset = load_dataset(
@@ -244,7 +244,7 @@ def get_readme_intro(
 
 | Metric | This model | Original model ({model_link}) |
 | :----- | :--------: | :---------------------------: |
-| **KL divergence** | {trial.user_attrs["kl_divergence"]:.2f} | 0 *(by definition)* |
+| **KL divergence** | {trial.user_attrs["kl_divergence"]:.4f} | 0 *(by definition)* |
 | **Refusals** | {trial.user_attrs["refusals"]}/{len(bad_prompts)} | {base_refusals}/{
         len(bad_prompts)
     } |

uv.lock

@@ -725,7 +725,7 @@ wheels = [
 
 [[package]]
 name = "heretic-llm"
-version = "1.0.1"
+version = "1.1.0"
 source = { editable = "." }
 dependencies = [
     { name = "accelerate" },
@@ -747,6 +747,7 @@ research = [
     { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
     { name = "numpy", version = "2.3.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" },
     { name = "pacmap" },
+    { name = "scikit-learn" },
 ]
 
 [package.dev-dependencies]
@@ -769,6 +770,7 @@ requires-dist = [
     { name = "pydantic-settings", specifier = ">=2.10.1" },
     { name = "questionary", specifier = ">=2.1.1" },
     { name = "rich", specifier = ">=14.1.0" },
+    { name = "scikit-learn", marker = "extra == 'research'", specifier = ">=1.7.2" },
     { name = "transformers", specifier = ">=4.55.2" },
 ]
 provides-extras = ["research"]