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]

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/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,6 +66,7 @@ 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
@@ -477,6 +479,48 @@ def run():
         print()
         print("Checking for common response prefix...")
         prefix_check_prompts = good_prompts[:100] + bad_prompts[:100]
+
+        # 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).
+
+        # 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,
+            ),
+        )
+
+        cot_skip_applied = False
+
+        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
@@ -488,15 +532,25 @@ def run():
             settings.response_prefix = commonprefix(responses).rstrip(" ")
 
             if settings.response_prefix:
-            print(f"* Prefix found: [bold]{settings.response_prefix!r}[/]")
+                print(
+                    f"* Prefix found: [bold]{escape(repr(settings.response_prefix))}[/]"
+                )
 
-            for cot_initializer, closed_cot_block in settings.chain_of_thought_skips:
+                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}[/]"
+                            f"* Closed Chain-of-Thought block: [bold]{escape(repr(settings.response_prefix))}[/]"
                         )
+                        cot_skip_applied = True
+                        break
+            else:
+                print("* None found")
 
+        if cot_skip_applied:
             # When using a Chain-of-Thought skip, we need to check that the prefix
             # is actually complete (e.g. not missing a trailing newline).
             print("* Rechecking with prefix...")
@@ -505,13 +559,9 @@ def run():
             if additional_prefix:
                 settings.response_prefix += additional_prefix
                 print(
-                            f"* Extended prefix found: [bold]{settings.response_prefix!r}[/]"
+                    f"* Extended prefix found: [bold]{escape(repr(settings.response_prefix))}[/]"
                 )
 
-                    break
-        else:
-            print("* None found")
-
     evaluator = Evaluator(settings, model)
 
     if settings.evaluate_model is not None:

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/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/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