vh/heretic
1
0
Fork 0
mirror of https://github.com/p-e-w/heretic.git synced 2026-09-27 22:41:24 -07:00
Code Issues Packages Projects Releases Wiki Activity

feat: support dataset specifications containing multiple individual datasets

Browse Source
This commit is contained in:
Philipp Emanuel Weidmann
2026-09-25 16:43:11 +05:30
parent dc6703788c
commit ffa66af2d4
13 changed files with 235 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
config.default.toml
+17
View File
@@ -102,6 +102,23 @@ max_shard_size = "5GB"
# System prompt to use when prompting the model.
system_prompt = "You are a helpful assistant."
# Dataset of prompts to use for automatically determining the optimal batch size.
[batch_size_test_prompts]
dataset = "mlabonne/harmless_alpaca"
split = "train[:256]"
column = "text"
# Dataset of prompts to use for automatically determining the response prefix.
[[response_prefix_test_prompts]]
dataset = "mlabonne/harmless_alpaca"
split = "train[:100]"
column = "text"
[[response_prefix_test_prompts]]
dataset = "mlabonne/harmful_behaviors"
split = "train[:100]"
column = "text"
# Plugin-specific settings live in top-level TOML tables.
#
# For scorer plugins, use: `[scorer.<ClassName>]` (and optionally `[scorer.<ClassName>_<instance_name>]` for instance-related config).
src/heretic/config.py
+38 -2
View File
@@ -2,7 +2,7 @@
# Copyright (C) 2025-2026 Philipp Emanuel Weidmann <pew@worldwidemann.com> + contributors
from enum import Enum
from typing import Dict, Literal
from typing import Dict, Literal, TypeAlias
from pydantic import (
BaseModel,
@@ -37,7 +37,7 @@ class ExportStrategy(str, Enum):
ADAPTER = "adapter"
class DatasetSpecification(BaseModel):
class SingleDatasetSpecification(BaseModel):
dataset: str = Field(
description="Hugging Face dataset ID, or path to dataset on disk."
)
@@ -81,6 +81,11 @@ class DatasetSpecification(BaseModel):
)
DatasetSpecification: TypeAlias = (
SingleDatasetSpecification | list[SingleDatasetSpecification]
)
class ScorerConfig(BaseModel):
"""
Configuration for a scorer plugin.
@@ -282,6 +287,18 @@ class Settings(BaseSettings):
exclude=True,
)
batch_size_test_prompts: DatasetSpecification = Field(
default=SingleDatasetSpecification(
dataset="mlabonne/harmless_alpaca",
split="train[:256]",
column="text",
),
description="Dataset of prompts to use for automatically determining the optimal batch size.",
# When storing a settings object, the batch size is already fixed,
# either determined by the automatic mechanism or by explicit user choice.
exclude=True,
)
max_response_length: PositiveInt = Field(
default=100,
description="Maximum number of tokens to generate for each response.",
@@ -296,6 +313,25 @@ class Settings(BaseSettings):
),
)
response_prefix_test_prompts: DatasetSpecification = Field(
default=[
SingleDatasetSpecification(
dataset="mlabonne/harmless_alpaca",
split="train[:100]",
column="text",
),
SingleDatasetSpecification(
dataset="mlabonne/harmful_behaviors",
split="train[:100]",
column="text",
),
],
description="Dataset of prompts to use for automatically determining the response prefix.",
# When storing a settings object, the response prefix is already fixed,
# either determined by the automatic mechanism or by explicit user choice.
exclude=True,
)
chain_of_thought_skips: list[tuple[str, str]] = Field(
default=[
# Most thinking models.
src/heretic/main.py
+30 -30
View File
@@ -86,10 +86,12 @@ from .reproduce import (
from .system import empty_cache, get_accelerator_info
from .utils import (
ask_if_unset,
format_dataset_specification,
format_duration,
format_exception,
get_file_sha256,
get_readme_intro,
is_dataset_specification_reproducible,
is_hf_path,
load_prompts,
print,
@@ -412,30 +414,17 @@ def run():
print()
print_memory_usage()
# TODO: Introduce a dedicated dataset setting for test prompts.
good_prompts_dataset = DatasetSpecification(
dataset="mlabonne/harmless_alpaca",
split="train[:5]",
column="text",
)
bad_prompts_dataset = DatasetSpecification(
dataset="mlabonne/harmful_behaviors",
split="train[:5]",
column="text",
)
print()
print(f"Loading good prompts from [bold]{good_prompts_dataset.dataset}[/]...")
good_prompts = load_prompts(settings, good_prompts_dataset)
print(f"* [bold]{len(good_prompts)}[/] prompts loaded")
print()
print(f"Loading bad prompts from [bold]{bad_prompts_dataset.dataset}[/]...")
bad_prompts = load_prompts(settings, bad_prompts_dataset)
print(f"* [bold]{len(bad_prompts)}[/] prompts loaded")
if settings.batch_size == 0:
print()
print(
f"Loading batch size test prompts from [bold]{format_dataset_specification(settings.batch_size_test_prompts)}[/]..."
)
batch_size_test_prompts = load_prompts(
settings,
settings.batch_size_test_prompts,
)
print(f"* [bold]{len(batch_size_test_prompts)}[/] prompts loaded")
print()
print("Determining optimal batch size...")
@@ -446,7 +435,9 @@ def run():
while batch_size <= settings.max_batch_size:
print(f"* Trying batch size [bold]{batch_size}[/]... ", end="")
prompts = good_prompts * math.ceil(batch_size / len(good_prompts))
prompts = batch_size_test_prompts * math.ceil(
batch_size / len(batch_size_test_prompts)
)
prompts = prompts[:batch_size]
try:
@@ -487,9 +478,18 @@ def run():
print(f"* Chosen batch size: [bold]{settings.batch_size}[/]")
if settings.response_prefix is None:
print()
print(
f"Loading response prefix test prompts from [bold]{format_dataset_specification(settings.response_prefix_test_prompts)}[/]..."
)
response_prefix_test_prompts = load_prompts(
settings,
settings.response_prefix_test_prompts,
)
print(f"* [bold]{len(response_prefix_test_prompts)}[/] prompts loaded")
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.
@@ -532,7 +532,7 @@ def run():
# 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)
responses = model.get_responses_batched(response_prefix_test_prompts)
# Despite being located in os.path, commonprefix actually performs
# a naive string operation without any path-specific logic,
@@ -565,7 +565,7 @@ def run():
# 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)
responses = model.get_responses_batched(response_prefix_test_prompts)
additional_prefix = commonprefix(responses).rstrip(" ")
if additional_prefix:
settings.response_prefix += additional_prefix
@@ -1057,15 +1057,14 @@ def run():
# dataset was likely loaded from a local cache), and that
# only built-in plugins are used (external plugins cannot
# be resolved when reproducing).
dataset_specifications = [
dataset_specifications: list[DatasetSpecification] = [
*evaluator.get_dataset_specifications(),
*modifier.get_dataset_specifications(),
]
is_reproducible = (
is_hf_path(settings.model)
and all(
is_hf_path(specification.dataset)
and specification.commit is not None
is_dataset_specification_reproducible(specification)
for specification in dataset_specifications
)
and evaluator.all_scorers_reproducible()
@@ -1190,6 +1189,7 @@ def run():
upload_reproduce_folder(
repo_id,
settings,
dataset_specifications,
token,
checkpoint_path=study_checkpoint_file,
trial=trial,
src/heretic/modifiers/abliteration.py
+6 -6
View File
@@ -19,9 +19,9 @@ from pydantic import (
)
from torch import Tensor
from heretic.config import DatasetSpecification
from heretic.config import DatasetSpecification, SingleDatasetSpecification
from heretic.modifier import Context, Modifier, Serializable
from heretic.utils import print
from heretic.utils import format_dataset_specification, print
@dataclass
@@ -77,7 +77,7 @@ class RowNormalization(str, Enum):
class Settings(BaseModel):
good_prompts: DatasetSpecification = Field(
default=DatasetSpecification(
default=SingleDatasetSpecification(
dataset="mlabonne/harmless_alpaca",
split="train[:400]",
column="text",
@@ -86,7 +86,7 @@ class Settings(BaseModel):
)
bad_prompts: DatasetSpecification = Field(
default=DatasetSpecification(
default=SingleDatasetSpecification(
dataset="mlabonne/harmful_behaviors",
split="train[:400]",
column="text",
@@ -158,14 +158,14 @@ class Abliteration(Modifier[Parameters]):
print()
print(
f"Loading good prompts from [bold]{self.settings.good_prompts.dataset}[/]..."
f"Loading good prompts from [bold]{format_dataset_specification(self.settings.good_prompts)}[/]..."
)
self.good_prompts = ctx.load_prompts(self.settings.good_prompts)
print(f"* [bold]{len(self.good_prompts)}[/] prompts loaded")
print()
print(
f"Loading bad prompts from [bold]{self.settings.bad_prompts.dataset}[/]..."
f"Loading bad prompts from [bold]{format_dataset_specification(self.settings.bad_prompts)}[/]..."
)
self.bad_prompts = ctx.load_prompts(self.settings.bad_prompts)
print(f"* [bold]{len(self.bad_prompts)}[/] prompts loaded")
src/heretic/plugin.py
+6 -2
View File
@@ -13,7 +13,7 @@ from typing import Annotated, Any, TypeVar, Union, get_args, get_origin, get_typ
from pydantic import BaseModel
from torch import Tensor
from .config import DatasetSpecification
from .config import DatasetSpecification, SingleDatasetSpecification
from .config import Settings as HereticSettings
from .model import Model
from .utils import Prompt, deep_merge_dicts, load_prompts
@@ -347,7 +347,11 @@ class Plugin:
return []
specifications = []
for value in dict(self.settings).values():
if isinstance(value, DatasetSpecification):
if isinstance(value, SingleDatasetSpecification) or (
isinstance(value, list)
and len(value) > 0
and isinstance(value[0], SingleDatasetSpecification)
):
specifications.append(value)
return specifications
src/heretic/scorers/keyword_rate.py
+4 -4
View File
@@ -3,9 +3,9 @@
from pydantic import BaseModel, Field
from heretic.config import DatasetSpecification
from heretic.config import DatasetSpecification, SingleDatasetSpecification
from heretic.scorer import Context, Score, Scorer
from heretic.utils import print
from heretic.utils import format_dataset_specification, print
REFUSAL_MARKERS = [
"disclaimer",
@@ -51,7 +51,7 @@ class Settings(BaseModel):
)
prompts: DatasetSpecification = Field(
default=DatasetSpecification(
default=SingleDatasetSpecification(
dataset="mlabonne/harmful_behaviors",
split="test[:100]",
column="text",
@@ -90,7 +90,7 @@ class KeywordRate(Scorer):
def init(self, ctx: Context) -> None:
print()
print(
f"Loading {self.settings.score_name} evaluation prompts from [bold]{self.settings.prompts.dataset}[/]..."
f"Loading {self.settings.score_name} evaluation prompts from [bold]{format_dataset_specification(self.settings.prompts)}[/]..."
)
self.prompts = ctx.load_prompts(self.settings.prompts)
print(f"* [bold]{len(self.prompts)}[/] prompts loaded")
src/heretic/scorers/kl_divergence.py
+4 -4
View File
@@ -4,15 +4,15 @@
import torch.nn.functional as F
from pydantic import BaseModel, Field
from heretic.config import DatasetSpecification
from heretic.config import DatasetSpecification, SingleDatasetSpecification
from heretic.plugin import Context
from heretic.scorer import Score, Scorer
from heretic.utils import print
from heretic.utils import format_dataset_specification, print
class Settings(BaseModel):
prompts: DatasetSpecification = Field(
default=DatasetSpecification(
default=SingleDatasetSpecification(
dataset="mlabonne/harmless_alpaca",
split="test[:100]",
column="text",
@@ -42,7 +42,7 @@ class KLDivergence(Scorer):
def init(self, ctx: Context) -> None:
print()
print(
f"Loading KL divergence evaluation prompts from [bold]{self.settings.prompts.dataset}[/]..."
f"Loading KL divergence evaluation prompts from [bold]{format_dataset_specification(self.settings.prompts)}[/]..."
)
self.prompts = ctx.load_prompts(self.settings.prompts)
print(f"* [bold]{len(self.prompts)}[/] prompts loaded")
src/heretic/utils.py
+70 -4
View File
@@ -30,7 +30,7 @@ from psutil import Process
from questionary import Question
from rich.console import Console
from .config import DatasetSpecification, Settings
from .config import DatasetSpecification, Settings, SingleDatasetSpecification
from .system import (
get_accelerator_info_dict,
get_cpu_info_dict,
@@ -169,9 +169,9 @@ def get_split_slice(split_str: str, length: int) -> tuple[int, int]:
return absolute_instruction.from_, absolute_instruction.to
def load_prompts(
def _load_prompts_single(
settings: Settings,
specification: DatasetSpecification,
specification: SingleDatasetSpecification,
) -> list[Prompt]:
path = specification.dataset
split_str = specification.split
@@ -261,6 +261,44 @@ def load_prompts(
]
def load_prompts(
settings: Settings,
specification: DatasetSpecification,
) -> list[Prompt]:
if isinstance(specification, SingleDatasetSpecification):
return _load_prompts_single(settings, specification)
else:
return [
prompt
for single_specification in specification
for prompt in _load_prompts_single(settings, single_specification)
]
def format_dataset_specification(specification: DatasetSpecification) -> str:
if isinstance(specification, SingleDatasetSpecification):
return specification.dataset
else:
return (
"\\["
+ ", ".join(
single_specification.dataset for single_specification in specification
)
+ "]"
)
def is_dataset_specification_reproducible(specification: DatasetSpecification) -> bool:
if isinstance(specification, SingleDatasetSpecification):
return is_hf_path(specification.dataset) and specification.commit is not None
else:
return all(
is_hf_path(single_specification.dataset)
and single_specification.commit is not None
for single_specification in specification
)
def batchify(items: list[T], batch_size: int) -> list[list[T]]:
return [items[i : i + batch_size] for i in range(0, len(items), batch_size)]
@@ -367,6 +405,7 @@ def format_hf_link(
def generate_reproduce_readme(
settings: Settings,
dataset_specifications: list[DatasetSpecification],
checkpoint_filename: str,
trial: Trial | FrozenTrial,
include_system_information: bool,
@@ -483,6 +522,29 @@ def generate_reproduce_readme(
f" --index-url https://download.pytorch.org/whl/{suffix}"
)
formatted_datasets = set()
for specification in dataset_specifications:
if isinstance(specification, SingleDatasetSpecification):
formatted_datasets.add(
format_hf_link(
specification.dataset,
specification.commit,
is_dataset=True,
)
)
else:
for single_specification in specification:
formatted_datasets.add(
format_hf_link(
single_specification.dataset,
single_specification.commit,
is_dataset=True,
)
)
dataset_lines = "\n".join(
f"- {formatted_dataset}" for formatted_dataset in sorted(formatted_datasets)
)
trial_scores = trial.user_attrs["scores"]
score_lines = "\n".join(
(
@@ -502,7 +564,7 @@ This directory contains the necessary information and assets to reproduce the re
## Datasets
- TODO: Collect all datasets from scorers and modifiers.
{dataset_lines}
## Selected trial
@@ -619,6 +681,7 @@ def get_file_sha256(file_path: str | Path) -> str:
def create_reproduce_folder(
path: Path,
settings: Settings,
dataset_specifications: list[DatasetSpecification],
checkpoint_path: str | Path,
trial: Trial | FrozenTrial,
uploaded_model_hashes: dict[str, str],
@@ -667,6 +730,7 @@ def create_reproduce_folder(
(reproduce_dir / "README.md").write_text(
generate_reproduce_readme(
settings,
dataset_specifications,
checkpoint_filename,
trial,
include_system_information=include_system_information,
@@ -683,6 +747,7 @@ def create_reproduce_folder(
def upload_reproduce_folder(
repo_id: str,
settings: Settings,
dataset_specifications: list[DatasetSpecification],
token: str,
checkpoint_path: str | Path,
trial: Trial | FrozenTrial,
@@ -711,6 +776,7 @@ def upload_reproduce_folder(
create_reproduce_folder(
tmp_path,
settings,
dataset_specifications,
checkpoint_path=checkpoint_path,
trial=trial,
uploaded_model_hashes=uploaded_model_hashes,
tests/gemma-4e/config.toml
+12
View File
@@ -18,6 +18,18 @@ trial_index = 0
model_action = "save"
save_directory = "model"
[[response_prefix_test_prompts]]
dataset = "mlabonne/harmless_alpaca"
commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f"
split = "train[:5]"
column = "text"
[[response_prefix_test_prompts]]
dataset = "mlabonne/harmful_behaviors"
commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7"
split = "train[:5]"
column = "text"
[scorer.KLDivergence.prompts]
dataset = "mlabonne/harmless_alpaca"
commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f"
tests/minicpm5/config.toml
+12
View File
@@ -18,6 +18,18 @@ trial_index = 0
model_action = "save"
save_directory = "model"
[[response_prefix_test_prompts]]
dataset = "mlabonne/harmless_alpaca"
commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f"
split = "train[:5]"
column = "text"
[[response_prefix_test_prompts]]
dataset = "mlabonne/harmful_behaviors"
commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7"
split = "train[:5]"
column = "text"
[scorer.KLDivergence.prompts]
dataset = "mlabonne/harmless_alpaca"
commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f"
tests/mistral-3/config.toml
+12
View File
@@ -18,6 +18,18 @@ trial_index = 0
model_action = "save"
save_directory = "model"
[[response_prefix_test_prompts]]
dataset = "mlabonne/harmless_alpaca"
commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f"
split = "train[:5]"
column = "text"
[[response_prefix_test_prompts]]
dataset = "mlabonne/harmful_behaviors"
commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7"
split = "train[:5]"
column = "text"
[scorer.KLDivergence.prompts]
dataset = "mlabonne/harmless_alpaca"
commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f"
tests/qwen2.5/config.toml
+12
View File
@@ -18,6 +18,18 @@ trial_index = 0
model_action = "save"
save_directory = "model"
[[response_prefix_test_prompts]]
dataset = "mlabonne/harmless_alpaca"
commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f"
split = "train[:5]"
column = "text"
[[response_prefix_test_prompts]]
dataset = "mlabonne/harmful_behaviors"
commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7"
split = "train[:5]"
column = "text"
[scorer.KLDivergence.prompts]
dataset = "mlabonne/harmless_alpaca"
commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f"
tests/qwen3.5-moe/config.toml
+12
View File
@@ -18,6 +18,18 @@ trial_index = 0
model_action = "save"
save_directory = "model"
[[response_prefix_test_prompts]]
dataset = "mlabonne/harmless_alpaca"
commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f"
split = "train[:5]"
column = "text"
[[response_prefix_test_prompts]]
dataset = "mlabonne/harmful_behaviors"
commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7"
split = "train[:5]"
column = "text"
[scorer.KLDivergence.prompts]
dataset = "mlabonne/harmless_alpaca"
commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f"