fix: prevent infinite loops

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.",
     )

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()