@contextmanager
    def __call__(self, model: PreTrainedModel) -> Generator:
        """
        Context manager that handles both initial prefilling and KVzip scoring/compression.
        
        This overrides the base class __call__ method to implement the full KVzip algorithm:
        1. First yield: allows initial prefilling with context
        2. After yield: performs KVzip scoring and compression using context reconstruction
        """
        if not isinstance(model, SUPPORTED_MODELS):
            logger.warning(f"Model {type(model)} not tested, supported models: {SUPPORTED_MODELS}")

        if isinstance(model, Gemma3ForCausalLM):
            logger.warning("Compression in Gemma3 is only applied to layer without sliding window attention")

        # Store model reference for later use
        tokenizer = AutoTokenizer.from_pretrained(model.config.name_or_path)

        # Get suffix_ids directly using tokenizer's chat template (do this once, not in hook)
        if tokenizer.chat_template is None:
            suffix_text = "\n"  # Default suffix for models without chat template
        else:
            # Use a dummy context to extract the question suffix from chat template
            dummy_context = "dummy context"
            separator = "\n" + "#" * len(dummy_context)
            temp_context = tokenizer.apply_chat_template(
                [{"role": "user", "content": dummy_context + separator}],
                add_generation_prompt=True,
                tokenize=False
            )
            _, suffix_text = temp_context.split(separator)
        
        # Tokenize suffix directly to ids
        self._suffix_ids = tokenizer.encode(suffix_text, return_tensors="pt", add_special_tokens=False)

        # Register embedding hook to capture context information
        hooks = []
        try:
            # First yield: Initial prefilling phase (no compression hooks yet)
            embedding_hook = model.model.embed_tokens.register_forward_hook(self._forward_hook_embedding,
                                                                            with_kwargs=True)
            yield
            # Remove embedding hook since we no longer need it
            embedding_hook.remove()

            # After yield: KVzip scoring and compression phase
            if self.compression_ratio > 0 and self._context_ids is not None:
                # Now register attention hooks for compression
                for layer in model.model.layers:
                    if isinstance(model, Gemma3ForCausalLM) and layer.is_sliding:
                        continue
                    layer.self_attn.rotary_emb = model.model.rotary_emb
                    hooks.append(layer.self_attn.register_forward_hook(self.forward_hook, with_kwargs=True))

                self._perform_kvzip_compression(model, tokenizer)
        finally:
            for hook in hooks:
                hook.remove()

    def _forward_hook_embedding(self, module: nn.Module, input: list[torch.Tensor], kwargs: dict, output: list):
        """
        Hook for embedding layer to capture context information from the first forward pass.

        """
        self._context_ids = input[0]
        self._cache = ... # fetch from kwargs

        return output


    def _perform_kvzip_compression(self, model: PreTrainedModel, tokenizer: PreTrainedTokenizer):
        """
        Perform the KVzip scoring and compression algorithm.
        """
        context_length = self._context_ids.shape[1]
        self.context_length = context_length

        # Prepare chunked inputs for context reconstruction
        input_ids = self.prepare(model, tokenizer, context_length)

        # Reset start_idx for scoring
        self.start_idx = 0

        # Perform scoring through context reconstruction
        # Use the stored cache from the initial forward pass
        for prefill_ids, repeat_ids in input_ids:
            self.end_idx = self.start_idx + prefill_ids.shape[1]
            # Pass the cache that was used in the initial forward pass
            model(
                input_ids=repeat_ids.to(model.device),
                past_key_values=self._cache,
                num_logits_to_keep=1,
            )
            self.start_idx = self.end_idx

        # Verify tokenization consistency
        assert self.end_idx == context_length, "Tokenization is not consistent"

        # Perform final compression
        self.compress_post(model)