feat: headless operation + end-to-end tests (#392)

* fix: remove notebook input shims

Closes #280

* feat: support headless operation (no interactive input)

* fix: prevent infinite loops

* feat: add end-to-end tests

* ci: run tests in CI

* ci: fix test output ordering

* fix: replace home-cooked `set_seed` function with Transformers builtin

* feat: print PyTorch config when running tests

* feat: print additional information

* experiment: try to standardize test environment

* fix: revert environment changes

* feat: support multiple valid hashes for each output file

* feat: add test output hashes for CI

* feat: add test output hashes for CI (alternative environment)

* feat: add hashes for Windows (#394)

* fix: Hash on windows

* trigger ci

* fix: prefer .yaml (used widely than .toml for model configs)

* use removeprefix

* docs: restore commet

* use removeprefix again

* tests: Add windows hash files for all test models

* trigger ci

* fix: minor cleanup

* clean merge mismatch

* remove unnecessary CRLF replace, now that we support more SUMS files

* fix: use binary mode for hashes everywhere

---------

Co-authored-by: Vinay Umrethe <umrethevinay@gmail.com>

.github/workflows/ci.yml

@@ -40,6 +40,11 @@ jobs:
       - name: Check typing
         run: uv run ty check --output-format=github --error-on-warning .
 
+      - name: Run tests
+        env:
+          PYTHONUNBUFFERED: "1"
+        run: uv run tests/run_tests.py 2>&1
+
       - name: Build package
         run: uv build
 

.gitignore

@@ -15,11 +15,14 @@ wheels/
 # Editors
 /.vscode/
 
-# Configuration files
+# Configuration file (root only, not ignored in test directories)
 /config.toml
 
 # Study checkpoints
-/checkpoints/
+checkpoints/
 
 # Residual plots
-/plots/
+plots/
+
+# Models generated by tests
+/tests/*/model/

README.md

@@ -86,7 +86,7 @@ models with Heretic.
 Prepare a Python 3.10+ environment with PyTorch 2.2+ installed as appropriate
 for your hardware. Then run:
 
-```
+```sh
 pip install -U heretic-llm
 heretic Qwen/Qwen3-4B-Instruct-2507
 ```
@@ -134,7 +134,7 @@ provides features designed to support research into the semantics of model inter
 (interpretability). To use those features, you need to install Heretic with the
 optional `research` extra:
 
-```
+```sh
 pip install -U heretic-llm[research]
 ```
 

config.default.toml

@@ -71,6 +71,9 @@ chain_of_thought_skips = [
 # Whether to print prompt/response pairs when counting refusals.
 print_responses = false
 
+# Whether to print additional information that can help with debugging.
+print_debug_information = false
+
 # Whether to print detailed information about residuals and refusal directions.
 print_residual_geometry = false
 

pyproject.toml

@@ -38,6 +38,8 @@ dependencies = [
     "questionary~=2.1",
     "rich~=14.3",
     "tomli-w~=1.2",
+    "torch",  # version deliberately unspecified
+    "torchvision",  # version deliberately unspecified
     "tqdm~=4.67",
     "transformers[kernels]~=5.6",
 ]

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.",
     )
@@ -247,6 +252,12 @@ class Settings(BaseSettings):
         exclude=True,
     )
 
+    print_debug_information: bool = Field(
+        default=False,
+        description="Whether to print additional information that can help with debugging.",
+        exclude=True,
+    )
+
     print_residual_geometry: bool = Field(
         default=False,
         description="Whether to print detailed information about residuals and refusal directions.",
@@ -311,7 +322,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 +343,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,14 +429,61 @@ class Settings(BaseSettings):
         exclude=True,
     )
 
+    max_shard_size: PositiveInt | str = Field(
+        default="5GB",
+        description="Maximum size for individual safetensors files generated when exporting a model.",
+    )
+
     export_strategy: ExportStrategy | None = Field(
         default=None,
         description='How to export the model: "merge", "adapter", or unset to prompt the user.',
     )
 
-    max_shard_size: int | str = Field(
-        default="5GB",
-        description="Maximum size for individual safetensors files generated when exporting a model.",
+    checkpoint_action: str | None = Field(
+        default=None,
+        description='Action to take in case a checkpoint exists: "continue", "restart", or unset to prompt the user.',
+    )
+
+    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: PositiveInt | None = Field(
+        default=None,
+        description="Number of additional trials to run, or unset to prompt the user.",
+    )
+
+    model_action: str | None = Field(
+        default=None,
+        description='Action to take with the decensored model: "save", "upload", or unset to prompt the user.',
+    )
+
+    save_directory: str | None = Field(
+        default=None,
+        description="Directory to save the model to, or unset to prompt the user.",
+        exclude=True,
+    )
+
+    upload_repo_id: str | None = Field(
+        default=None,
+        description="Name of the Hugging Face repository to upload the model to, or unset to prompt the user.",
+        exclude=True,
+    )
+
+    upload_repo_private: bool | None = Field(
+        default=None,
+        description="Whether the Hugging Face repository to upload the model to should be private, or unset to prompt the user.",
+    )
+
+    upload_reproducibility_information: str | None = Field(
+        default=None,
+        description='Which reproducibility information to add to the Hugging Face repository: "full", "basic", "none", or unset to prompt the user.',
+    )
+
+    ignore_mismatches: bool | None = Field(
+        default=None,
+        description="Whether to attempt to reproduce the model even if there are environment mismatches, or unset to prompt the user.",
     )
 
     refusal_markers: list[str] = Field(

src/heretic/main.py

@@ -80,6 +80,7 @@ from .reproduce import (
 )
 from .system import empty_cache, get_accelerator_info
 from .utils import (
+    ask_if_unset,
     format_duration,
     format_exception,
     get_file_sha256,
@@ -89,11 +90,6 @@ from .utils import (
     load_prompts,
     print,
     print_memory_usage,
-    prompt_password,
-    prompt_path,
-    prompt_select,
-    prompt_text,
-    set_seed,
     upload_reproduce_folder,
 )
 
@@ -108,10 +104,10 @@ def obtain_export_strategy(
     Returns an export strategy, or None if cancelled.
     """
 
-    if settings.export_strategy is not None:
-        return settings.export_strategy
-
-    if settings.quantization == QuantizationMethod.BNB_4BIT:
+    if (
+        settings.quantization == QuantizationMethod.BNB_4BIT
+        and settings.export_strategy is None
+    ):
         print()
         print(
             "The model was loaded with quantization. Merging requires reloading the base model."
@@ -155,27 +151,29 @@ def obtain_export_strategy(
 
         print()
 
-    strategy = prompt_select(
-        "How do you want to export the model?",
-        choices=[
-            Choice(
-                title="Merge the abliteration LoRA and export the full model"
-                + (
-                    ""
-                    if settings.quantization == QuantizationMethod.NONE
-                    else " (requires sufficient RAM)"
+    return ask_if_unset(
+        settings.export_strategy,
+        questionary.select(
+            "How do you want to export the model?",
+            choices=[
+                Choice(
+                    title="Merge the abliteration LoRA and export the full model"
+                    + (
+                        ""
+                        if settings.quantization == QuantizationMethod.NONE
+                        else " (requires sufficient RAM)"
+                    ),
+                    value=ExportStrategy.MERGE,
                 ),
-                value=ExportStrategy.MERGE,
-            ),
-            Choice(
-                title="Export the abliteration LoRA only (can be merged later)",
-                value=ExportStrategy.ADAPTER,
-            ),
-        ],
+                Choice(
+                    title="Export the abliteration LoRA only (can be merged later)",
+                    value=ExportStrategy.ADAPTER,
+                ),
+            ],
+            style=Style([("highlighted", "reverse")]),
+        ),
     )
 
-    return strategy
-
 
 def run():
     # Enable expandable segments to reduce memory fragmentation on multi-GPU setups.
@@ -254,7 +252,7 @@ def run():
             )
             return
 
-        if not check_environment(reproduction_information):
+        if not check_environment(settings, reproduction_information):
             return
 
         print()
@@ -266,10 +264,22 @@ def run():
     if settings.seed is None:
         settings.seed = random.randint(0, 2**32 - 1)
 
-    set_seed(settings.seed)
+    transformers.set_seed(settings.seed)
 
     print(get_accelerator_info())
 
+    if settings.print_debug_information:
+        print()
+        print(torch.__config__.show().strip())
+        print()
+        print(
+            f"torch.backends.mkldnn.enabled = [bold]{torch.backends.mkldnn.enabled}[/]"
+        )
+        print(f"torch.get_num_threads() = [bold]{torch.get_num_threads()}[/]")
+        print(
+            f"torch.get_num_interop_threads() = [bold]{torch.get_num_interop_threads()}[/]"
+        )
+
     # We don't need gradients as we only do inference.
     torch.set_grad_enabled(False)
 
@@ -320,15 +330,17 @@ def run():
         choices = []
 
         if existing_study.user_attrs["finished"]:
-            print()
-            print(
-                (
-                    "[green]You have already processed this model.[/] "
-                    "You can show the results from the previous run, allowing you to export models or to run additional trials. "
-                    "Alternatively, you can ignore the previous run and start from scratch. "
-                    "This will delete the checkpoint file and all results from the previous run."
+            if settings.checkpoint_action is None:
+                print()
+                print(
+                    (
+                        "[green]You have already processed this model.[/] "
+                        "You can show the results from the previous run, allowing you to export models or to run additional trials. "
+                        "Alternatively, you can ignore the previous run and start from scratch. "
+                        "This will delete the checkpoint file and all results from the previous run."
+                    )
                 )
-            )
+
             choices.append(
                 Choice(
                     title="Show the results from the previous run",
@@ -336,15 +348,17 @@ def run():
                 )
             )
         else:
-            print()
-            print(
-                (
-                    "[yellow]You have already processed this model, but the run was interrupted.[/] "
-                    "You can continue the previous run from where it stopped. This will override any specified settings. "
-                    "Alternatively, you can ignore the previous run and start from scratch. "
-                    "This will delete the checkpoint file and all results from the previous run."
+            if settings.checkpoint_action is None:
+                print()
+                print(
+                    (
+                        "[yellow]You have already processed this model, but the run was interrupted.[/] "
+                        "You can continue the previous run from where it stopped. This will override any specified settings. "
+                        "Alternatively, you can ignore the previous run and start from scratch. "
+                        "This will delete the checkpoint file and all results from the previous run."
+                    )
                 )
-            )
+
             choices.append(
                 Choice(
                     title="Continue the previous run",
@@ -366,19 +380,29 @@ def run():
             )
         )
 
-        print()
-        choice = prompt_select("How would you like to proceed?", choices)
+        if settings.checkpoint_action is None:
+            print()
 
-        if choice == "continue":
+        action = ask_if_unset(
+            settings.checkpoint_action,
+            questionary.select(
+                "How would you like to proceed?",
+                choices=choices,
+                style=Style([("highlighted", "reverse")]),
+            ),
+        )
+
+        if action is None or action == "":
+            return
+
+        if action == "continue":
             settings = Settings.model_validate_json(
                 existing_study.user_attrs["settings"]
             )
-        elif choice == "restart":
+        elif action == "restart":
             os.unlink(study_checkpoint_file)
             backend = JournalFileBackend(study_checkpoint_file, lock_obj=lock_obj)
             storage = JournalStorage(backend)
-        elif choice is None or choice == "":
-            return
 
     model = Model(settings)
     print()
@@ -619,7 +643,7 @@ def run():
             min_weight_distance = trial.suggest_float(
                 f"{component}.min_weight_distance",
                 1.0,
-                0.6 * last_layer_index,
+                max(0.6 * last_layer_index, 1.0),
             )
 
             parameters[component] = AbliterationParameters(
@@ -709,7 +733,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
@@ -766,18 +792,24 @@ def run():
 
             print()
             print("[bold green]Optimization finished![/]")
-            print()
-            print(
-                (
-                    "The following trials resulted in Pareto optimal combinations of refusals and KL divergence. "
-                    "After selecting a trial, you will be able to save the model, upload it to Hugging Face, "
-                    "chat with it to test how well it works, or run standard benchmarks on it. "
-                    "You can return to this menu later to select a different trial. "
-                    "[yellow]Note that KL divergence values above 0.5 usually indicate significant damage to the original model's capabilities.[/]"
-                )
-            )
 
-        while True:
+            if settings.trial_index is None:
+                print()
+                print(
+                    (
+                        "The following trials resulted in Pareto optimal combinations of refusals and KL divergence. "
+                        "After selecting a trial, you will be able to save the model, upload it to Hugging Face, "
+                        "chat with it to test how well it works, or run standard benchmarks on it. "
+                        "You can return to this menu later to select a different trial. "
+                        "[yellow]Note that KL divergence values above 0.5 usually indicate significant damage to the original model's capabilities.[/]"
+                    )
+                )
+
+        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"]
@@ -797,8 +829,19 @@ def run():
                 print()
                 print("Restoring model from reproduction information...")
             else:
-                print()
-                trial = prompt_select("Which trial do you want to use?", choices)
+                if settings.trial_index is None:
+                    print()
+
+                trial = ask_if_unset(
+                    None
+                    if settings.trial_index is None
+                    else best_trials[settings.trial_index],
+                    questionary.select(
+                        "Which trial do you want to use?",
+                        choices=choices,
+                        style=Style([("highlighted", "reverse")]),
+                    ),
+                )
 
                 if trial is None or trial == "":
                     return
@@ -806,8 +849,11 @@ def run():
                 if trial == "continue":
                     while True:
                         try:
-                            n_additional_trials = prompt_text(
-                                "How many additional trials do you want to run?"
+                            n_additional_trials = ask_if_unset(
+                                settings.n_additional_trials,
+                                questionary.text(
+                                    "How many additional trials do you want to run?"
+                                ),
                             )
                             if n_additional_trials is None or n_additional_trials == "":
                                 n_additional_trials = 0
@@ -866,22 +912,46 @@ def run():
 
             reset_trial_model()
 
-            while True:
-                print()
-                action = prompt_select(
-                    "What do you want to do with the decensored model?",
-                    [
-                        "Save the model to a local folder",
-                        "Upload the model to Hugging Face",
-                        "Chat with the model",
-                        "Benchmark the model",
-                        Choice(
-                            title="Exit program"
-                            if reproduction_mode
-                            else "Return to the trial selection menu",
-                            value="",
-                        ),
-                    ],
+            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()
+
+                action = ask_if_unset(
+                    settings.model_action,
+                    questionary.select(
+                        "What do you want to do with the decensored model?",
+                        choices=[
+                            Choice(
+                                title="Save the model to a local folder",
+                                value="save",
+                            ),
+                            Choice(
+                                title="Upload the model to Hugging Face",
+                                value="upload",
+                            ),
+                            Choice(
+                                title="Chat with the model",
+                                value="chat",
+                            ),
+                            Choice(
+                                title="Benchmark the model",
+                                value="benchmark",
+                            ),
+                            Choice(
+                                title="Exit program"
+                                if reproduction_mode
+                                else "Return to the trial selection menu",
+                                value="",
+                            ),
+                        ],
+                        style=Style([("highlighted", "reverse")]),
+                    ),
                 )
 
                 if action is None or action == "":
@@ -895,8 +965,14 @@ def run():
                 # the optimized model.
                 try:
                     match action:
-                        case "Save the model to a local folder":
-                            save_directory = prompt_path("Path to the folder:")
+                        case "save":
+                            save_directory = ask_if_unset(
+                                settings.save_directory,
+                                questionary.path(
+                                    "Path to the folder:",
+                                    only_directories=True,
+                                ),
+                            )
                             if not save_directory:
                                 continue
 
@@ -951,13 +1027,20 @@ def run():
                                             f"[bold]{filename}:[/] [red]File not found[/]"
                                         )
 
-                        case "Upload the model to Hugging Face":
+                        case "upload":
                             # We don't use huggingface_hub.login() because that stores the token on disk,
                             # and since this program will often be run on rented or shared GPU servers,
                             # it's better to not persist credentials.
                             token = huggingface_hub.get_token()
                             if not token:
-                                token = prompt_password("Hugging Face access token:")
+                                # NOTE: Unlike for most other values obtained from interactive inputs, it is
+                                #       not possible to set the token via the settings. This is a security
+                                #       precaution to prevent exporting the token under all circumstances.
+                                #       For scripting, the correct way to set the token is through the HF_TOKEN
+                                #       environment variable, or through the HF token file.
+                                token = questionary.password(
+                                    "Hugging Face access token:"
+                                ).ask()
                             if not token:
                                 continue
 
@@ -969,17 +1052,32 @@ def run():
                             email = user.get("email", "no email found")
                             print(f"Logged in as [bold]{fullname} ({email})[/]")
 
-                            repo_id = prompt_text(
-                                "Name of repository:",
-                                default=f"{user['name']}/{Path(settings.model).name}-heretic",
+                            repo_id = ask_if_unset(
+                                settings.upload_repo_id,
+                                questionary.text(
+                                    "Name of repository:",
+                                    default=f"{user['name']}/{Path(settings.model).name}-heretic",
+                                ),
                             )
+                            if not repo_id:
+                                continue
 
-                            visibility = prompt_select(
-                                "Should the repository be public or private?",
-                                [
-                                    "Public",
-                                    "Private",
-                                ],
+                            visibility = ask_if_unset(
+                                None
+                                if settings.upload_repo_private is None
+                                else (
+                                    "Private"
+                                    if settings.upload_repo_private
+                                    else "Public"
+                                ),
+                                questionary.select(
+                                    "Should the repository be public or private?",
+                                    choices=[
+                                        "Public",
+                                        "Private",
+                                    ],
+                                    style=Style([("highlighted", "reverse")]),
+                                ),
                             )
                             if visibility is None:
                                 continue
@@ -1004,31 +1102,37 @@ def run():
                             )
 
                             if is_reproducible:
-                                print(
-                                    (
-                                        "Heretic can add information to the repository that allows others to reproduce the model. "
-                                        "This is optional, but valuable to the community as both a learning tool and to preserve computational work already done. "
-                                        "Guaranteeing reproducibility requires basic system information (Python and OS version, CPU and GPU/accelerator info) "
-                                        "as tensor operations can give different results in different system environments. "
-                                        "[bold]The information does not include any file system paths or other private data.[/]"
+                                if settings.upload_reproducibility_information is None:
+                                    print(
+                                        (
+                                            "Heretic can add information to the repository that allows others to reproduce the model. "
+                                            "This is optional, but valuable to the community as both a learning tool and to preserve computational work already done. "
+                                            "Guaranteeing reproducibility requires basic system information (Python and OS version, CPU and GPU/accelerator info) "
+                                            "as tensor operations can give different results in different system environments. "
+                                            "[bold]The information does not include any file system paths or other private data.[/]"
+                                        )
                                     )
-                                )
-                                reproducibility_information = prompt_select(
-                                    "Which reproducibility information do you want to add?",
-                                    [
-                                        Choice(
-                                            title="Full: Settings, package versions, and system information",
-                                            value="full",
-                                        ),
-                                        Choice(
-                                            title="Basic: Settings and package versions",
-                                            value="basic",
-                                        ),
-                                        Choice(
-                                            title="Don't add any reproducibility information",
-                                            value="none",
-                                        ),
-                                    ],
+
+                                reproducibility_information = ask_if_unset(
+                                    settings.upload_reproducibility_information,
+                                    questionary.select(
+                                        "Which reproducibility information do you want to add?",
+                                        choices=[
+                                            Choice(
+                                                title="Full: Settings, package versions, and system information",
+                                                value="full",
+                                            ),
+                                            Choice(
+                                                title="Basic: Settings and package versions",
+                                                value="basic",
+                                            ),
+                                            Choice(
+                                                title="Don't add any reproducibility information",
+                                                value="none",
+                                            ),
+                                        ],
+                                        style=Style([("highlighted", "reverse")]),
+                                    ),
                                 )
                                 if reproducibility_information is None:
                                     continue
@@ -1174,7 +1278,7 @@ def run():
                                             f"[bold]{filename}:[/] [red]File not found[/]"
                                         )
 
-                        case "Chat with the model":
+                        case "chat":
                             print()
                             print(
                                 "[cyan]Press Ctrl+C at any time to return to the menu.[/]"
@@ -1186,11 +1290,10 @@ def run():
 
                             while True:
                                 try:
-                                    message = prompt_text(
+                                    message = questionary.text(
                                         "User:",
                                         qmark=">",
-                                        unsafe=True,
-                                    )
+                                    ).unsafe_ask()
                                     if not message:
                                         break
                                     chat.append({"role": "user", "content": message})
@@ -1204,7 +1307,7 @@ def run():
                                     # Ctrl+C/Ctrl+D
                                     break
 
-                        case "Benchmark the model":
+                        case "benchmark":
                             benchmarks = questionary.checkbox(
                                 "Which benchmarks do you want to run?",
                                 [
@@ -1219,16 +1322,17 @@ def run():
                             if not benchmarks:
                                 continue
 
-                            scope = prompt_select(
+                            scope = questionary.select(
                                 (
                                     "Do you want to benchmark the original model along with the decensored model? "
                                     "Benchmarking both models allows you to compare the scores, but it takes twice as much time."
                                 ),
-                                [
+                                choices=[
                                     "Benchmark only the decensored model",
                                     "Benchmark both models",
                                 ],
-                            )
+                                style=Style([("highlighted", "reverse")]),
+                            ).ask()
                             if scope is None:
                                 continue
                             benchmark_original_model = scope == "Benchmark both models"

src/heretic/model.py

@@ -586,11 +586,16 @@ class Model:
                         W = W - W_org
                         # Use a low-rank SVD to get an approximation of the matrix.
                         r = self.peft_config.r
+
                         # svd_lowrank is randomized:
                         # https://github.com/pytorch/pytorch/blob/20919052303c0b5ba87f8bf7e19237dc33ab09d3/torch/_lowrank.py#L108-L109
                         # Reseed immediately before the call so restoring a trial is independent of RNG history.
                         torch.manual_seed(self.settings.seed)
+                        # "It's safe to call this function if CUDA is not available;
+                        # in that case, it is silently ignored."
+                        torch.cuda.manual_seed_all(self.settings.seed)  # ty:ignore[invalid-argument-type]
                         U, S, Vh = torch.svd_lowrank(W, q=2 * r + 4, niter=6)
+
                         # Truncate it to the part we want to store in the LoRA adapter.
                         # Note: svd_lowrank actually returns V, so transpose it to get Vh.
                         U = U[:, :r]

src/heretic/reproduce.py

@@ -12,6 +12,7 @@ from typing import Any, cast
 from urllib.request import urlopen
 
 import cpuinfo
+import questionary
 import torch
 from huggingface_hub import HfApi, hf_hub_download
 from huggingface_hub.utils import (
@@ -19,15 +20,16 @@ from huggingface_hub.utils import (
     disable_progress_bars,
     enable_progress_bars,
 )
-from questionary import Choice
+from questionary import Choice, Style
 from rich.table import Table
 
+from .config import Settings
 from .system import (
     get_accelerator_info_dict,
     get_heretic_version_info,
     get_requirements_dict,
 )
-from .utils import print, prompt_select
+from .utils import ask_if_unset, print
 
 
 def collect_reproducibles(path: str):
@@ -192,7 +194,10 @@ def format_version_information(version_information: dict[str, Any]) -> str:
         return f"{version}-unknown-{random.randint(2**16, 2**17)}"
 
 
-def check_environment(reproduction_information: dict[str, Any]) -> bool:
+def check_environment(
+    settings: Settings,
+    reproduction_information: dict[str, Any],
+) -> bool | None:
     mismatch_severity: MismatchSeverity | None = None
 
     system_mismatches = []
@@ -361,22 +366,26 @@ def check_environment(reproduction_information: dict[str, Any]) -> bool:
             )
         )
 
-        print()
-        choice = prompt_select(
-            "How would you like to proceed?",
-            [
-                Choice(
-                    title="Attempt to reproduce the model anyway",
-                    value=True,
-                ),
-                Choice(
-                    title="Exit program",
-                    value=False,
-                ),
-            ],
-        )
+        if settings.ignore_mismatches is None:
+            print()
 
-        return choice
+        return ask_if_unset(
+            settings.ignore_mismatches,
+            questionary.select(
+                "How would you like to proceed?",
+                choices=[
+                    Choice(
+                        title="Attempt to reproduce the model anyway",
+                        value=True,
+                    ),
+                    Choice(
+                        title="Exit program",
+                        value=False,
+                    ),
+                ],
+                style=Style([("highlighted", "reverse")]),
+            ),
+        )
     else:
         # There are no mismatches at all, so there is nothing to confirm.
         return True

src/heretic/utils.py

@@ -1,23 +1,19 @@
 # SPDX-License-Identifier: AGPL-3.0-or-later
 # Copyright (C) 2025-2026  Philipp Emanuel Weidmann <pew@worldwidemann.com> + contributors
 
-import getpass
 import hashlib
 import json
 import os
 import platform
-import random
 import tempfile
 import traceback
 from dataclasses import dataclass
 from datetime import datetime, timezone
 from importlib.metadata import version
 from pathlib import Path
-from typing import Any, TypeVar
+from typing import TypeVar
 
 import huggingface_hub
-import numpy as np
-import questionary
 import tomli_w
 import torch
 from datasets import DatasetDict, ReadInstruction, load_dataset, load_from_disk
@@ -28,7 +24,7 @@ from huggingface_hub.utils import validate_repo_id
 from optuna import Trial
 from optuna.trial import FrozenTrial
 from psutil import Process
-from questionary import Choice, Style
+from questionary import Question
 from rich.console import Console
 
 from .config import DatasetSpecification, Settings
@@ -41,6 +37,9 @@ from .system import (
     is_xpu_available,
 )
 
+T = TypeVar("T")
+
+
 print = Console(highlight=False).print
 
 
@@ -67,99 +66,6 @@ def print_memory_usage():
         p("Driver (reserved) MPS memory", torch.mps.driver_allocated_memory())
 
 
-def is_notebook() -> bool:
-    # Check for specific environment variables (Colab, Kaggle).
-    # This is necessary because when running as a subprocess (e.g. !heretic),
-    # get_ipython() might not be available or might not reflect the notebook environment.
-    if os.getenv("COLAB_GPU") or os.getenv("KAGGLE_KERNEL_RUN_TYPE"):
-        return True
-
-    # Check IPython shell type (for library usage).
-    try:
-        from IPython import get_ipython  # ty:ignore[unresolved-import]
-
-        shell = get_ipython()
-        if shell is None:
-            return False
-
-        shell_name = shell.__class__.__name__
-        if shell_name in ["ZMQInteractiveShell", "Shell"]:
-            return True
-
-        if "google.colab" in str(shell.__class__):
-            return True
-
-        return False
-    except (ImportError, NameError, AttributeError):
-        return False
-
-
-def prompt_select(message: str, choices: list[Any]) -> Any:
-    if is_notebook():
-        print()
-        print(message)
-        real_choices = []
-
-        for i, choice in enumerate(choices, 1):
-            if isinstance(choice, Choice):
-                print(f"[{i}] {choice.title}")
-                real_choices.append(choice.value)
-            else:
-                print(f"[{i}] {choice}")
-                real_choices.append(choice)
-
-        while True:
-            try:
-                selection = input("Enter number: ")
-                index = int(selection) - 1
-                if 0 <= index < len(real_choices):
-                    return real_choices[index]
-                print(
-                    f"[red]Please enter a number between 1 and {len(real_choices)}[/]"
-                )
-            except ValueError:
-                print("[red]Invalid input. Please enter a number.[/]")
-    else:
-        return questionary.select(
-            message,
-            choices=choices,
-            style=Style([("highlighted", "reverse")]),
-        ).ask()
-
-
-def prompt_text(
-    message: str,
-    default: str = "",
-    qmark: str = "?",
-    unsafe: bool = False,
-) -> str:
-    if is_notebook():
-        print()
-        result = input(f"{message} [{default}]: " if default else f"{message}: ")
-        return result if result else default
-    else:
-        question = questionary.text(message, default=default, qmark=qmark)
-        if unsafe:
-            return question.unsafe_ask()
-        else:
-            return question.ask()
-
-
-def prompt_path(message: str) -> str:
-    if is_notebook():
-        return prompt_text(message)
-    else:
-        return questionary.path(message, only_directories=True).ask()
-
-
-def prompt_password(message: str) -> str:
-    if is_notebook():
-        print()
-        return getpass.getpass(message)
-    else:
-        return questionary.password(message).ask()
-
-
 def format_duration(seconds: float) -> str:
     seconds = round(seconds)
     hours, seconds = divmod(seconds, 3600)
@@ -186,6 +92,16 @@ def format_exception(error: Exception) -> str:
     return traceback.format_exc().strip()
 
 
+def ask_if_unset(value: T, question: Question, unsafe: bool = False) -> T:
+    if value is None:
+        if unsafe:
+            return question.unsafe_ask()
+        else:
+            return question.ask()
+    else:
+        return value
+
+
 def is_hf_path(path: str) -> bool:
     """Checks whether a path likely refers to a Hugging Face repository."""
 
@@ -297,9 +213,6 @@ def load_prompts(
     ]
 
 
-T = TypeVar("T")
-
-
 def batchify(items: list[T], batch_size: int) -> list[list[T]]:
     return [items[i : i + batch_size] for i in range(0, len(items), batch_size)]
 
@@ -386,14 +299,6 @@ def generate_requirements_txt() -> str:
     return "\n".join(requirements) + "\n"
 
 
-def set_seed(seed: int):
-    """Sets the seed for all RNGs."""
-
-    random.seed(seed)
-    np.random.seed(seed)
-    torch.manual_seed(seed)
-
-
 def format_hf_link(
     path: str,
     commit: str | None = None,

tests/README.md

@@ -0,0 +1,17 @@
+Run the tests with
+
+```sh
+uv run run_tests.py
+```
+
+To update the hashes after a logic change, run the tests, then execute
+
+```sh
+cd TEST_DIR/model
+sha256sum -b * > ../SHA256SUMS.LABEL
+```
+
+where `LABEL` describes the type of system you are running the tests on.
+Since PyTorch does not guarantee exact cross-system reproducibility regardless of configuration,
+multiple valid hashes can be provided for each output file. The above update must be performed
+for each `TEST_DIR` and on each type of system.

tests/gemma-4e/SHA256SUMS.ci

@@ -0,0 +1,7 @@
+2f1b4d75d067bae3fe44e676721c7f077d243bc007156cb9c2f8b5836613d082 *chat_template.jinja
+ca80080dfa4ec6ba87152fa2b9afe70b90c400e5c4b1d6bdc3aa3114467ca68f *config.json
+70070bac883cf9c39b5992450d6b23cd160eaf33099e24c654e0359d2f87c760 *generation_config.json
+f3f4ec19504f182486459cf4e255ece265c25f827840d63b6a9d4058b8e4877a *model.safetensors
+32bdf45d2ad4cc29a0822ddd157a182de76644f0419a6228d151495256e9813c *processor_config.json
+cc8d3a0ce36466ccc1278bf987df5f71db1719b9ca6b4118264f45cb627bfe0f *tokenizer.json
+a1bab8c81ed15fa6ce912ec993c66cb49392e0487fb1ea5f5f11ea3618683627 *tokenizer_config.json

tests/gemma-4e/SHA256SUMS.ci2

@@ -0,0 +1,7 @@
+2f1b4d75d067bae3fe44e676721c7f077d243bc007156cb9c2f8b5836613d082 *chat_template.jinja
+ca80080dfa4ec6ba87152fa2b9afe70b90c400e5c4b1d6bdc3aa3114467ca68f *config.json
+70070bac883cf9c39b5992450d6b23cd160eaf33099e24c654e0359d2f87c760 *generation_config.json
+53c4ee891dce23c0ac85bebc2c4d48301469750fafbb3e6e024c15786d94db8b *model.safetensors
+32bdf45d2ad4cc29a0822ddd157a182de76644f0419a6228d151495256e9813c *processor_config.json
+cc8d3a0ce36466ccc1278bf987df5f71db1719b9ca6b4118264f45cb627bfe0f *tokenizer.json
+a1bab8c81ed15fa6ce912ec993c66cb49392e0487fb1ea5f5f11ea3618683627 *tokenizer_config.json

tests/gemma-4e/SHA256SUMS.linux

@@ -0,0 +1,7 @@
+2f1b4d75d067bae3fe44e676721c7f077d243bc007156cb9c2f8b5836613d082 *chat_template.jinja
+ca80080dfa4ec6ba87152fa2b9afe70b90c400e5c4b1d6bdc3aa3114467ca68f *config.json
+70070bac883cf9c39b5992450d6b23cd160eaf33099e24c654e0359d2f87c760 *generation_config.json
+effe36925f85ecb1e29bba84501a456bb49df21e4047be8b7ea3f6f88181fb65 *model.safetensors
+32bdf45d2ad4cc29a0822ddd157a182de76644f0419a6228d151495256e9813c *processor_config.json
+cc8d3a0ce36466ccc1278bf987df5f71db1719b9ca6b4118264f45cb627bfe0f *tokenizer.json
+a1bab8c81ed15fa6ce912ec993c66cb49392e0487fb1ea5f5f11ea3618683627 *tokenizer_config.json

tests/gemma-4e/SHA256SUMS.windows

@@ -0,0 +1,7 @@
+b16d3228a775c549ba97af41233a54e9de8dd2b65250f78346661d18b936a8b5 *chat_template.jinja
+0094ad598a8043f84d82ad5c886547bca1d1d7f302d82f1491f83d388e89acd4 *config.json
+1a019c5d688d54cf01318eab88cb4345dfa52135eb1d83c2f54125469eb88d5c *generation_config.json
+effe36925f85ecb1e29bba84501a456bb49df21e4047be8b7ea3f6f88181fb65 *model.safetensors
+24d00232e58cfa179fe8b3911c788d4aad9a6279d778ebe4c72e82623b6197f9 *processor_config.json
+cc8d3a0ce36466ccc1278bf987df5f71db1719b9ca6b4118264f45cb627bfe0f *tokenizer.json
+8044bbbddaee8dc47e6b5660e013ba92224d4a5392b2939c59699aa0105f5c8b *tokenizer_config.json

tests/gemma-4e/config.toml

@@ -0,0 +1,41 @@
+model = "tiny-random/gemma-4e"
+model_commit = "3a207ada2c2cd95e9671942e84cf47ea58f0f6af"
+
+seed = 12345
+print_debug_information = true
+
+batch_size = 2
+max_response_length = 10
+kl_divergence_target = 0
+n_trials = 2
+n_startup_trials = 1
+
+export_strategy = "merge"
+checkpoint_action = "restart"
+trial_index = 0
+model_action = "save"
+save_directory = "model"
+
+[good_prompts]
+dataset = "mlabonne/harmless_alpaca"
+commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f"
+split = "train[:5]"
+column = "text"
+
+[bad_prompts]
+dataset = "mlabonne/harmful_behaviors"
+commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7"
+split = "train[:5]"
+column = "text"
+
+[good_evaluation_prompts]
+dataset = "mlabonne/harmless_alpaca"
+commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f"
+split = "test[:5]"
+column = "text"
+
+[bad_evaluation_prompts]
+dataset = "mlabonne/harmful_behaviors"
+commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7"
+split = "test[:5]"
+column = "text"

tests/mistral-3/SHA256SUMS.ci

@@ -0,0 +1,7 @@
+39f03c383413f531fd302c06c7e982ad98c83f0657a8339ae25478ccb81fdcda *chat_template.jinja
+f69f84977a47c8fea9ce9fc26b7de379216cb01146ea726a87996d3554cfcd19 *config.json
+34dfa6012ca9ac5f57e5521d8dbaecbc7ab7f7ab0fd96ec020b543aab5f265d9 *generation_config.json
+876c6691eb85e3e5e11771e589529830fb454ab26344e1271ae550661e312b50 *model.safetensors
+84be30b124b50749c56d25fdbec5ccedf564446f6b3b035e88e1e07b986d2491 *processor_config.json
+c3a8d92e371b92a2cd6e678e31ebc27d0235e929a51fbf290f74742b341fa96f *tokenizer.json
+7b29c843c0043622d28fd4638451cbb0a609d99a0762ffbff3b92b4b2fee4d94 *tokenizer_config.json

tests/mistral-3/SHA256SUMS.ci2

@@ -0,0 +1,7 @@
+39f03c383413f531fd302c06c7e982ad98c83f0657a8339ae25478ccb81fdcda *chat_template.jinja
+f69f84977a47c8fea9ce9fc26b7de379216cb01146ea726a87996d3554cfcd19 *config.json
+34dfa6012ca9ac5f57e5521d8dbaecbc7ab7f7ab0fd96ec020b543aab5f265d9 *generation_config.json
+6febb813086f253e5ec0fcda02fdfc849c551a7dba54681b37ac5bc402e4eed6 *model.safetensors
+84be30b124b50749c56d25fdbec5ccedf564446f6b3b035e88e1e07b986d2491 *processor_config.json
+c3a8d92e371b92a2cd6e678e31ebc27d0235e929a51fbf290f74742b341fa96f *tokenizer.json
+7b29c843c0043622d28fd4638451cbb0a609d99a0762ffbff3b92b4b2fee4d94 *tokenizer_config.json

tests/mistral-3/SHA256SUMS.linux

@@ -0,0 +1,7 @@
+39f03c383413f531fd302c06c7e982ad98c83f0657a8339ae25478ccb81fdcda *chat_template.jinja
+f69f84977a47c8fea9ce9fc26b7de379216cb01146ea726a87996d3554cfcd19 *config.json
+34dfa6012ca9ac5f57e5521d8dbaecbc7ab7f7ab0fd96ec020b543aab5f265d9 *generation_config.json
+29aff97d5633dead9e1ccd29a2cc153b4b7431d22f63c8d6cf60bc6547681cc9 *model.safetensors
+84be30b124b50749c56d25fdbec5ccedf564446f6b3b035e88e1e07b986d2491 *processor_config.json
+c3a8d92e371b92a2cd6e678e31ebc27d0235e929a51fbf290f74742b341fa96f *tokenizer.json
+7b29c843c0043622d28fd4638451cbb0a609d99a0762ffbff3b92b4b2fee4d94 *tokenizer_config.json

tests/mistral-3/SHA256SUMS.windows

@@ -0,0 +1,7 @@
+72f84af4ea36b82409c35e31b584361534305ef7c0d90fce20d0dc38a7efead8 *chat_template.jinja
+e4c5278b361c57621253c27a2c3db358e1580aec8a14be8e19d4420a224137cf *config.json
+8dde85c000ae807be907421465826c7c63a39f6acf6d04a5a84efaf116ed4ef7 *generation_config.json
+29aff97d5633dead9e1ccd29a2cc153b4b7431d22f63c8d6cf60bc6547681cc9 *model.safetensors
+20e7a6dcde0a6f60ea3b4fb08f6f7afa62532dda93a3111e28384ba5150575f9 *processor_config.json
+c3a8d92e371b92a2cd6e678e31ebc27d0235e929a51fbf290f74742b341fa96f *tokenizer.json
+60a8042e29b4b20e884e48375aa1b9ac0025547371d50e60f6d55e6a9675e868 *tokenizer_config.json

tests/mistral-3/config.toml

@@ -0,0 +1,41 @@
+model = "tiny-random/mistral-3"
+model_commit = "931aa2e5c9668fc3679e56aa44972fe18597d55d"
+
+seed = 12345
+print_debug_information = true
+
+batch_size = 2
+max_response_length = 10
+kl_divergence_target = 0
+n_trials = 2
+n_startup_trials = 1
+
+export_strategy = "merge"
+checkpoint_action = "restart"
+trial_index = 0
+model_action = "save"
+save_directory = "model"
+
+[good_prompts]
+dataset = "mlabonne/harmless_alpaca"
+commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f"
+split = "train[:5]"
+column = "text"
+
+[bad_prompts]
+dataset = "mlabonne/harmful_behaviors"
+commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7"
+split = "train[:5]"
+column = "text"
+
+[good_evaluation_prompts]
+dataset = "mlabonne/harmless_alpaca"
+commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f"
+split = "test[:5]"
+column = "text"
+
+[bad_evaluation_prompts]
+dataset = "mlabonne/harmful_behaviors"
+commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7"
+split = "test[:5]"
+column = "text"

tests/qwen3.5-moe/SHA256SUMS.ci

@@ -0,0 +1,7 @@
+a4aee8afcf2e0711942cf848899be66016f8d14a889ff9ede07bca099c28f715 *chat_template.jinja
+749b56d1b1e08081981169db6f2c44ab0be4fd6ebb452d15baafa5e09c21586a *config.json
+4625d1d64d41d1fa9dae7af4ba1e1d7e65a194073d4efa58acb266a916eaaa74 *generation_config.json
+5fb94c65bcd9d736735a45e50c2b0bfafd3bb09a444c49b8cff2e131ed35797e *model.safetensors
+01562eddd6f9e9ec4bc31656a3b7055284cafbf889acc6c4348dca431ae31f68 *processor_config.json
+87a7830d63fcf43bf241c3c5242e96e62dd3fdc29224ca26fed8ea333db72de4 *tokenizer.json
+2e31d1126e81bddf8d15c3f95260fb487b48c5131b24fcbb5bb9d2537e7afac0 *tokenizer_config.json

tests/qwen3.5-moe/SHA256SUMS.linux

@@ -0,0 +1,7 @@
+a4aee8afcf2e0711942cf848899be66016f8d14a889ff9ede07bca099c28f715 *chat_template.jinja
+749b56d1b1e08081981169db6f2c44ab0be4fd6ebb452d15baafa5e09c21586a *config.json
+4625d1d64d41d1fa9dae7af4ba1e1d7e65a194073d4efa58acb266a916eaaa74 *generation_config.json
+5e0fb0ac724cf079b693fc76a515e60bc16de72c32b36c107b9f078061c4f2ef *model.safetensors
+01562eddd6f9e9ec4bc31656a3b7055284cafbf889acc6c4348dca431ae31f68 *processor_config.json
+87a7830d63fcf43bf241c3c5242e96e62dd3fdc29224ca26fed8ea333db72de4 *tokenizer.json
+2e31d1126e81bddf8d15c3f95260fb487b48c5131b24fcbb5bb9d2537e7afac0 *tokenizer_config.json

tests/qwen3.5-moe/SHA256SUMS.windows

@@ -0,0 +1,7 @@
+a92e1dd97cb1cb175c9b70c0828e146bea4371c2643319b661b777e89811972e *chat_template.jinja
+b75e911805663da79fb9fbbbcc917b8f1a285d2da54d95c2c63ea7c1ffe9a05a *config.json
+2cbd9df0e99570efcced23b8d777bdf1fc692efda54b21eb59ad56ade76c9db6 *generation_config.json
+5f099b32807d0b84ed90765ca0ed53f8771da4738767bc1940486fec954570cf *model.safetensors
+0c29f9491e769aabbc389ad5912127cf6d9d5fceda2db8767f73d48131348c81 *processor_config.json
+87a7830d63fcf43bf241c3c5242e96e62dd3fdc29224ca26fed8ea333db72de4 *tokenizer.json
+4796e48d790a26d65f167bec8fc742beaa71f79f9468a6cd8b3ffa97f6e2a198 *tokenizer_config.json

tests/qwen3.5-moe/config.toml

@@ -0,0 +1,41 @@
+model = "tiny-random/qwen3.5-moe"
+model_commit = "2ebfa8d9717238c5dda927008104fa172a149050"
+
+seed = 12345
+print_debug_information = true
+
+batch_size = 2
+max_response_length = 10
+kl_divergence_target = 0
+n_trials = 2
+n_startup_trials = 1
+
+export_strategy = "merge"
+checkpoint_action = "restart"
+trial_index = 0
+model_action = "save"
+save_directory = "model"
+
+[good_prompts]
+dataset = "mlabonne/harmless_alpaca"
+commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f"
+split = "train[:5]"
+column = "text"
+
+[bad_prompts]
+dataset = "mlabonne/harmful_behaviors"
+commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7"
+split = "train[:5]"
+column = "text"
+
+[good_evaluation_prompts]
+dataset = "mlabonne/harmless_alpaca"
+commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f"
+split = "test[:5]"
+column = "text"
+
+[bad_evaluation_prompts]
+dataset = "mlabonne/harmful_behaviors"
+commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7"
+split = "test[:5]"
+column = "text"

tests/run_tests.py

@@ -0,0 +1,87 @@
+# SPDX-License-Identifier: AGPL-3.0-or-later
+# Copyright (C) 2025-2026  Philipp Emanuel Weidmann <pew@worldwidemann.com> + contributors
+
+import hashlib
+import subprocess
+import sys
+from pathlib import Path
+
+
+# TODO: Replace this with hashlib.file_digest when we drop support for Python 3.10.
+def get_file_sha256(file_path: str | Path) -> str:
+    hash = hashlib.sha256()
+
+    with open(file_path, "rb") as file:
+        # Read the file in 64 kB blocks.
+        for block in iter(lambda: file.read(65536), b""):
+            hash.update(block)
+
+    return hash.hexdigest()
+
+
+script_directory = Path(__file__).resolve().parent
+
+project_directory = script_directory.parent
+
+tests_failed = False
+
+for test_directory in script_directory.iterdir():
+    if test_directory.is_dir():
+        config_file = test_directory / "config.toml"
+        hash_files = list(test_directory.glob("SHA256SUMS.*"))
+
+        if config_file.is_file() and hash_files:
+            print("#" * 50)
+            print(f"Running test {test_directory.name}")
+            print("#" * 50)
+            print()
+
+            subprocess.run(
+                [
+                    "uv",
+                    "run",
+                    "--project",
+                    project_directory,
+                    "--directory",
+                    test_directory,
+                    "heretic",
+                ],
+                check=True,
+            )
+
+            print()
+
+            valid_hashes: dict[str, list[str]] = {}
+
+            for hash_file in hash_files:
+                with open(hash_file, "r", encoding="utf-8") as file:
+                    for line in file:
+                        if line.strip():
+                            sha256, filename = line.split()
+                            filename = filename.removeprefix("*")
+
+                            if filename not in valid_hashes:
+                                valid_hashes[filename] = []
+
+                            valid_hashes[filename].append(sha256.lower())
+
+            for filename in valid_hashes:
+                sha256 = get_file_sha256(test_directory / "model" / filename)
+
+                if sha256.lower() not in valid_hashes[filename]:
+                    print(
+                        (
+                            f"Test {test_directory.name} has FAILED!\n"
+                            f"Output file {filename} doesn't match any valid hash.\n\n"
+                            f"Valid hashes:\n"
+                            f"{chr(10).join(valid_hashes[filename])}\n\n"
+                            f"Actual hash:\n"
+                            f"{sha256}\n"
+                        )
+                    )
+                    tests_failed = True
+
+if tests_failed:
+    sys.exit("Tests failed.")
+else:
+    print("All tests passed.")

uv.lock

@@ -968,6 +968,8 @@ dependencies = [
     { name = "questionary" },
     { name = "rich" },
     { name = "tomli-w" },
+    { name = "torch" },
+    { name = "torchvision" },
     { name = "tqdm" },
     { name = "transformers", extra = ["kernels"] },
 ]
@@ -1011,6 +1013,8 @@ requires-dist = [
     { name = "rich", specifier = "~=14.3" },
     { name = "scikit-learn", marker = "extra == 'research'", specifier = "~=1.7" },
     { name = "tomli-w", specifier = "~=1.2" },
+    { name = "torch" },
+    { name = "torchvision" },
     { name = "tqdm", specifier = "~=4.67" },
     { name = "transformers", extras = ["kernels"], specifier = "~=5.6" },
 ]
@@ -3738,6 +3742,47 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/db/2b/f7818f6ec88758dfd21da46b6cd46af9d1b3433e53ddbb19ad1e0da17f9b/torch-2.9.1-cp314-cp314t-win_amd64.whl", hash = "sha256:c88d3299ddeb2b35dcc31753305612db485ab6f1823e37fb29451c8b2732b87e", size = 111163659, upload-time = "2025-11-12T15:23:20.009Z" },
 ]
 
+[[package]]
+name = "torchvision"
+version = "0.24.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
+    { name = "numpy", version = "2.3.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" },
+    { name = "pillow" },
+    { name = "torch" },
+]
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/f7/09/d51aadf8591138e08b74c64a6eb783630c7a31ca2634416277115a9c3a2b/torchvision-0.24.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:ded5e625788572e4e1c4d155d1bbc48805c113794100d70e19c76e39e4d53465", size = 1891441, upload-time = "2025-11-12T15:25:01.687Z" },
+    { url = "https://files.pythonhosted.org/packages/6b/49/a35df863e7c153aad82af7505abd8264a5b510306689712ef86bea862822/torchvision-0.24.1-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:54ed17c3d30e718e08d8da3fd5b30ea44b0311317e55647cb97077a29ecbc25b", size = 2386226, upload-time = "2025-11-12T15:25:05.449Z" },
+    { url = "https://files.pythonhosted.org/packages/49/20/f2d7cd1eea052887c1083afff0b8df5228ec93b53e03759f20b1a3c6d22a/torchvision-0.24.1-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:f476da4e085b7307aaab6f540219617d46d5926aeda24be33e1359771c83778f", size = 8046093, upload-time = "2025-11-12T15:25:09.425Z" },
+    { url = "https://files.pythonhosted.org/packages/d8/cf/0ff4007c09903199307da5f53a192ff5d62b45447069e9ef3a19bdc5ff12/torchvision-0.24.1-cp310-cp310-win_amd64.whl", hash = "sha256:fbdbdae5e540b868a681240b7dbd6473986c862445ee8a138680a6a97d6c34ff", size = 3696202, upload-time = "2025-11-12T15:25:10.657Z" },
+    { url = "https://files.pythonhosted.org/packages/e7/69/30f5f03752aa1a7c23931d2519b31e557f3f10af5089d787cddf3b903ecf/torchvision-0.24.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:056c525dc875f18fe8e9c27079ada166a7b2755cea5a2199b0bc7f1f8364e600", size = 1891436, upload-time = "2025-11-12T15:25:04.3Z" },
+    { url = "https://files.pythonhosted.org/packages/0c/69/49aae86edb75fe16460b59a191fcc0f568c2378f780bb063850db0fe007a/torchvision-0.24.1-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:1e39619de698e2821d71976c92c8a9e50cdfd1e993507dfb340f2688bfdd8283", size = 2387757, upload-time = "2025-11-12T15:25:06.795Z" },
+    { url = "https://files.pythonhosted.org/packages/11/c9/1dfc3db98797b326f1d0c3f3bb61c83b167a813fc7eab6fcd2edb8c7eb9d/torchvision-0.24.1-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:a0f106663e60332aa4fcb1ca2159ef8c3f2ed266b0e6df88de261048a840e0df", size = 8047682, upload-time = "2025-11-12T15:25:21.125Z" },
+    { url = "https://files.pythonhosted.org/packages/fa/bb/cfc6a6f6ccc84a534ed1fdf029ae5716dd6ff04e57ed9dc2dab38bf652d5/torchvision-0.24.1-cp311-cp311-win_amd64.whl", hash = "sha256:a9308cdd37d8a42e14a3e7fd9d271830c7fecb150dd929b642f3c1460514599a", size = 4037588, upload-time = "2025-11-12T15:25:14.402Z" },
+    { url = "https://files.pythonhosted.org/packages/f0/af/18e2c6b9538a045f60718a0c5a058908ccb24f88fde8e6f0fc12d5ff7bd3/torchvision-0.24.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:e48bf6a8ec95872eb45763f06499f87bd2fb246b9b96cb00aae260fda2f96193", size = 1891433, upload-time = "2025-11-12T15:25:03.232Z" },
+    { url = "https://files.pythonhosted.org/packages/9d/43/600e5cfb0643d10d633124f5982d7abc2170dfd7ce985584ff16edab3e76/torchvision-0.24.1-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:7fb7590c737ebe3e1c077ad60c0e5e2e56bb26e7bccc3b9d04dbfc34fd09f050", size = 2386737, upload-time = "2025-11-12T15:25:08.288Z" },
+    { url = "https://files.pythonhosted.org/packages/93/b1/db2941526ecddd84884132e2742a55c9311296a6a38627f9e2627f5ac889/torchvision-0.24.1-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:66a98471fc18cad9064123106d810a75f57f0838eee20edc56233fd8484b0cc7", size = 8049868, upload-time = "2025-11-12T15:25:13.058Z" },
+    { url = "https://files.pythonhosted.org/packages/69/98/16e583f59f86cd59949f59d52bfa8fc286f86341a229a9d15cbe7a694f0c/torchvision-0.24.1-cp312-cp312-win_amd64.whl", hash = "sha256:4aa6cb806eb8541e92c9b313e96192c6b826e9eb0042720e2fa250d021079952", size = 4302006, upload-time = "2025-11-12T15:25:16.184Z" },
+    { url = "https://files.pythonhosted.org/packages/e4/97/ab40550f482577f2788304c27220e8ba02c63313bd74cf2f8920526aac20/torchvision-0.24.1-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:8a6696db7fb71eadb2c6a48602106e136c785642e598eb1533e0b27744f2cce6", size = 1891435, upload-time = "2025-11-12T15:25:28.642Z" },
+    { url = "https://files.pythonhosted.org/packages/30/65/ac0a3f9be6abdbe4e1d82c915d7e20de97e7fd0e9a277970508b015309f3/torchvision-0.24.1-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:db2125c46f9cb25dc740be831ce3ce99303cfe60439249a41b04fd9f373be671", size = 2338718, upload-time = "2025-11-12T15:25:26.19Z" },
+    { url = "https://files.pythonhosted.org/packages/10/b5/5bba24ff9d325181508501ed7f0c3de8ed3dd2edca0784d48b144b6c5252/torchvision-0.24.1-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:f035f0cacd1f44a8ff6cb7ca3627d84c54d685055961d73a1a9fb9827a5414c8", size = 8049661, upload-time = "2025-11-12T15:25:22.558Z" },
+    { url = "https://files.pythonhosted.org/packages/5c/ec/54a96ae9ab6a0dd66d4bba27771f892e36478a9c3489fa56e51c70abcc4d/torchvision-0.24.1-cp313-cp313-win_amd64.whl", hash = "sha256:16274823b93048e0a29d83415166a2e9e0bf4e1b432668357b657612a4802864", size = 4319808, upload-time = "2025-11-12T15:25:17.318Z" },
+    { url = "https://files.pythonhosted.org/packages/d5/f3/a90a389a7e547f3eb8821b13f96ea7c0563cdefbbbb60a10e08dda9720ff/torchvision-0.24.1-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:e3f96208b4bef54cd60e415545f5200346a65024e04f29a26cd0006dbf9e8e66", size = 2005342, upload-time = "2025-11-12T15:25:11.871Z" },
+    { url = "https://files.pythonhosted.org/packages/a9/fe/ff27d2ed1b524078164bea1062f23d2618a5fc3208e247d6153c18c91a76/torchvision-0.24.1-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:f231f6a4f2aa6522713326d0d2563538fa72d613741ae364f9913027fa52ea35", size = 2341708, upload-time = "2025-11-12T15:25:25.08Z" },
+    { url = "https://files.pythonhosted.org/packages/b1/b9/d6c903495cbdfd2533b3ef6f7b5643ff589ea062f8feb5c206ee79b9d9e5/torchvision-0.24.1-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:1540a9e7f8cf55fe17554482f5a125a7e426347b71de07327d5de6bfd8d17caa", size = 8177239, upload-time = "2025-11-12T15:25:18.554Z" },
+    { url = "https://files.pythonhosted.org/packages/4f/2b/ba02e4261369c3798310483028495cf507e6cb3f394f42e4796981ecf3a7/torchvision-0.24.1-cp313-cp313t-win_amd64.whl", hash = "sha256:d83e16d70ea85d2f196d678bfb702c36be7a655b003abed84e465988b6128938", size = 4251604, upload-time = "2025-11-12T15:25:34.069Z" },
+    { url = "https://files.pythonhosted.org/packages/42/84/577b2cef8f32094add5f52887867da4c2a3e6b4261538447e9b48eb25812/torchvision-0.24.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:cccf4b4fec7fdfcd3431b9ea75d1588c0a8596d0333245dafebee0462abe3388", size = 2005319, upload-time = "2025-11-12T15:25:23.827Z" },
+    { url = "https://files.pythonhosted.org/packages/5f/34/ecb786bffe0159a3b49941a61caaae089853132f3cd1e8f555e3621f7e6f/torchvision-0.24.1-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:1b495edd3a8f9911292424117544f0b4ab780452e998649425d1f4b2bed6695f", size = 2338844, upload-time = "2025-11-12T15:25:32.625Z" },
+    { url = "https://files.pythonhosted.org/packages/51/99/a84623786a6969504c87f2dc3892200f586ee13503f519d282faab0bb4f0/torchvision-0.24.1-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:ab211e1807dc3e53acf8f6638df9a7444c80c0ad050466e8d652b3e83776987b", size = 8175144, upload-time = "2025-11-12T15:25:31.355Z" },
+    { url = "https://files.pythonhosted.org/packages/6d/ba/8fae3525b233e109317ce6a9c1de922ab2881737b029a7e88021f81e068f/torchvision-0.24.1-cp314-cp314-win_amd64.whl", hash = "sha256:18f9cb60e64b37b551cd605a3d62c15730c086362b40682d23e24b616a697d41", size = 4234459, upload-time = "2025-11-12T15:25:19.859Z" },
+    { url = "https://files.pythonhosted.org/packages/50/33/481602c1c72d0485d4b3a6b48c9534b71c2957c9d83bf860eb837bf5a620/torchvision-0.24.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:ec9d7379c519428395e4ffda4dbb99ec56be64b0a75b95989e00f9ec7ae0b2d7", size = 2005336, upload-time = "2025-11-12T15:25:27.225Z" },
+    { url = "https://files.pythonhosted.org/packages/d0/7f/372de60bf3dd8f5593bd0d03f4aecf0d1fd58f5bc6943618d9d913f5e6d5/torchvision-0.24.1-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:af9201184c2712d808bd4eb656899011afdfce1e83721c7cb08000034df353fe", size = 2341704, upload-time = "2025-11-12T15:25:29.857Z" },
+    { url = "https://files.pythonhosted.org/packages/36/9b/0f3b9ff3d0225ee2324ec663de0e7fb3eb855615ca958ac1875f22f1f8e5/torchvision-0.24.1-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:9ef95d819fd6df81bc7cc97b8f21a15d2c0d3ac5dbfaab5cbc2d2ce57114b19e", size = 8177422, upload-time = "2025-11-12T15:25:37.357Z" },
+    { url = "https://files.pythonhosted.org/packages/d6/ab/e2bcc7c2f13d882a58f8b30ff86f794210b075736587ea50f8c545834f8a/torchvision-0.24.1-cp314-cp314t-win_amd64.whl", hash = "sha256:480b271d6edff83ac2e8d69bbb4cf2073f93366516a50d48f140ccfceedb002e", size = 4335190, upload-time = "2025-11-12T15:25:35.745Z" },
+]
+
 [[package]]
 name = "tqdm"
 version = "4.67.1"