Code refactor 🤖🤖🤖

evaluation/benchmarks/infinite_bench/calculate_metrics.py

@@ -14,12 +14,39 @@
 
 def calculate_metrics(df):
     preds = df["predicted_answer"].tolist()
-    labels = df["answer"].tolist()
     task = df["task"].tolist()[0]
+    labels = [to_upstream_label(row, task) for _, row in df.iterrows()]
     model = "dummy_model"  # not really used
     return get_score(labels, preds, task, model)
 
 
+# The Hub dataset stores every answer as a list of strings and keeps only the letter of multiple-choice answers,
+# while the upstream scorers below expect numbers for code_run, math_find and math_calc, and [option_text, letter]
+# for code_debug and longbook_choice_eng (see get_answer in InfiniteBench/src/eval_utils.py).
+OPTIONS_PATTERN = re.compile(r"\nA\. (.*)\nB\. (.*)\nC\. (.*)\nD\. (.*)\Z")
+
+
+def to_upstream_label(row, task: str) -> list:
+    label = list(row["answer"])
+    if task in ["code_run", "math_find", "math_calc"]:
+        return [to_number(value) for value in label]
+    if task in ["code_debug", "longbook_choice_eng"]:
+        # The options end the question, which query-aware evaluation moves to the end of the context
+        match = OPTIONS_PATTERN.search(row["question"] or row["context"])
+        if match is None:
+            raise ValueError(f"Could not find the A-D options of the {task} question")
+        letter = label[0]
+        return [match.group("ABCD".index(letter) + 1), letter]
+    return label
+
+
+def to_number(value: str) -> int | float:
+    try:
+        return int(value)
+    except ValueError:
+        return float(value)
+
+
 def normalize_answer(s: str) -> str:
     """Lower text and remove punctuation, articles and extra whitespace."""
 

evaluation/benchmarks/loogle/create_huggingface_dataset.py

@@ -38,9 +38,7 @@
 
     if task == "shortdep_cloze":
         df["answer"] = df["answer"].apply(
-            lambda values: json.dumps(
-                {f"<mask-{i}>": value for i, value in enumerate(values)}, ensure_ascii=False
-            )
+            lambda values: json.dumps({f"<mask-{i}>": value for i, value in enumerate(values)}, ensure_ascii=False)
         )
 
     df["context"] = df["context"].apply(lambda x: context_prompt[task].format(input=x))

evaluation/benchmarks/needle_in_haystack/utils.py

@@ -51,16 +51,21 @@ def insert_needle_in_haystack(
     """
 
     # Store the original context and needle to be reused for each depth
-    original_context = df["context"][0]
-    needle_text = needle_text or df["needle"][0]
-    question_text = question_text or df["question"][0]
-    answer_prefix = answer_prefix or df["answer_prefix"][0]
-    max_new_tokens = df["max_new_tokens"][0]
+    original_context = df["context"].iloc[0]
+    needle_text = needle_text or df["needle"].iloc[0]
+    question_text = question_text or df["question"].iloc[0]
+    answer_prefix = answer_prefix or df["answer_prefix"].iloc[0]
+    max_new_tokens = df["max_new_tokens"].iloc[0]
 
     logger.info(f"Preparing dataset for inference. Needle: {needle_text}")
     tokenized_needle = tokenizer.encode(needle_text, add_special_tokens=False)
     # Account for system prompts and other overhead
     context_length_limit = max_context_length - len(tokenized_needle) - 150
+    if context_length_limit <= 0:
+        raise ValueError(
+            f"max_context_length ({max_context_length}) must exceed the needle length ({len(tokenized_needle)} "
+            "tokens) plus 150 tokens reserved for the prompt"
+        )
     # Tokenize the original context once
     tokenized_context = tokenizer.encode(original_context, add_special_tokens=False)[:context_length_limit]
     # Initialize a list to hold the new rows

evaluation/evaluate.py

@@ -5,7 +5,7 @@
 import logging
 import random
 import sys
-from dataclasses import asdict, dataclass
+from dataclasses import asdict, dataclass, fields
 from pathlib import Path
 from typing import Any, Dict, Optional
 
@@ -80,6 +80,16 @@ class EvaluationConfig:
 
     def __post_init__(self):
         """Validate configuration after initialization."""
+        # Fire parses "--trust_remote_code false" as the (truthy) string "false"
+        for config_field in fields(self):
+            if config_field.type is bool:
+                value = getattr(self, config_field.name)
+                if isinstance(value, str) and value.lower() in ["true", "false"]:
+                    value = value.lower() == "true"
+                    setattr(self, config_field.name, value)
+                if not isinstance(value, bool):
+                    raise ValueError(f"{config_field.name} must be a boolean, got {value!r}")
+
         # Validate dataset
         assert self.dataset in DATASET_REGISTRY, f"No dataset found for {self.dataset}"
         assert self.dataset in SCORER_REGISTRY, f"No scorer found for {self.dataset}"
@@ -173,7 +183,7 @@ def save_config(self, config_filename: Path):
         if self.head_compression_ratio is None:
             config_dict.pop("head_compression_ratio", None)
         with open(str(config_filename), "w") as f:
-            yaml.dump(config_dict, f, default_flow_style=False, indent=2, sort_keys=False)
+            yaml.safe_dump(config_dict, f, default_flow_style=False, indent=2, sort_keys=False)
 
 
 def _load_yaml_config(path: str | Path) -> dict:
@@ -327,7 +337,9 @@ def _load_and_prepare_dataset(self):
         logger.info(f"Loading dataset: {DATASET_REGISTRY[dataset_name]} (data_dir: {data_dir})")
         df = load_dataset(DATASET_REGISTRY[dataset_name], data_dir=data_dir, split="test").to_pandas()
 
-        if fraction < 1.0:
+        if fraction < 1.0 and dataset_name == "needle_in_haystack":
+            logger.warning("Ignoring fraction for needle_in_haystack: the haystack dataset has a single row.")
+        elif fraction < 1.0:
             original_len = len(df)
             df = df.sample(frac=fraction, random_state=self.config.seed)
             logger.info(f"Sampled {len(df)} samples ({fraction:.2f}) from original {original_len} samples.")
@@ -366,7 +378,7 @@ def _setup_model_pipeline(self):
             device = "auto" if torch.cuda.is_available() else "cpu"
             logger.info(f"No device specified, auto-detected device: {device}")
 
-        model_kwargs = self.config.model_kwargs or {}
+        model_kwargs = dict(self.config.model_kwargs or {})
 
         if self.config.fp8:
             model_kwargs["quantization_config"] = FineGrainedFP8Config()
@@ -375,6 +387,8 @@ def _setup_model_pipeline(self):
         if isinstance(self.press, ObservedAttentionPress):
             model_kwargs["attn_implementation"] = "eager"
             logger.info("ObservedAttentionPress detected, setting attn_implementation to 'eager'.")
+        elif model_kwargs.get("attn_implementation") is not None:
+            logger.info(f"Using attn_implementation='{model_kwargs['attn_implementation']}' from model_kwargs.")
         else:
             try:
                 import flash_attn  # noqa: F401

evaluation/evaluate_registry.py

@@ -125,7 +125,6 @@
     "no_press": None,
     "cam_streaming_llm": CAMPress(base_press=StreamingLLMPress()),
     "cam_knorm": CAMPress(base_press=KnormPress()),
-    "cam_adakv_snapkv": CAMPress(base_press=AdaKVPress(SnapKVPress())),
     "cam_tova": CAMPress(base_press=TOVAPress()),
     "decoding_knorm": DecodingPress(base_press=KnormPress()),
     "decoding_streaming_llm": DecodingPress(base_press=StreamingLLMPress()),

kvpress/attention_patch.py

@@ -48,6 +48,18 @@ def search_hyperplane(X, max_iter: int = 1000):
     raise ValueError("Could not find fake keys such that for every query q, exp(<q, k>) = 0")
 
 
+def check_masked_key_indices_support(module):
+    """
+    Raise a ValueError if the attention implementation of the module ignores module.masked_key_indices.
+    Eager attention is not registered in ALL_ATTENTION_FUNCTIONS, so it is not patched by patch_attention_functions.
+    """
+    if module.config._attn_implementation == "eager":
+        raise ValueError(
+            "Head-wise compression is not supported with attn_implementation='eager' as masked keys would be ignored. "
+            "Please use another attention implementation such as 'sdpa' or 'flash_attention_2'."
+        )
+
+
 def attention_patch(func):
     """
     Decorator to update the keys before the attention computation at the indices provided in module.masked_key_indices
@@ -70,7 +82,7 @@ def wrapper(module, query, key, value, attention_mask, dropout, **kwargs):
         if query.shape[2] == key.shape[2]:
             # Prefilling
             module.masked_key_indices = None
-        elif getattr(module, "masked_key_indices", None) is not None:
+        elif getattr(module, "masked_key_indices", None) is not None and module.masked_key_indices[-1].numel() > 0:
             # Decoding: build fake keys k s.t. exp(<q, k>) = 0
             bsz, num_heads, seq_len, head_dim = query.shape
             num_key_value_heads = key.shape[1]

kvpress/pipeline.py

@@ -138,7 +138,7 @@ def preprocess(
 
         # Apply chat template if available
         if self.tokenizer.chat_template is None:
-            bos_token = getattr(self.tokenizer, "bos_token", "")
+            bos_token = self.tokenizer.bos_token or ""
             context = bos_token + context
             question_suffix = "\n"  # to separate the question from the answer
         else:
@@ -199,9 +199,15 @@ def _forward(
         list[str]
             Generated answers for each input question.
         """
-        if isinstance(press, (DecodingPress, PrefillDecodingPress)) and len(input_tensors["questions_ids"]) > 1:
+        compresses_during_decoding = (
+            isinstance(press, DecodingPress)
+            or (isinstance(press, PrefillDecodingPress) and press.decoding_press is not None)
+            or (isinstance(press, DMSPress) and press.decoding)
+        )
+        if compresses_during_decoding and len(input_tensors["questions_ids"]) > 1:
             raise ValueError(
-                "DecodingPress is not compatible with multiple questions. Please specify a single question."
+                f"{type(press).__name__} compresses the cache during decoding and is not compatible with multiple "
+                "questions. Please specify a single question."
             )
 
         context_ids = input_tensors["context_ids"].to(self.model.device)
@@ -305,9 +311,12 @@ def generate_answer(
         position_ids = position_ids[:, -1:] + 1
         generated_ids = [outputs.logits[0, -1].argmax()]
 
-        should_stop_token_ids = self.model.generation_config.eos_token_id
-        if not isinstance(should_stop_token_ids, list):
-            should_stop_token_ids = [should_stop_token_ids]
+        eos_token_ids = self.model.generation_config.eos_token_id
+        if not isinstance(eos_token_ids, list):
+            eos_token_ids = [eos_token_ids]
+        should_stop_token_ids = {
+            token_id for token_id in [*eos_token_ids, self.tokenizer.eos_token_id] if token_id is not None
+        }
 
         for i in range(max_new_tokens - 1):
             outputs = self.model(

kvpress/presses/adakv_press.py

@@ -6,6 +6,7 @@
 
 import torch
 
+from kvpress.attention_patch import check_masked_key_indices_support
 from kvpress.presses.base_press import BasePress
 from kvpress.presses.scorer_press import ScorerPress
 from kvpress.utils import compute_n_kept
@@ -55,7 +56,7 @@ def compress(self, module, hidden_states, keys, values, attentions, kwargs):
         if self.compression_ratio == 0:
             return keys, values
 
-        assert module.config._attn_implementation != "eager", "eager mode not supported"
+        check_masked_key_indices_support(module)
 
         # Compute scores
         scores = self.press.score(module, hidden_states, keys, values, attentions, kwargs)

kvpress/presses/base_press.py

@@ -15,12 +15,11 @@
     MistralForCausalLM,
     Phi3ForCausalLM,
     PreTrainedModel,
-    QuantizedCache,
     Qwen2ForCausalLM,
     Qwen3ForCausalLM,
 )
 
-from kvpress.utils import extract_keys_and_values
+from kvpress.utils import extract_keys_and_values, set_keys_and_values
 
 logger = logging.getLogger(__name__)
 
@@ -138,7 +137,6 @@ def forward_hook(self, module: nn.Module, input: list[torch.Tensor], kwargs: dic
 
         hidden_states = kwargs["hidden_states"]
         cache = kwargs["past_key_values"]
-        cache_layer = cache.layers[module.layer_idx]
         q_len = hidden_states.shape[1]
 
         # Don't compress after pre-filling
@@ -149,15 +147,7 @@ def forward_hook(self, module: nn.Module, input: list[torch.Tensor], kwargs: dic
 
         keys, values = self.compress(module, hidden_states, keys, values, output[1], kwargs)
 
-        if isinstance(cache, QuantizedCache):
-            cache_layer._quantized_keys = cache_layer._quantize(keys, axis=cache_layer.axis_key)
-            cache_layer._quantized_values = cache_layer._quantize(values, axis=cache_layer.axis_value)
-            cache_layer.keys = torch.zeros(0, dtype=keys.dtype, device=keys.device)  # type: ignore[index]
-            cache_layer.values = torch.zeros(0, dtype=keys.dtype, device=keys.device)  # type: ignore[index]
-            cache_layer.cumulative_length = keys.shape[2]
-        else:
-            cache_layer.keys = keys
-            cache_layer.values = values
+        set_keys_and_values(cache, module.layer_idx, keys, values)
 
         return output
 
@@ -200,6 +190,9 @@ def __call__(self, model: PreTrainedModel) -> Generator:
                     # Skip layers with sliding window attention, only for Gemma3
                     continue
                 layer.self_attn.rotary_emb = language_model.rotary_emb
+                # Masks set by head-wise presses during prefilling must survive later press contexts (e.g. decoding)
+                if not hasattr(layer.self_attn, "masked_key_indices"):
+                    layer.self_attn.masked_key_indices = None
                 hooks.append(layer.self_attn.register_forward_hook(self.forward_hook, with_kwargs=True))
             yield
         finally: