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

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

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
 ```

elk/extraction/extraction.py

@@ -157,9 +157,7 @@ def extract_hiddens(
     world_size: int = 1,
 ) -> Iterable[dict]:
     first_device = (
-        devices
-        if not isinstance(devices, ModelDevices)
-        else devices.first_device
+        devices if not isinstance(devices, ModelDevices) else devices.first_device
     )
     """Run inference on a model with a set of prompts, yielding the hidden states."""
     os.environ["TOKENIZERS_PARALLELISM"] = "false"
@@ -174,15 +172,13 @@ def extract_hiddens(
     ds_names = cfg.datasets
     assert len(ds_names) == 1, "Can only extract hiddens from one dataset at a time."
 
-
     model = instantiate_model_with_devices(
         cfg=cfg, device_config=devices, is_verbose=is_verbose
     )
     tokenizer = instantiate_tokenizer(
         cfg.model, truncation_side="left", verbose=is_verbose
     )
 
-
     is_enc_dec = model.config.is_encoder_decoder
     if is_enc_dec and cfg.use_encoder_states:
         assert hasattr(model, "get_encoder") and callable(model.get_encoder)

elk/utils/multi_gpu.py

@@ -4,7 +4,6 @@
 
 import torch
 from accelerate import infer_auto_device_map, init_empty_weights
-from torch import dtype
 from torch.nn import Module
 from transformers import PreTrainedModel
 

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)