[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Author: pre-commit-ci[bot]

README.md

@@ -44,7 +44,7 @@ gpus per model.
 
 This is an example to run a single 8bit llama-65b model on 2 A40s that have
 ~50 GB of memory each.
- 
+
 ```
 elk elicit huggyllama/llama-65b imdb --num_gpus 2 --gpus_per_model 2 --int8
 ```

tests/test_smoke_elicit.py

@@ -7,7 +7,7 @@
 
 def test_smoke_elicit_run_tiny_gpt2_ccs(tmp_path: Path):
     # we need about 5 mb of gpu memory to run this test
-    model_path, min_mem = "sshleifer/tiny-gpt2", 10 * 1024 ** 2
+    model_path, min_mem = "sshleifer/tiny-gpt2", 10 * 1024**2
     dataset_name = "imdb"
     elicit = Elicit(
         data=Extract(
@@ -38,7 +38,7 @@ def test_smoke_elicit_run_tiny_gpt2_ccs(tmp_path: Path):
 
 def test_smoke_elicit_run_tiny_gpt2_eigen(tmp_path: Path):
     # we need about 5 mb of gpu memory to run this test
-    model_path, min_mem = "sshleifer/tiny-gpt2", 10 * 1024 ** 2
+    model_path, min_mem = "sshleifer/tiny-gpt2", 10 * 1024**2
     dataset_name = "imdb"
     elicit = Elicit(
         data=Extract(

tests/test_smoke_eval.py

@@ -19,7 +19,7 @@ def setup_elicit(
     tmp_path: Path,
     dataset_name="imdb",
     model_path="sshleifer/tiny-gpt2",
-    min_mem=10 * 1024 ** 2,
+    min_mem=10 * 1024**2,
     is_ccs: bool = True,
 ) -> Elicit:
     """Setup elicit config for testing, execute elicit, and save output to tmp_path.

tests/test_split_devices.py

@@ -5,39 +5,30 @@ def test_split_2_devices_1_gpu_per_model():
     devices = ["a", "b"]
     gpus_per_model = 1
     models_to_create = 2
-    assert (
-        split_devices_into_model_devices(
-            devices=devices,
-            gpus_per_model=gpus_per_model,
-            models_to_create=models_to_create,
-        )
-        == [ModelDevices("a", []), ModelDevices("b", [])]
-    )
+    assert split_devices_into_model_devices(
+        devices=devices,
+        gpus_per_model=gpus_per_model,
+        models_to_create=models_to_create,
+    ) == [ModelDevices("a", []), ModelDevices("b", [])]
 
 
 def test_split_4_devices_2_gpus_per_model():
     devices = ["a", "b", "c", "d"]
     gpus_per_model = 2
     models_to_create = 2
-    assert (
-        split_devices_into_model_devices(
-            devices=devices,
-            gpus_per_model=gpus_per_model,
-            models_to_create=models_to_create,
-        )
-        == [ModelDevices("a", ["b"]), ModelDevices("c", ["d"])]
-    )
+    assert split_devices_into_model_devices(
+        devices=devices,
+        gpus_per_model=gpus_per_model,
+        models_to_create=models_to_create,
+    ) == [ModelDevices("a", ["b"]), ModelDevices("c", ["d"])]
 
 
 def test_split_7_devices_3_gpus_per_model():
     devices = ["a", "b", "c", "d", "e", "f", "g"]
     gpus_per_model = 3
     models_to_create = 2
-    assert (
-        split_devices_into_model_devices(
-            devices=devices,
-            gpus_per_model=gpus_per_model,
-            models_to_create=models_to_create,
-        )
-        == [ModelDevices("a", ["b", "c"]), ModelDevices("d", ["e", "f"])]
-    )
+    assert split_devices_into_model_devices(
+        devices=devices,
+        gpus_per_model=gpus_per_model,
+        models_to_create=models_to_create,
+    ) == [ModelDevices("a", ["b", "c"]), ModelDevices("d", ["e", "f"])]

tests/test_truncated_eigh.py

@@ -11,7 +11,7 @@ def random_symmetric_matrix(n: int, k: int) -> torch.Tensor:
     assert k <= n, "Rank k should be less than or equal to the matrix size n."
 
     # Generate random n x k matrix A with elements drawn from a uniform distribution
-    A = torch.rand(n, k) / k ** 0.5
+    A = torch.rand(n, k) / k**0.5
 
     # Create a diagonal matrix D with k eigenvalues evenly distributed around zero
     eigenvalues = torch.linspace(-1, 1, k)