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