diff --git a/config.default.toml b/config.default.toml index ebcf96d..c0d6d51 100644 --- a/config.default.toml +++ b/config.default.toml @@ -137,6 +137,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. # Dataset of prompts that tend to not result in refusals (used for calculating residual directions). [good_prompts] diff --git a/src/heretic/config.py b/src/heretic/config.py index a119b24..e3ad23e 100644 --- a/src/heretic/config.py +++ b/src/heretic/config.py @@ -54,6 +54,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/utils.py b/src/heretic/utils.py index 3b4149e..73873eb 100644 --- a/src/heretic/utils.py +++ b/src/heretic/utils.py @@ -208,6 +208,7 @@ def load_prompts( ) dataset = load_dataset( path, + name=specification.config, revision=specification.commit, split=split_str, ) @@ -225,6 +226,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,