File size: 1,851 Bytes
2c65306 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 | """
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)
|