mirror of
https://github.com/p-e-w/heretic.git
synced 2026-09-15 16:45:54 -07:00
fix: save only essential settings
This commit is contained in:
@@ -1,6 +1,7 @@
|
||||
# SPDX-License-Identifier: AGPL-3.0-or-later
|
||||
# Copyright (C) 2025-2026 Philipp Emanuel Weidmann <pew@worldwidemann.com> + contributors
|
||||
|
||||
import copy
|
||||
from enum import Enum
|
||||
from typing import Dict
|
||||
|
||||
@@ -13,6 +14,12 @@ from pydantic_settings import (
|
||||
TomlConfigSettingsSource,
|
||||
)
|
||||
|
||||
# !!!IMPORTANT!!!
|
||||
#
|
||||
# Any settings added to the classes defined in this module
|
||||
# must be evaluated for privacy implications and added to
|
||||
# the logic in get_essential_settings if required.
|
||||
|
||||
|
||||
class QuantizationMethod(str, Enum):
|
||||
NONE = "none"
|
||||
@@ -460,3 +467,47 @@ class Settings(BaseSettings):
|
||||
file_secret_settings,
|
||||
TomlConfigSettingsSource(settings_cls, toml_file="config.toml"),
|
||||
)
|
||||
|
||||
|
||||
def get_essential_settings(settings: Settings) -> Settings:
|
||||
"""
|
||||
Returns a stripped-down version of the settings object that only contains
|
||||
settings that directly influence the results of the abliteration run.
|
||||
In particular, this object contains no file system paths other than (possibly)
|
||||
paths to local models and datasets.
|
||||
"""
|
||||
|
||||
essential_settings = copy.deepcopy(settings)
|
||||
|
||||
del essential_settings.evaluate_model
|
||||
|
||||
# We always use the default for security reasons.
|
||||
del essential_settings.trust_remote_code
|
||||
|
||||
# When storing a settings object, the batch size is already fixed,
|
||||
# either determined by the automatic mechanism or by explicit user choice.
|
||||
del essential_settings.max_batch_size
|
||||
|
||||
# When storing a settings object, the response prefix is already fixed,
|
||||
# either determined by the automatic mechanism or by explicit user choice.
|
||||
del essential_settings.chain_of_thought_skips
|
||||
|
||||
del essential_settings.print_responses
|
||||
del essential_settings.print_residual_geometry
|
||||
del essential_settings.plot_residuals
|
||||
del essential_settings.residual_plot_path
|
||||
del essential_settings.residual_plot_title
|
||||
del essential_settings.residual_plot_style
|
||||
del essential_settings.study_checkpoint_dir
|
||||
del essential_settings.benchmarks
|
||||
|
||||
for dataset in [
|
||||
essential_settings.good_prompts,
|
||||
essential_settings.bad_prompts,
|
||||
essential_settings.good_evaluation_prompts,
|
||||
essential_settings.bad_evaluation_prompts,
|
||||
]:
|
||||
del dataset.residual_plot_label
|
||||
del dataset.residual_plot_color
|
||||
|
||||
return essential_settings
|
||||
|
||||
+9
-3
@@ -45,7 +45,7 @@ from rich.table import Table
|
||||
from rich.traceback import install
|
||||
|
||||
from .analyzer import Analyzer
|
||||
from .config import QuantizationMethod, Settings
|
||||
from .config import QuantizationMethod, Settings, get_essential_settings
|
||||
from .evaluator import Evaluator
|
||||
from .model import AbliterationParameters, Model, get_model_class
|
||||
from .system import empty_cache, get_accelerator_info
|
||||
@@ -570,7 +570,10 @@ def run():
|
||||
load_if_exists=True,
|
||||
)
|
||||
|
||||
study.set_user_attr("settings", settings.model_dump_json())
|
||||
study.set_user_attr(
|
||||
"settings",
|
||||
get_essential_settings(settings).model_dump_json(exclude_none=True),
|
||||
)
|
||||
study.set_user_attr("finished", False)
|
||||
|
||||
def count_completed_trials() -> int:
|
||||
@@ -684,7 +687,10 @@ def run():
|
||||
continue
|
||||
|
||||
settings.n_trials += n_additional_trials
|
||||
study.set_user_attr("settings", settings.model_dump_json())
|
||||
study.set_user_attr(
|
||||
"settings",
|
||||
get_essential_settings(settings).model_dump_json(exclude_none=True),
|
||||
)
|
||||
study.set_user_attr("finished", False)
|
||||
|
||||
try:
|
||||
|
||||
@@ -26,7 +26,7 @@ from psutil import Process
|
||||
from questionary import Choice, Style
|
||||
from rich.console import Console
|
||||
|
||||
from .config import DatasetSpecification, Settings
|
||||
from .config import DatasetSpecification, Settings, get_essential_settings
|
||||
from .system import (
|
||||
get_accelerator_info_dict,
|
||||
get_cpu_info_dict,
|
||||
@@ -312,7 +312,7 @@ def get_readme_intro(settings: Settings, trial: Trial) -> str:
|
||||
def generate_config_toml(settings: Settings) -> str:
|
||||
"""Serializes the full Settings object to TOML."""
|
||||
|
||||
return tomli_w.dumps(settings.model_dump(exclude_none=True))
|
||||
return tomli_w.dumps(get_essential_settings(settings).model_dump(exclude_none=True))
|
||||
|
||||
|
||||
def generate_requirements_txt() -> str:
|
||||
@@ -551,7 +551,7 @@ def generate_reproduce_json(
|
||||
"pytorch_version": torch.__version__,
|
||||
"requirements": get_requirements_dict(),
|
||||
},
|
||||
"settings": settings.model_dump(exclude_none=True),
|
||||
"settings": get_essential_settings(settings).model_dump(exclude_none=True),
|
||||
"parameters": {
|
||||
"direction_index": trial.user_attrs["direction_index"],
|
||||
"abliteration_parameters": trial.user_attrs["parameters"],
|
||||
|
||||
Reference in New Issue
Block a user