From 8593a5b4164ffaf28cfa52952cbddb21f55708ba Mon Sep 17 00:00:00 2001 From: Philipp Emanuel Weidmann Date: Tue, 23 Jun 2026 10:18:34 +0530 Subject: [PATCH] feat: add end-to-end tests --- .gitignore | 9 +++-- pyproject.toml | 2 + src/heretic/main.py | 6 ++- tests/mistral-3/SHA256SUMS | 7 ++++ tests/mistral-3/config.toml | 40 ++++++++++++++++++++ tests/run_tests.py | 73 +++++++++++++++++++++++++++++++++++++ uv.lock | 45 +++++++++++++++++++++++ 7 files changed, 177 insertions(+), 5 deletions(-) create mode 100644 tests/mistral-3/SHA256SUMS create mode 100644 tests/mistral-3/config.toml create mode 100644 tests/run_tests.py diff --git a/.gitignore b/.gitignore index 1241cea..851d494 100644 --- a/.gitignore +++ b/.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/ diff --git a/pyproject.toml b/pyproject.toml index b8074b5..cca44c0 100644 --- a/pyproject.toml +++ b/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", ] diff --git a/src/heretic/main.py b/src/heretic/main.py index 243841b..832c730 100644 --- a/src/heretic/main.py +++ b/src/heretic/main.py @@ -612,7 +612,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( @@ -798,7 +798,9 @@ def run(): print() print("Restoring model from reproduction information...") else: - print() + if settings.trial_index is None: + print() + trial = ask_if_unset( None if settings.trial_index is None diff --git a/tests/mistral-3/SHA256SUMS b/tests/mistral-3/SHA256SUMS new file mode 100644 index 0000000..62bf8c2 --- /dev/null +++ b/tests/mistral-3/SHA256SUMS @@ -0,0 +1,7 @@ +39f03c383413f531fd302c06c7e982ad98c83f0657a8339ae25478ccb81fdcda chat_template.jinja +f69f84977a47c8fea9ce9fc26b7de379216cb01146ea726a87996d3554cfcd19 config.json +34dfa6012ca9ac5f57e5521d8dbaecbc7ab7f7ab0fd96ec020b543aab5f265d9 generation_config.json +26b0435167dba8138d21eff9c511e909d501361805f946828de0fd19ada434dc model.safetensors +84be30b124b50749c56d25fdbec5ccedf564446f6b3b035e88e1e07b986d2491 processor_config.json +c3a8d92e371b92a2cd6e678e31ebc27d0235e929a51fbf290f74742b341fa96f tokenizer.json +7b29c843c0043622d28fd4638451cbb0a609d99a0762ffbff3b92b4b2fee4d94 tokenizer_config.json diff --git a/tests/mistral-3/config.toml b/tests/mistral-3/config.toml new file mode 100644 index 0000000..5600222 --- /dev/null +++ b/tests/mistral-3/config.toml @@ -0,0 +1,40 @@ +model = "tiny-random/mistral-3" +model_commit = "931aa2e5c9668fc3679e56aa44972fe18597d55d" + +batch_size = 2 +max_response_length = 10 +kl_divergence_target = 0 +n_trials = 2 +n_startup_trials = 1 + +seed = 12345 + +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" diff --git a/tests/run_tests.py b/tests/run_tests.py new file mode 100644 index 0000000..c4d7dd9 --- /dev/null +++ b/tests/run_tests.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: AGPL-3.0-or-later +# Copyright (C) 2025-2026 Philipp Emanuel Weidmann + 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 + +for test_directory in script_directory.iterdir(): + if test_directory.is_dir(): + config_file = test_directory / "config.toml" + hash_file = test_directory / "SHA256SUMS" + + if config_file.is_file() and hash_file.is_file(): + 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() + + # To update the hashes after a logic change, run the tests, then execute + # + # cd /model + # sha256sum * > ../SHA256SUMS + + with open(hash_file, "r", encoding="utf-8") as file: + for line in file: + if line.strip(): + original_sha256, filename = line.split() + sha256 = get_file_sha256(test_directory / "model" / filename) + + if sha256.lower() != original_sha256.lower(): + sys.exit( + ( + f"Test {test_directory.name} has FAILED!\n" + f"Output file {filename} doesn't match.\n" + f"Expected hash: {original_sha256}\n" + f"Actual hash: {sha256}" + ) + ) + +print("All tests passed!") diff --git a/uv.lock b/uv.lock index 9f75d44..28a70bf 100644 --- a/uv.lock +++ b/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"