fix: Check response prefix.

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. <think>) by using a dummy prompt.
+        # If found, then we extract the end of it (e.g. </think>) 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 "<think>\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)