diff --git a/config.default.toml b/config.default.toml index c9e79ef..c3f70c5 100644 --- a/config.default.toml +++ b/config.default.toml @@ -128,6 +128,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. [scorer.KeywordRate] # Name that describes what the configured keyword rate measures. diff --git a/src/heretic/config.py b/src/heretic/config.py index d0a519a..cc17c42 100644 --- a/src/heretic/config.py +++ b/src/heretic/config.py @@ -47,6 +47,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.", diff --git a/src/heretic/main.py b/src/heretic/main.py index 08fb61e..a4b3d5e 100644 --- a/src/heretic/main.py +++ b/src/heretic/main.py @@ -39,12 +39,13 @@ import logging import math import os import random +import re import time import warnings 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 @@ -63,6 +64,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 @@ -488,40 +490,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. ) 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. ) 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) diff --git a/src/heretic/utils.py b/src/heretic/utils.py index a8d3595..08065ca 100644 --- a/src/heretic/utils.py +++ b/src/heretic/utils.py @@ -212,6 +212,7 @@ def load_prompts( ) dataset = load_dataset( path, + name=specification.config, revision=specification.commit, split=split_str, ) @@ -229,6 +230,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, diff --git a/tests/mistral-3/SHA256SUMS.windows b/tests/mistral-3/SHA256SUMS.windows index 0ec00df..2d10d08 100644 --- a/tests/mistral-3/SHA256SUMS.windows +++ b/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 diff --git a/tests/qwen3.5-moe/SHA256SUMS.ci b/tests/qwen3.5-moe/SHA256SUMS.ci index 44c28f5..1fa8e5d 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 +2b3e575ac065f11ae5d4a7c3740efccbed294b646f1645239191ee8393354e03 *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 8836f91..24d015f 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 +6061519a9595326df41abcdd093892463793d4d026d6fd23548f1792f622a252 *model.safetensors 0c29f9491e769aabbc389ad5912127cf6d9d5fceda2db8767f73d48131348c81 *processor_config.json 87a7830d63fcf43bf241c3c5242e96e62dd3fdc29224ca26fed8ea333db72de4 *tokenizer.json 4796e48d790a26d65f167bec8fc742beaa71f79f9468a6cd8b3ffa97f6e2a198 *tokenizer_config.json