feat: A better logic for detecting the thinking prefix in the response. (#423)

* fix: Check response prefix.

In case if the model adds <think> at the end of user prompt and only generates </think> end, then the detection goes wrong.

* fix: Handle whitespace in CoT, remove redundant checks

* fix: Type checker

* fix: try looking at tag positions, not text end

regex

* ruff, fix import sorting

* fix: rich markup

* fix: I missed other ones

* feat: Update SHA256SUMS file hashes in the tests.

A major change that affects reproducibility.

* fix: It is now sensible to also update the extra SHA256SUMS.ci2 file.

* fix: Only consider whitespace, no other text or instructions.

because mistral-3 as additional reasoning instructions in its chat template. And I suppose many other models can have it too.

* fix: Update windows hashes.

* fix: Update CI hashes too

* as always update case two of mistral-3 (ci2 hash)

* docs: update comment

* feat: Handle the edge case for models having additional instructions.

* fix: Update windows hash for mistral-3

* docs: remove a line from the comments

because I'm not sure about GPT-OSS models' thinking tags and it cannot be confirmed using an untrained tiny GPT-OSS model. And inference fallback would of course generate gibberish as the model cannot understand additional instructions about 'how to generate response and how to think' from the chat_template.

* fix: a few things.

* fix: Update hash for qwen3.5 after the whitespace fix for its response prefix.

* docs: Update comment

* fix: Update qwen3.5 hash for CI

* fix: Remove Case 2 which only serves tests

unnecessary

* fix: Hash

* fix: concern is valid enough, so we use a small text.

add a comment too

* docs: minor

config.default.toml

@@ -137,6 +137,8 @@ system_prompt = "You are a helpful assistant."
 # or a path to a plain text file with one prompt per line (empty lines are ignored).
 # For text files, "column" is ignored and "split" is optional; when given, it selects
 # a subset of the lines using slice notation (e.g. "[:400]").
+# "config" specifies a dataset's specific config/subset name (e.g. "english", "hindi").
+# Leave unset for datasets with a single configuration.
 
 # Dataset of prompts that tend to not result in refusals (used for calculating residual directions).
 [good_prompts]
@@ -157,6 +159,9 @@ residual_plot_color = "darkorange"
 # Plugin-specific settings live in a top-level TOML table.
 # For scorer plugins, use: `[scorer.<ClassName>]` (and optionally `[scorer.<ClassName>_<instance_name>]` for instance-related config).
 [scorer.KeywordRate]
+# Name that describes what the configured keyword rate measures.
+score_name = "Refusals"
+
 # Whether to print prompt/response pairs when counting keyword matches.
 print_responses = false
 

config.nohumor.toml

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

config.noslop.toml

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

config.piqa.toml

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

pyproject.toml

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

src/heretic/config.py

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

src/heretic/evaluator.py

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

src/heretic/main.py

@@ -39,13 +39,14 @@ import logging
 import math
 import os
 import random
+import re
 import time
 import warnings
 from dataclasses import asdict
 from importlib.metadata import version
 from os.path import commonprefix
 from pathlib import Path
-from typing import Any
+from typing import Any, cast
 
 import huggingface_hub
 import lm_eval
@@ -65,7 +66,9 @@ from optuna.storages.journal import JournalFileBackend, JournalFileOpenLock
 from optuna.trial import FrozenTrial, TrialState, create_trial
 from pydantic import ValidationError
 from questionary import Choice, Style
+from rich.markup import escape
 from rich.table import Table
+from rich.text import Text
 from rich.traceback import install
 
 from .analyzer import Analyzer
@@ -476,40 +479,88 @@ def run():
         print()
         print("Checking for common response prefix...")
         prefix_check_prompts = good_prompts[:100] + bad_prompts[:100]
-        responses = model.get_responses_batched(prefix_check_prompts)
 
-        # Despite being located in os.path, commonprefix actually performs
-        # a naive string operation without any path-specific logic,
-        # which is exactly what we need here. Trailing spaces are removed
-        # to avoid issues where multiple different tokens that all start
-        # with a space character lead to the common prefix ending with
-        # a space, which would result in an uncommon tokenization.
-        settings.response_prefix = commonprefix(responses).rstrip(" ")
+        # Detect if the model's chat template inserts a reasoning tag on its own
+        # at the end of user's prompt (e.g. <think>) by using a dummy prompt.
+        # If found, then we use the full closed CoT as the response prefix.
+        # LiquidAI's LFM models do this (Lfm2ForCausalLM).
 
-        if settings.response_prefix:
-            print(f"* Prefix found: [bold]{settings.response_prefix!r}[/]")
+        # This cast is valid because str is the return type
+        # for a single chat operation with tokenize=False.
+        dummy_prompt = cast(
+            str,
+            model.tokenizer.apply_chat_template(
+                [{"role": "user", "content": "This is a dummy prompt."}],
+                add_generation_prompt=True,
+                tokenize=False,
+            ),
+        )
 
-            for cot_initializer, closed_cot_block in settings.chain_of_thought_skips:
-                if settings.response_prefix.startswith(cot_initializer):
-                    settings.response_prefix = closed_cot_block
-                    print(
-                        f"* Closed Chain-of-Thought block: [bold]{settings.response_prefix!r}[/]"
-                    )
+        cot_skip_applied = False
 
-                    # When using a Chain-of-Thought skip, we need to check that the prefix
-                    # is actually complete (e.g. not missing a trailing newline).
-                    print("* Rechecking with prefix...")
-                    responses = model.get_responses_batched(prefix_check_prompts)
-                    additional_prefix = commonprefix(responses).rstrip(" ")
-                    if additional_prefix:
-                        settings.response_prefix += additional_prefix
+        for cot_initializer, closed_cot_block in settings.chain_of_thought_skips:
+            # Match the tag and ignore any whitespace characters following it at the end
+            # (if any), including spaces, tabs, and linebreaks. This is required for models
+            # having whitespaces after the tags.
+            pattern = rf"{re.escape(cot_initializer)}\s*$"
+            match = re.search(pattern, dummy_prompt)
+
+            if match:
+                # We use only the closed CoT block here. Any whitespaces
+                # will be handled by the 'Rechecking with prefix' logic below.
+                settings.response_prefix = closed_cot_block
+                print(
+                    f"* Closed Chain-of-Thought block: [bold]{escape(repr(settings.response_prefix))}[/]"
+                )
+                cot_skip_applied = True
+                break
+
+        # Fallback to inference for models like mistral-3 which are specifically
+        # instructed to generate thinking tags using the system prompt in their
+        # chat template, instead of inserting a prefix tag (e.g. <think>) at
+        # the end of user prompt like the case above. We expect the model to
+        # generate those tags.
+        if settings.response_prefix is None:
+            responses = model.get_responses_batched(prefix_check_prompts)
+
+            # Despite being located in os.path, commonprefix actually performs
+            # a naive string operation without any path-specific logic,
+            # which is exactly what we need here. Trailing spaces are removed
+            # to avoid issues where multiple different tokens that all start
+            # with a space character lead to the common prefix ending with
+            # a space, which would result in an uncommon tokenization.
+            settings.response_prefix = commonprefix(responses).rstrip(" ")
+
+            if settings.response_prefix:
+                print(
+                    f"* Prefix found: [bold]{escape(repr(settings.response_prefix))}[/]"
+                )
+
+                for (
+                    cot_initializer,
+                    closed_cot_block,
+                ) in settings.chain_of_thought_skips:
+                    if settings.response_prefix.startswith(cot_initializer):
+                        settings.response_prefix = closed_cot_block
                         print(
-                            f"* Extended prefix found: [bold]{settings.response_prefix!r}[/]"
+                            f"* Closed Chain-of-Thought block: [bold]{escape(repr(settings.response_prefix))}[/]"
                         )
+                        cot_skip_applied = True
+                        break
+            else:
+                print("* None found")
 
-                    break
-        else:
-            print("* None found")
+        if cot_skip_applied:
+            # When using a Chain-of-Thought skip, we need to check that the prefix
+            # is actually complete (e.g. not missing a trailing newline).
+            print("* Rechecking with prefix...")
+            responses = model.get_responses_batched(prefix_check_prompts)
+            additional_prefix = commonprefix(responses).rstrip(" ")
+            if additional_prefix:
+                settings.response_prefix += additional_prefix
+                print(
+                    f"* Extended prefix found: [bold]{escape(repr(settings.response_prefix))}[/]"
+                )
 
     evaluator = Evaluator(settings, model)
 
@@ -519,10 +570,8 @@ def run():
         settings.model = settings.evaluate_model
         model.reset_model()
         print("* Evaluating...")
-        print()
-        print("[bold]Metrics:[/]")
-        for score_name, score in evaluator.get_scores():
-            print(f"  * {score_name}: [bold]{score.rich_display}[/]")
+        for name, score in evaluator.get_scores():
+            print(f"  * [bold]{name}:[/] [green]{score.rich_display}[/]")
         return
 
     if not reproduction_mode and not evaluator.get_objective_names():
@@ -673,7 +722,7 @@ def run():
 
         print()
         print(
-            f"Running trial [bold]{trial_index}[/] of [bold]{settings.n_trials}[/]..."
+            f"[magenta]Running trial [bold]{trial_index}[/] of [bold]{settings.n_trials}[/]...[/]"
         )
         print("* Parameters:")
         for name, value in get_trial_parameters(trial).items():
@@ -685,10 +734,8 @@ def run():
         print("* Evaluating...")
         scores = evaluator.get_scores()
         objective_values = evaluator.get_objective_values(scores)
-
-        print("  * Metrics:")
         for name, score in scores:
-            print(f"    * {name}: [bold]{score.rich_display}[/]")
+            print(f"  * [bold]{name}:[/] [green]{score.rich_display}[/]")
 
         elapsed_time = time.perf_counter() - start_time
         remaining_time = (elapsed_time / (trial_index - start_index)) * (
@@ -793,7 +840,7 @@ def run():
                 score_parts: list[str] = []
                 for score in trial.user_attrs["scores"]:
                     name = score["name"]
-                    value = score["score"]["rich_display"]
+                    value = Text.from_markup(score["score"]["rich_display"]).plain
                     score_parts.append(f"{name}: {value}")
 
                 return f"{prefix} " + ", ".join(score_parts)
@@ -828,7 +875,6 @@ def run():
                         "After selecting a trial, you will be able to save the model, upload it to Hugging Face, "
                         "chat with it to test how well it works, or run standard benchmarks on it. "
                         "You can return to this menu later to select a different trial. "
-                        "[yellow]Note that KL divergence values above 0.5 usually indicate significant damage to the original model's capabilities.[/]"
                     )
                 )
 

src/heretic/scorers/benchmark_score.py

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

src/heretic/scorers/keyword_rate.py

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

src/heretic/scorers/kl_divergence.py

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

src/heretic/utils.py

@@ -208,6 +208,7 @@ def load_prompts(
                     )
             dataset = load_dataset(
                 path,
+                name=specification.config,
                 revision=specification.commit,
                 split=split_str,
             )
@@ -225,6 +226,7 @@ def load_prompts(
             # Path should be a local directory.
             dataset = load_dataset(
                 path,
+                name=specification.config,
                 split=split_str,
                 # Don't require the number of examples (lines) per split to be pre-defined.
                 verification_mode=VerificationMode.NO_CHECKS,

tests/minicpm5/config.toml

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

tests/mistral-3/SHA256SUMS.windows

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

tests/qwen2.5/config.toml

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

tests/qwen3.5-moe/SHA256SUMS.ci

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

tests/qwen3.5-moe/SHA256SUMS.windows

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

tests/run_tests.py

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

uv.lock

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