diff --git a/src/heretic/main.py b/src/heretic/main.py
index fbc6cd59..a26be5b2 100644
--- a/src/heretic/main.py
+++ b/src/heretic/main.py
@@ -39,6 +39,7 @@ def _is_help_invocation() -> bool:
import math
import os
import random
+import re
import time
import warnings
from dataclasses import asdict
@@ -65,6 +66,7 @@ def _is_help_invocation() -> bool:
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.traceback import install
@@ -476,40 +478,81 @@ 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(" ")
-
- if settings.response_prefix:
- print(f"* Prefix found: [bold]{settings.response_prefix!r}[/]")
-
- 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}[/]"
- )
- # 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
+ # Detect if the model's chat template inserts a reasoning tag on its own
+ # in the end of user's prompt (e.g. ) by using a dummy prompt.
+ # If found, then we extract the end of it (e.g. ) as the response prefix.
+ # Otherwise we fall back to real prompts inference.
+ # LiquidAI's LFM models do this (Lfm2ForCausalLM).
+ dummy_prompt = model.tokenizer.apply_chat_template(
+ [{"role": "user", "content": ""}],
+ add_generation_prompt=True,
+ tokenize=False,
+ )
+
+ cot_skip_applied = False
+
+ # Some chat-templates add whitespace after the tag (e.g. Qwen3.5 adds "\n"),
+ # so we just strip any whitespace from the end before checking.
+ assert isinstance(dummy_prompt, str)
+
+ for cot_initializer, closed_cot_block in settings.chain_of_thought_skips:
+ # Match the tag and capture any text or whitespace following it in the end.
+ pattern = rf"{re.escape(cot_initializer)}(.*)$"
+ match = re.search(pattern, dummy_prompt, re.DOTALL)
+
+ if match:
+ trailing_content = match.group(1)
+ # Get the closing tag and add the spaces or newlines found.
+ base_closed_block = closed_cot_block[len(cot_initializer) :]
+ settings.response_prefix = base_closed_block + trailing_content
+ print(
+ f"* Closed Chain-of-Thought block: [bold]{escape(repr(settings.response_prefix))}[/]"
+ )
+ cot_skip_applied = True
+ break
+
+ 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))}[/]"
)
-
- break
- else:
- print("* None found")
+ 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...")
+ 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)
diff --git a/tests/mistral-3/SHA256SUMS.ci b/tests/mistral-3/SHA256SUMS.ci
index d9a60544..e793f62d 100644
--- a/tests/mistral-3/SHA256SUMS.ci
+++ b/tests/mistral-3/SHA256SUMS.ci
@@ -1,7 +1,7 @@
39f03c383413f531fd302c06c7e982ad98c83f0657a8339ae25478ccb81fdcda *chat_template.jinja
f69f84977a47c8fea9ce9fc26b7de379216cb01146ea726a87996d3554cfcd19 *config.json
34dfa6012ca9ac5f57e5521d8dbaecbc7ab7f7ab0fd96ec020b543aab5f265d9 *generation_config.json
-876c6691eb85e3e5e11771e589529830fb454ab26344e1271ae550661e312b50 *model.safetensors
+caf6f7a04fd6a9c194ac7be4e92981ddb22a9003ff34fc31a247340edae27006 *model.safetensors
84be30b124b50749c56d25fdbec5ccedf564446f6b3b035e88e1e07b986d2491 *processor_config.json
c3a8d92e371b92a2cd6e678e31ebc27d0235e929a51fbf290f74742b341fa96f *tokenizer.json
7b29c843c0043622d28fd4638451cbb0a609d99a0762ffbff3b92b4b2fee4d94 *tokenizer_config.json
diff --git a/tests/mistral-3/SHA256SUMS.ci2 b/tests/mistral-3/SHA256SUMS.ci2
index 8729d400..e24d27f8 100644
--- a/tests/mistral-3/SHA256SUMS.ci2
+++ b/tests/mistral-3/SHA256SUMS.ci2
@@ -1,7 +1,7 @@
39f03c383413f531fd302c06c7e982ad98c83f0657a8339ae25478ccb81fdcda *chat_template.jinja
f69f84977a47c8fea9ce9fc26b7de379216cb01146ea726a87996d3554cfcd19 *config.json
34dfa6012ca9ac5f57e5521d8dbaecbc7ab7f7ab0fd96ec020b543aab5f265d9 *generation_config.json
-6febb813086f253e5ec0fcda02fdfc849c551a7dba54681b37ac5bc402e4eed6 *model.safetensors
+f4c9816506ad3c540d707fecb9e85c57acb424529a322b3f1e625c21f507b815 *model.safetensors
84be30b124b50749c56d25fdbec5ccedf564446f6b3b035e88e1e07b986d2491 *processor_config.json
c3a8d92e371b92a2cd6e678e31ebc27d0235e929a51fbf290f74742b341fa96f *tokenizer.json
7b29c843c0043622d28fd4638451cbb0a609d99a0762ffbff3b92b4b2fee4d94 *tokenizer_config.json
diff --git a/tests/qwen3.5-moe/SHA256SUMS.ci b/tests/qwen3.5-moe/SHA256SUMS.ci
index 44c28f5a..6bdeb97d 100644
--- a/tests/qwen3.5-moe/SHA256SUMS.ci
+++ b/tests/qwen3.5-moe/SHA256SUMS.ci
@@ -1,7 +1,7 @@
a4aee8afcf2e0711942cf848899be66016f8d14a889ff9ede07bca099c28f715 *chat_template.jinja
749b56d1b1e08081981169db6f2c44ab0be4fd6ebb452d15baafa5e09c21586a *config.json
4625d1d64d41d1fa9dae7af4ba1e1d7e65a194073d4efa58acb266a916eaaa74 *generation_config.json
-5fb94c65bcd9d736735a45e50c2b0bfafd3bb09a444c49b8cff2e131ed35797e *model.safetensors
+2951a008dc8e2bab778cd8b78996e4d51c628fd2648f493b123256d9ca55f72b *model.safetensors
01562eddd6f9e9ec4bc31656a3b7055284cafbf889acc6c4348dca431ae31f68 *processor_config.json
87a7830d63fcf43bf241c3c5242e96e62dd3fdc29224ca26fed8ea333db72de4 *tokenizer.json
2e31d1126e81bddf8d15c3f95260fb487b48c5131b24fcbb5bb9d2537e7afac0 *tokenizer_config.json
diff --git a/tests/qwen3.5-moe/SHA256SUMS.windows b/tests/qwen3.5-moe/SHA256SUMS.windows
index 8836f913..f92f6088 100644
--- a/tests/qwen3.5-moe/SHA256SUMS.windows
+++ b/tests/qwen3.5-moe/SHA256SUMS.windows
@@ -1,7 +1,7 @@
a92e1dd97cb1cb175c9b70c0828e146bea4371c2643319b661b777e89811972e *chat_template.jinja
b75e911805663da79fb9fbbbcc917b8f1a285d2da54d95c2c63ea7c1ffe9a05a *config.json
2cbd9df0e99570efcced23b8d777bdf1fc692efda54b21eb59ad56ade76c9db6 *generation_config.json
-5f099b32807d0b84ed90765ca0ed53f8771da4738767bc1940486fec954570cf *model.safetensors
+7117e1ccd8c8aaded235c7f9251e45b21b7be4ad045e9b3f2ad743ec3888586d *model.safetensors
0c29f9491e769aabbc389ad5912127cf6d9d5fceda2db8767f73d48131348c81 *processor_config.json
87a7830d63fcf43bf241c3c5242e96e62dd3fdc29224ca26fed8ea333db72de4 *tokenizer.json
4796e48d790a26d65f167bec8fc742beaa71f79f9468a6cd8b3ffa97f6e2a198 *tokenizer_config.json