1 Commits

Author SHA1 Message Date
Philipp Emanuel Weidmann 92ab7f09d5 feat: add modifier base class 2026-09-04 19:16:14 +05:30
7 changed files with 80 additions and 20 deletions
+53
View File
@@ -0,0 +1,53 @@
# SPDX-License-Identifier: AGPL-3.0-or-later
# Copyright (C) 2025-2026 Philipp Emanuel Weidmann <pew@worldwidemann.com> + 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`.
"""
+14 -8
View File
@@ -149,12 +149,8 @@ def load_plugin(
class Context: class Context:
""" """
Runtime context passed to plugins Runtime context passed to plugins.
Acts as a quasi-API for plugins to access Heretic functionality.
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.
""" """
def __init__(self, settings: HereticSettings, model: Model) -> None: def __init__(self, settings: HereticSettings, model: Model) -> None:
@@ -180,6 +176,13 @@ class Context:
def get_residuals(self, prompts: list[Prompt]) -> Tensor: def get_residuals(self, prompts: list[Prompt]) -> Tensor:
return self._model.get_residuals_batched(prompts) 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]: def load_prompts(self, specification: DatasetSpecification) -> list[Prompt]:
return load_prompts(self._settings, specification) return load_prompts(self._settings, specification)
@@ -211,8 +214,11 @@ class Plugin:
return False return False
def __init__( 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 # Plugins that declare a settings schema should always receive
# validated plugin settings from the evaluator. # validated plugin settings from the evaluator.
settings_model = self.__class__.get_settings_model() settings_model = self.__class__.get_settings_model()
+2 -2
View File
@@ -32,7 +32,7 @@ class Scorer(Plugin, ABC):
Scorers evaluate model behavior and return a Score. Scorers evaluate model behavior and return a Score.
Example: counting refusals, measuring KL divergence, etc. Examples: Counting refusals, measuring KL divergence, etc.
""" """
@property @property
@@ -47,7 +47,7 @@ class Scorer(Plugin, ABC):
self, self,
heretic_settings: HereticSettings, heretic_settings: HereticSettings,
settings: BaseModel | None = None, settings: BaseModel | None = None,
): ) -> None:
super().__init__(heretic_settings=heretic_settings, settings=settings) super().__init__(heretic_settings=heretic_settings, settings=settings)
@abstractmethod @abstractmethod
+7 -4
View File
@@ -41,9 +41,11 @@ class BenchmarkScore(Scorer):
return self.settings.score_name return self.settings.score_name
def init(self, ctx: Context) -> None: def init(self, ctx: Context) -> None:
model = ctx.get_model()
self.hflm = HFLM( self.hflm = HFLM(
pretrained=ctx._model.model, # ty:ignore[invalid-argument-type] pretrained=model.model, # ty:ignore[invalid-argument-type]
tokenizer=ctx._model.tokenizer, # ty:ignore[invalid-argument-type] tokenizer=model.tokenizer, # ty:ignore[invalid-argument-type]
batch_size="auto", batch_size="auto",
) )
@@ -52,8 +54,9 @@ class BenchmarkScore(Scorer):
# then update its internal model every time we calculate the score, # then update its internal model every time we calculate the score,
# is to get the benefits of batch size caching while allowing for # is to get the benefits of batch size caching while allowing for
# model reloads, e.g. when using --evaluate-model. # model reloads, e.g. when using --evaluate-model.
self.hflm.pretrained = ctx._model.model model = ctx.get_model()
self.hflm._model = ctx._model.model self.hflm.pretrained = model.model
self.hflm._model = model.model
results = lm_eval.simple_evaluate( results = lm_eval.simple_evaluate(
model=self.hflm, model=self.hflm,
-2
View File
@@ -43,8 +43,6 @@ T = TypeVar("T")
print = Console(highlight=False).print print = Console(highlight=False).print
T = TypeVar("T")
def deep_merge_dicts(base: dict[str, Any], override: dict[str, Any]) -> dict[str, Any]: def deep_merge_dicts(base: dict[str, Any], override: dict[str, Any]) -> dict[str, Any]:
""" """