Add post_init_from_model to BasePress (#163)

Author: SimJeg

evaluation/evaluate.py

@@ -329,9 +329,7 @@ def _load_and_prepare_dataset(self):
             # FinchPress uses a delimiter token to separate context and question
             # So we need to update the tokenizer and the model embeddings.
             logger.info("FinchPress detected, updating model and tokenizer with delimiter token.")
-            self.press.update_model_and_tokenizer(
-                self.pipeline.model, self.pipeline.tokenizer
-            )  # type: ignore[attr-defined]
+            self.press.update_model_and_tokenizer(self.pipeline.model, self.pipeline.tokenizer)  # type: ignore[attr-defined]
             df["context"] = df["context"] + self.press.delimiter_token  # type: ignore[attr-defined, index]
 
         if self.config.compress_questions:

kvpress/presses/adakv_press.py

@@ -39,6 +39,9 @@ def __post_init__(self):
         assert isinstance(self.press, ScorerPress), "AdaKVPress requires a ScorerPress as input"
         assert 0 <= self.alpha_safeguard <= 1, "alpha_safeguard should be in [0, 1]"
 
+    def post_init_from_model(self, model):
+        self.press.post_init_from_model(model)
+
     @property
     def compression_ratio(self):
         return self.press.compression_ratio

kvpress/presses/base_press.py

@@ -46,6 +46,12 @@ class BasePress:
     The compression is applied only during pre-filling (not during generation).
     """
 
+    def post_init_from_model(self, model: PreTrainedModel):
+        """
+        Optional method to initialize press parameters from the model
+        """
+        pass
+
     def compress(
         self,
         module: nn.Module,
@@ -179,6 +185,7 @@ def __call__(self, model: PreTrainedModel) -> Generator:
         if isinstance(model, Gemma3ForConditionalGeneration):
             logger.warning_once("Compression in Gemma3 is only applied to layer without sliding window attention")
 
+        self.post_init_from_model(model)
         hooks = []
         try:
             language_model = model.model.language_model if hasattr(model.model, "language_model") else model.model

kvpress/presses/block_press.py

@@ -36,6 +36,9 @@ class BlockPress(BasePress):
     def __post_init__(self):
         assert isinstance(self.press, ScorerPress), "BlockPress requires a ScorerPress"
 
+    def post_init_from_model(self, model):
+        self.press.post_init_from_model(model)
+
     @property
     def compression_ratio(self):
         return self.press.compression_ratio

kvpress/presses/chunk_press.py

@@ -36,6 +36,9 @@ class ChunkPress(BasePress):
     def __post_init__(self):
         assert isinstance(self.press, ScorerPress), "ChunkPress requires a ScorerPress as input"
 
+    def post_init_from_model(self, model):
+        self.press.post_init_from_model(model)
+
     @property
     def compression_ratio(self):
         return self.press.compression_ratio

kvpress/presses/chunkkv_press.py

@@ -37,6 +37,9 @@ class ChunkKVPress(BasePress):
     def __post_init__(self):
         assert isinstance(self.press, ScorerPress), "ChunkKVPress requires a ScorerPress as input"
 
+    def post_init_from_model(self, model):
+        self.press.post_init_from_model(model)
+
     @property
     def compression_ratio(self):
         return self.press.compression_ratio

kvpress/presses/composed_press.py

@@ -49,6 +49,10 @@ def __post_init__(self):
             isinstance(press, (AdaKVPress, KVzipPress)) for press in self.presses
         ), "ComposedPress cannot contains AdaKVPress or KVzipPress"
 
+    def post_init_from_model(self, model):
+        for press in self.presses:
+            press.post_init_from_model(model)
+
     def forward_hook(self, module, input, kwargs, output):
         retained_fraction = 1.0
         for press in self.presses:

kvpress/presses/criticalkv_press.py

@@ -43,6 +43,9 @@ def __init__(self, press: ScorerPress, epsilon: float = 1e-4, first_stage_ratio:
         if isinstance(self.press, ExpectedAttentionPress) and self.press.use_vnorm:
             logger.warning("use_vnorm should be disabled for CriticalKVPress")
 
+    def post_init_from_model(self, model):
+        self.press.post_init_from_model(model)
+
     @property  # type: ignore[misc]
     def compression_ratio(self):  #
         return self.press.compression_ratio

kvpress/presses/decoding_press.py

@@ -60,6 +60,9 @@ def __post_init__(self):
                 f"This will be overridden by the decoding press."
             )
 
+    def post_init_from_model(self, model):
+        self.base_press.post_init_from_model(model)
+
     def compress(
         self,
         module: nn.Module,

kvpress/presses/duo_attention_press.py

@@ -1,7 +1,6 @@
 # SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from contextlib import contextmanager
 from dataclasses import dataclass, field
 from io import StringIO
 
@@ -70,7 +69,7 @@ class DuoAttentionPress(BasePress):
     sink_size: int = field(init=False, default=None)
     streaming_mask: torch.Tensor = field(init=False, default=None)
 
-    def __post_init_from_model__(self, model):
+    def post_init_from_model(self, model):
         """
         Initialize sink_size, recent_size, and streaming_mask from a model
         """
@@ -101,7 +100,7 @@ def compress(self, module, hidden_states, keys, values, attentions, kwargs):
         assert module.config._attn_implementation != "eager", "eager mode not supported"
         if self.streaming_mask is None:
             raise ValueError(
-                "Streaming mask not initialized. Make sure to call __post_init_from_model__ to initialize this press."
+                "Streaming mask not initialized. Make sure to call post_init_from_model to initialize this press."
             )
         k_len = keys.shape[2]
 
@@ -141,12 +140,6 @@ def load_attention_pattern(model):
 
         return config["sink_size"], config["recent_size"], head_scores
 
-    @contextmanager
-    def __call__(self, model):
-        self.__post_init_from_model__(model)
-        with super().__call__(model):
-            yield
-
 
 @cached(cache, key=lambda model, num_samples=50, q_len=500: (model.config.name_or_path, num_samples, q_len))
 def duo_attention_on_the_fly(model, num_samples=50, q_len=500):