modeling_bitcache.py

"""
Model Architecture Hugging Face com Compressão Espectral BitCache O(N)
Carrega o modelo base e acopla a decimação espectral de Cheeger-Catalan no KV-Cache.
"""

import os
import sys
from pathlib import Path
from transformers import GPT2LMHeadModel
from .configuration_bitcache import BitCacheGPT2Config

# Garante que o motor fechado .pyc do BitCache seja carregado
CURRENT_DIR = Path(__file__).resolve().parent
sys.path.insert(0, str(CURRENT_DIR))
sys.path.insert(0, str(CURRENT_DIR / "bitcache_bin"))

try:
    from bitcache import patch_model, unpatch_model, get_bitcache_metrics
except ImportError:
    # Fallback se importado dentro do repositório
    from bitcache_bin.bitcache import patch_model, unpatch_model, get_bitcache_metrics


class BitCacheGPT2LMHeadModel(GPT2LMHeadModel):
    config_class = BitCacheGPT2Config

    def __init__(self, config: BitCacheGPT2Config):
        super().__init__(config)
        self._bitcache_active = False
        self._init_bitcache()

    def _init_bitcache(self):
        """Ativa a compressão espectral de KV-Cache BitCache no modelo."""
        if not self._bitcache_active:
            patch_model(
                self,
                profile=getattr(self.config, "profile", "cheeger_catalan"),
                retention_ratio=getattr(self.config, "retention_ratio", 0.35),
                protect_sink=getattr(self.config, "protect_sink", 4),
                protect_recent=getattr(self.config, "protect_recent", 16),
                boundary_window=getattr(self.config, "boundary_window", 16),
                boundary_expansion=getattr(self.config, "boundary_expansion", 6),
            )
            self._bitcache_active = True

    def get_compression_metrics(self):
        """Retorna as métricas de tokens comprimidos e economizados."""
        return get_bitcache_metrics(self)