fix format

Signed-off-by: Dominic789654 <xliu29@gmu.edu>

Author: Dominic789654

evaluation/evaluate.py

@@ -17,9 +17,11 @@
 from zero_scrolls.calculate_metrics import calculate_metrics as zero_scrolls_scorer
 
 from kvpress import (
-    CriticalKVPress,
-    CriticalAdaKVPress,
     AdaKVPress,
+    ChunkKVPress,
+    CriticalAdaKVPress,
+    CriticalKVPress,
+    DuoAttentionPress,
     ExpectedAttentionPress,
     KnormPress,
     ObservedAttentionPress,
@@ -28,8 +30,6 @@
     StreamingLLMPress,
     ThinKPress,
     TOVAPress,
-    DuoAttentionPress,
-    ChunkKVPress
 )
 
 logger = logging.getLogger(__name__)

kvpress/__init__.py

@@ -9,6 +9,8 @@
 from kvpress.presses.chunk_press import ChunkPress
 from kvpress.presses.chunkkv_press import ChunkKVPress
 from kvpress.presses.composed_press import ComposedPress
+from kvpress.presses.criticalkv_press import CriticalAdaKVPress, CriticalKVPress
+from kvpress.presses.duo_attention_press import DuoAttentionPress
 from kvpress.presses.expected_attention_press import ExpectedAttentionPress
 from kvpress.presses.key_rerotation_press import KeyRerotationPress
 from kvpress.presses.knorm_press import KnormPress
@@ -21,8 +23,6 @@
 from kvpress.presses.streaming_llm_press import StreamingLLMPress
 from kvpress.presses.think_press import ThinKPress
 from kvpress.presses.tova_press import TOVAPress
-from kvpress.presses.criticalkv_press import CriticalKVPress, CriticalAdaKVPress
-from kvpress.presses.duo_attention_press import DuoAttentionPress
 
 # Patch the attention functions to support head-wise compression
 patch_attention_functions()

kvpress/presses/chunkkv_press.py

@@ -51,7 +51,7 @@ def compress(
         assert attentions is None, "ChunkPress does not support attentions."
 
         kv_len = keys.shape[2]
-        
+
         # 1. Calculate global scores first
         global_scores = self.press.score(
             module,

kvpress/presses/criticalkv_press.py

@@ -8,8 +8,8 @@
 from transformers.models.llama.modeling_llama import repeat_kv
 
 from kvpress.presses.base_press import BasePress
-from kvpress.presses.scorer_press import ScorerPress
 from kvpress.presses.expected_attention_press import ExpectedAttentionPress
+from kvpress.presses.scorer_press import ScorerPress
 
 logger = logging.getLogger(__name__)
 
@@ -49,7 +49,7 @@ def vwl1norm(values, module):
         # Future kernel fusion optimization could eliminate this intermediate variables to enhance performance.
         head_WoV_norm_list = []
         for head in range(V.size(1)):
-            head_WoV = V[: , head, : , ...].matmul(Wo[head, ...].unsqueeze(0))
+            head_WoV = V[:, head, :, ...].matmul(Wo[head, ...].unsqueeze(0))
             head_WoV_norm = torch.norm(head_WoV, p=1, dim=-1)
             head_WoV_norm_list.append(head_WoV_norm)
 

kvpress/presses/duo_attention_press.py

@@ -1,17 +1,16 @@
 # SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from io import StringIO
-from dataclasses import dataclass, field
 from contextlib import contextmanager
+from dataclasses import dataclass, field
+from io import StringIO
 
-import torch
-import requests  # type: ignore[import-untyped]
 import numpy as np
+import requests  # type: ignore[import-untyped]
+import torch
 
 from kvpress.presses.base_press import BasePress
 
-
 PATTERNS_DICT = {
     "togethercomputer/Llama-2-7B-32K-Instruct": "Llama-2-7B-32K-Instruct/lr%3D0.02-reg%3D0.05-ctx%3D1000_32000-multi_passkey10",  # noqa: E501
     "gradientai//Llama-3-8B-Instruct-Gradient-1048k": "Llama-3-8B-Instruct-Gradient-1048k/lr%3D0.02-reg%3D0.05-ctx%3D1000_32000-multi_passkey10",  # noqa: E501

tests/default_presses.py

@@ -4,6 +4,7 @@
 import numpy as np
 
 from kvpress import (
+    DuoAttentionPress,
     ExpectedAttentionPress,
     KnormPress,
     RandomPress,
@@ -12,7 +13,6 @@
     StreamingLLMPress,
     ThinKPress,
     TOVAPress,
-    DuoAttentionPress,
 )
 
 

tests/presses/test_duo_attention_press.py

@@ -1,4 +1,4 @@
-from kvpress.presses.duo_attention_press import DuoAttentionPress, PATTERNS_DICT
+from kvpress.presses.duo_attention_press import PATTERNS_DICT, DuoAttentionPress
 
 
 def test_load_attention_pattern():

tests/presses/test_presses.py

@@ -8,12 +8,12 @@
 from transformers import DynamicCache
 
 from kvpress import (
-    CriticalKVPress,
-    CriticalAdaKVPress,
     AdaKVPress,
     ChunkKVPress,
     ChunkPress,
     ComposedPress,
+    CriticalAdaKVPress,
+    CriticalKVPress,
     KeyRerotationPress,
     KnormPress,
     ObservedAttentionPress,
@@ -57,8 +57,10 @@ def test_chunkkv_press(unit_test_model):  # noqa: F811
 
 
 @pytest.mark.parametrize("press_dict", default_presses)
-@pytest.mark.parametrize("wrapper_press", [None, ComposedPress, KeyRerotationPress, AdaKVPress, ChunkPress,
-                                           CriticalKVPress, CriticalAdaKVPress])
+@pytest.mark.parametrize(
+    "wrapper_press",
+    [None, ComposedPress, KeyRerotationPress, AdaKVPress, ChunkPress, CriticalKVPress, CriticalAdaKVPress],
+)
 def test_presses_run(unit_test_model, press_dict, wrapper_press):  # noqa: F811
     cls = press_dict["cls"]
     for kwargs in press_dict["kwargs"]: