Update config.py

config.py

@@ -7,7 +7,7 @@ import json
 import hashlib
 from dataclasses import dataclass, field, asdict
 from enum import Enum
-from typing import List, Optional, NamedTuple
+from typing import List, Optional, NamedTuple, Dict, Any
 from datetime import datetime
 import torch
 import transformers
@@ -17,6 +17,69 @@ import logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
+# Model configurations - NO HARDCODING
+SUPPORTED_MODELS: Dict[str, Dict[str, Any]] = {
+    "gpt2": {
+        "name": "gpt2",
+        "requires_auth": False,
+        "max_context": 1024,
+        "default_dtype": "float16"
+    },
+    "llama2-7b": {
+        "name": "meta-llama/Llama-2-7b-hf",
+        "requires_auth": True,
+        "max_context": 4096,
+        "default_dtype": "float16"
+    },
+    "mistral-7b": {
+        "name": "mistralai/Mistral-7B-v0.1",
+        "requires_auth": False,
+        "max_context": 8192,
+        "default_dtype": "float16"
+    },
+    "opt-1.3b": {
+        "name": "facebook/opt-1.3b",
+        "requires_auth": False,
+        "max_context": 2048,
+        "default_dtype": "float16"
+    }
+}
+
+# Benchmark configurations - NO HARDCODING
+BENCHMARK_CONFIGS: Dict[str, Dict[str, Any]] = {
+    "perplexity": {
+        "type": "perplexity",
+        "default_samples": 50,
+        "default_prefill": 512,
+        "default_generation": 64
+    },
+    "niah": {
+        "type": "needle_in_haystack",
+        "depths": [10, 25, 50, 75, 90],  # Percentage depths
+        "needle": "The secret password is BANANA",
+        "default_samples": 10,
+        "default_context": 4096
+    },
+    "ruler": {
+        "type": "ruler",
+        "max_seq_lengths": [1024, 2048, 4096, 8192],
+        "default_samples": 10,
+        "default_n_facts": 10
+    },
+    "scbench": {
+        "type": "shared_context",
+        "num_turns": [5, 10, 20],
+        "default_samples": 10,
+        "default_context": 2048
+    },
+    "longbench": {
+        "type": "longbench",
+        "subsets": ["narrativeqa", "qasper", "multifieldqa_en", "hotpotqa", "2wikimqa"],
+        "default_samples": 20,
+        "max_context": 8192
+    }
+}
+
 class CompressionType(Enum):
     """RocketKV-enhanced SPG methods with explicit validation."""
     NONE = "none"
@@ -184,6 +247,9 @@ class EnhancedSPGConfig:
     stage_compression_min: float = 2.0    # Minimum stage compression ratio
     stage_compression_max: float = 500.0  # Maximum stage compression ratio (INCREASED for 450x)
     
+    # Flash Attention support
+    use_flash_attention: bool = False  # Try to use Flash Attention if available
+    
     def __post_init__(self):
         """Validate all parameters - fail fast on invalid config."""
         constants = ResearchConstants()
@@ -304,6 +370,10 @@ class CompressionConfig:
     compression_type: CompressionType = CompressionType.ENHANCED_SPG
     seed: int = 42
     
+    # Model selection
+    model_key: str = "gpt2"  # Key into SUPPORTED_MODELS
+    model_name: str = field(init=False)  # Will be set in __post_init__
+    
     # Enhanced SPG configuration
     enhanced_spg_config: EnhancedSPGConfig = field(default_factory=EnhancedSPGConfig)
     
@@ -327,10 +397,25 @@ class CompressionConfig:
     dataset_config: str = "wikitext-2-raw-v1"
     dataset_split: str = "test"
     
+    # Benchmark configuration
+    benchmark_type: str = "perplexity"  # perplexity, niah, ruler, scbench, longbench
+    benchmark_subset: Optional[str] = None  # For longbench subsets
+    
+    # NIAH-specific parameters
+    niah_needle: str = field(default_factory=lambda: BENCHMARK_CONFIGS["niah"]["needle"])
+    niah_depth_percent: float = 50.0
+    
+    # RULER-specific parameters
+    ruler_max_seq_length: int = 4096
+    
+    # SCBench-specific parameters
+    scbench_num_turns: int = 10
+    
     # Memory and system settings
     clear_cache_between_runs: bool = True
     use_memory_snapshot: bool = True
     fail_on_cpu_fallback: bool = True  # CHANGED: Default to True for strict compliance
+    use_flash_attention: bool = False  # Try to use Flash Attention if available
     
     # Output settings
     generate_latex: bool = True
@@ -347,6 +432,15 @@ class CompressionConfig:
         """Comprehensive validation - fail fast on any invalid parameter."""
         constants = ResearchConstants()
         
+        # Set model name from key
+        if self.model_key not in SUPPORTED_MODELS:
+            raise ValueError(f"model_key {self.model_key} not in SUPPORTED_MODELS: {list(SUPPORTED_MODELS.keys())}")
+        self.model_name = SUPPORTED_MODELS[self.model_key]["name"]
+        
+        # Validate benchmark type
+        if self.benchmark_type not in BENCHMARK_CONFIGS:
+            raise ValueError(f"benchmark_type {self.benchmark_type} not in BENCHMARK_CONFIGS: {list(BENCHMARK_CONFIGS.keys())}")
+        
         # Validate core parameters
         if not isinstance(self.seed, int) or self.seed < 0:
             raise ValueError(f"seed must be non-negative integer, got {self.seed}")
@@ -371,6 +465,9 @@ class CompressionConfig:
         if not 100 <= self.n_bootstrap <= 10000:
             logger.warning(f"n_bootstrap {self.n_bootstrap} outside recommended range [100, 10000]")
         
+        # Pass Flash Attention setting to EnhancedSPGConfig
+        self.enhanced_spg_config.use_flash_attention = self.use_flash_attention
+        
         logger.info("RocketKV-enhanced SPG config validated successfully")
     
     def to_json(self) -> str: