diff --git a/src/heretic/config.py b/src/heretic/config.py index d91e5cb..17710d0 100644 --- a/src/heretic/config.py +++ b/src/heretic/config.py @@ -4,7 +4,12 @@ from enum import Enum from typing import Dict -from pydantic import BaseModel, Field +from pydantic import ( + BaseModel, + Field, + NonNegativeInt, + PositiveInt, +) from pydantic_settings import ( BaseSettings, CliSettingsSource, @@ -181,12 +186,12 @@ class Settings(BaseSettings): ), ) - batch_size: int = Field( + batch_size: NonNegativeInt = Field( default=0, # auto description="Number of input sequences to process in parallel (0 = auto).", ) - max_batch_size: int = Field( + max_batch_size: PositiveInt = Field( default=128, description="Maximum batch size to try when automatically determining the optimal batch size.", # When storing a settings object, the batch size is already fixed, @@ -194,7 +199,7 @@ class Settings(BaseSettings): exclude=True, ) - max_response_length: int = Field( + max_response_length: PositiveInt = Field( default=100, description="Maximum number of tokens to generate for each response.", ) @@ -311,7 +316,7 @@ class Settings(BaseSettings): ), ) - full_normalization_lora_rank: int = Field( + full_normalization_lora_rank: PositiveInt = Field( default=3, description=( 'The rank of the LoRA adapter to use when "full" row normalization is used. ' @@ -332,12 +337,12 @@ class Settings(BaseSettings): ), ) - n_trials: int = Field( + n_trials: PositiveInt = Field( default=200, description="Number of abliteration trials to run during optimization.", ) - n_startup_trials: int = Field( + n_startup_trials: NonNegativeInt = Field( default=60, description="Number of trials that use random sampling for the purpose of exploration.", ) @@ -418,7 +423,7 @@ class Settings(BaseSettings): exclude=True, ) - max_shard_size: int | str = Field( + max_shard_size: PositiveInt | str = Field( default="5GB", description="Maximum size for individual safetensors files generated when exporting a model.", ) @@ -433,12 +438,12 @@ class Settings(BaseSettings): description='Action to take in case a checkpoint exists: "continue", "restart", or unset to prompt the user.', ) - trial_index: int | None = Field( + trial_index: NonNegativeInt | None = Field( default=None, description="Index (in the sorted Pareto front) of the trial to use, or unset to prompt the user.", ) - n_additional_trials: int | None = Field( + n_additional_trials: PositiveInt | None = Field( default=None, description="Number of additional trials to run, or unset to prompt the user.", ) diff --git a/src/heretic/main.py b/src/heretic/main.py index ceec04e..243841b 100644 --- a/src/heretic/main.py +++ b/src/heretic/main.py @@ -702,7 +702,9 @@ def run(): if len(study.trials) == settings.n_trials: study.set_user_attr("finished", True) - while True: + trial_loop_active = True + + while trial_loop_active: if not reproduction_mode: # If no trials at all have been evaluated, the study must have been stopped # by pressing Ctrl+C while the first trial was running. In this case, we just @@ -772,7 +774,11 @@ def run(): ) ) - while True: + while trial_loop_active: + # Ensure a predefined trial is only processed once. + if settings.trial_index is not None: + trial_loop_active = False + if reproduction_mode: parameters = reproduction_information["parameters"] metrics = reproduction_information["metrics"] @@ -873,7 +879,13 @@ def run(): reset_trial_model() - while True: + action_loop_active = True + + while action_loop_active: + # Ensure a predefined action is only executed once. + if settings.model_action is not None: + action_loop_active = False + if settings.model_action is None: print()