diff --git a/config.default.toml b/config.default.toml index ebcf96d..7bc876d 100644 --- a/config.default.toml +++ b/config.default.toml @@ -91,8 +91,8 @@ residual_plot_style = "dark_background" # { plugin = , optimization = , instance_name = } # where is one of "minimize", "maximize", "none" (do not optimize) scorers = [ - { plugin = "heretic.scorers.keyword_rate.KeywordRate", optimization = "minimize"}, - { plugin = "heretic.scorers.kl_divergence.KLDivergence", optimization = "minimize"}, + { plugin = "heretic.scorers.keyword_rate.KeywordRate", optimization = "minimize" }, + { plugin = "heretic.scorers.kl_divergence.KLDivergence", optimization = "minimize" }, ] # Whether to adjust the residual directions so that only the component that is diff --git a/config.piqa.toml b/config.piqa.toml index a903a3a..76f08cb 100644 --- a/config.piqa.toml +++ b/config.piqa.toml @@ -2,6 +2,6 @@ # that you run Heretic from, and edit the configuration to your liking. scorers = [ - { plugin = "heretic.scorers.keyword_rate.KeywordRate", optimization = "minimize"}, - { plugin = "heretic.scorers.benchmark_score.BenchmarkScore", optimization = "maximize"}, + { plugin = "heretic.scorers.keyword_rate.KeywordRate", optimization = "minimize" }, + { plugin = "heretic.scorers.benchmark_score.BenchmarkScore", optimization = "maximize" }, ] diff --git a/src/heretic/modifier.py b/src/heretic/modifier.py new file mode 100644 index 0000000..a82c276 --- /dev/null +++ b/src/heretic/modifier.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: AGPL-3.0-or-later +# Copyright (C) 2025-2026 Philipp Emanuel Weidmann + contributors + +from abc import ABC, abstractmethod +from typing import Generic, TypeVar + +from optuna import Trial +from pydantic import BaseModel + +from heretic.plugin import Context, Plugin + +from .config import Settings as HereticSettings + +Parameters = TypeVar("Parameters") + + +class Modifier(Generic[Parameters], Plugin, ABC): + """ + Abstract base class for modifier plugins. + + Modifiers modify models based on an implementation-dependent set of optimizable parameters. + + Examples: Standard abliteration, ARA, SOMA, etc. + """ + + def __init__( + self, + heretic_settings: HereticSettings, + settings: BaseModel | None = None, + ) -> None: + super().__init__(heretic_settings=heretic_settings, settings=settings) + + @abstractmethod + def suggest_parameters(self, ctx: Context, trial: Trial) -> Parameters: + """ + Sample parameters for a trial using the trial's `suggest_*` methods, + collect them in an implementation-dependent parameters object, and + return that object. + """ + + @abstractmethod + def modify_model(self, ctx: Context, parameters: Parameters) -> None: + """ + Modify the model (obtainable via `ctx.get_model()`) + according to the provided parameters. + """ + + @abstractmethod + def reset_model(self, ctx: Context) -> None: + """ + Reset the model (obtainable via `ctx.get_model()`), + undoing any changes made by `modify_model`. + """ diff --git a/src/heretic/plugin.py b/src/heretic/plugin.py index 4b57fbf..0c62845 100644 --- a/src/heretic/plugin.py +++ b/src/heretic/plugin.py @@ -149,12 +149,8 @@ def load_plugin( class Context: """ - Runtime context passed to plugins - - Provides plugin-safe access to the model. - - Plugins must use `get_responses(...)`, `get_logits(...)`, etc. - Direct access to the underlying Model is intentionally not exposed. + Runtime context passed to plugins. + Acts as a quasi-API for plugins to access Heretic functionality. """ def __init__(self, settings: HereticSettings, model: Model) -> None: @@ -180,6 +176,13 @@ class Context: def get_residuals(self, prompts: list[Prompt]) -> Tensor: return self._model.get_residuals_batched(prompts) + def get_model(self) -> Model: + """ + Prefer managed methods (`get_responses` etc.) unless you + actually need access to the model object. + """ + return self._model + def load_prompts(self, specification: DatasetSpecification) -> list[Prompt]: return load_prompts(self._settings, specification) @@ -211,8 +214,11 @@ class Plugin: return False def __init__( - self, *, heretic_settings: HereticSettings, settings: BaseModel | None = None - ): + self, + *, + heretic_settings: HereticSettings, + settings: BaseModel | None = None, + ) -> None: # Plugins that declare a settings schema should always receive # validated plugin settings from the evaluator. settings_model = self.__class__.get_settings_model() diff --git a/src/heretic/scorer.py b/src/heretic/scorer.py index e61a309..b433ec4 100644 --- a/src/heretic/scorer.py +++ b/src/heretic/scorer.py @@ -32,7 +32,7 @@ class Scorer(Plugin, ABC): Scorers evaluate model behavior and return a Score. - Example: counting refusals, measuring KL divergence, etc. + Examples: Counting refusals, measuring KL divergence, etc. """ @property @@ -47,7 +47,7 @@ class Scorer(Plugin, ABC): self, heretic_settings: HereticSettings, settings: BaseModel | None = None, - ): + ) -> None: super().__init__(heretic_settings=heretic_settings, settings=settings) @abstractmethod diff --git a/src/heretic/scorers/benchmark_score.py b/src/heretic/scorers/benchmark_score.py index 9bc04b1..15aaf8c 100644 --- a/src/heretic/scorers/benchmark_score.py +++ b/src/heretic/scorers/benchmark_score.py @@ -41,9 +41,11 @@ class BenchmarkScore(Scorer): return self.settings.score_name def init(self, ctx: Context) -> None: + model = ctx.get_model() + self.hflm = HFLM( - pretrained=ctx._model.model, # ty:ignore[invalid-argument-type] - tokenizer=ctx._model.tokenizer, # ty:ignore[invalid-argument-type] + pretrained=model.model, # ty:ignore[invalid-argument-type] + tokenizer=model.tokenizer, # ty:ignore[invalid-argument-type] batch_size="auto", ) @@ -52,8 +54,9 @@ class BenchmarkScore(Scorer): # then update its internal model every time we calculate the score, # is to get the benefits of batch size caching while allowing for # model reloads, e.g. when using --evaluate-model. - self.hflm.pretrained = ctx._model.model - self.hflm._model = ctx._model.model + model = ctx.get_model() + self.hflm.pretrained = model.model + self.hflm._model = model.model results = lm_eval.simple_evaluate( model=self.hflm, diff --git a/src/heretic/utils.py b/src/heretic/utils.py index 3b4149e..5107394 100644 --- a/src/heretic/utils.py +++ b/src/heretic/utils.py @@ -43,8 +43,6 @@ T = TypeVar("T") print = Console(highlight=False).print -T = TypeVar("T") - def deep_merge_dicts(base: dict[str, Any], override: dict[str, Any]) -> dict[str, Any]: """