Upload runtime\difficulty_gate.py with huggingface_hub

runtime//difficulty_gate.py

@@ -0,0 +1,319 @@
+#!/usr/bin/env python
+"""
+Difficulty Gate for ContinuumAgent Project
+Smart routing system to determine whether to use patches based on query complexity
+"""
+
+import os
+import json
+from typing import Dict, Any, List, Optional, Tuple
+import numpy as np
+from llama_cpp import Llama
+
+class DifficultyGate:
+    """
+    Smart routing system to determine whether to use patches based on query complexity
+    Uses a simple heuristic approach for initial implementation, can be replaced with a learned classifier
+    """
+    
+    def __init__(self, 
+                 model_path: str,
+                 gate_threshold: float = 0.7,
+                 cache_dir: str = "models/gates",
+                 n_gpu_layers: int = 0):
+        """
+        Initialize the difficulty gate
+        
+        Args:
+            model_path: Path to GGUF model file
+            gate_threshold: Threshold for routing to patched model (0.0-1.0)
+            cache_dir: Directory for caching gate decisions
+            n_gpu_layers: Number of layers to offload to GPU
+        """
+        self.model_path = model_path
+        self.gate_threshold = gate_threshold
+        self.cache_dir = cache_dir
+        self.n_gpu_layers = n_gpu_layers
+        
+        # Create cache directory if it doesn't exist
+        os.makedirs(cache_dir, exist_ok=True)
+        
+        # Cache file path
+        self.cache_file = os.path.join(cache_dir, "gate_cache.json")
+        
+        # Load cache
+        self.decision_cache = self._load_cache()
+        
+        # Initialize gate model (small context for efficiency)
+        self._init_gate_model()
+        
+    def _init_gate_model(self) -> None:
+        """Initialize small gate model"""
+        try:
+            print(f"Loading gate model from {self.model_path}...")
+            self.gate_model = Llama(
+                model_path=self.model_path,
+                n_gpu_layers=self.n_gpu_layers,
+                n_ctx=512  # Small context for efficiency
+            )
+        except Exception as e:
+            print(f"Error loading gate model: {e}")
+            self.gate_model = None
+    
+    def _load_cache(self) -> Dict[str, Any]:
+        """
+        Load decision cache from file
+        
+        Returns:
+            Cache dictionary
+        """
+        if os.path.exists(self.cache_file):
+            try:
+                with open(self.cache_file, "r") as f:
+                    cache = json.load(f)
+                print(f"Loaded {len(cache.get('queries', []))} cached gate decisions")
+                return cache
+            except Exception as e:
+                print(f"Error loading cache: {e}")
+        
+        # Return empty cache
+        return {"queries": {}}
+    
+    def _save_cache(self) -> None:
+        """Save decision cache to file"""
+        try:
+            with open(self.cache_file, "w") as f:
+                json.dump(self.decision_cache, f, indent=2)
+        except Exception as e:
+            print(f"Error saving cache: {e}")
+    
+    def _query_hash(self, query: str) -> str:
+        """
+        Create simple hash for query caching
+        
+        Args:
+            query: Query string
+            
+        Returns:
+            Query hash
+        """
+        # Simple hash method, can be improved
+        import hashlib
+        return hashlib.md5(query.strip().lower().encode()).hexdigest()
+    
+    def _heuristic_features(self, query: str) -> Dict[str, float]:
+        """
+        Extract heuristic features from query
+        
+        Args:
+            query: Query string
+            
+        Returns:
+            Dictionary of feature values
+        """
+        # Lowercase query for consistent processing
+        query_lower = query.lower()
+        
+        # Feature 1: Query length
+        length = len(query)
+        norm_length = min(1.0, length / 200.0)  # Normalize to 0-1 (capped at 200 chars)
+        
+        # Feature 2: Presence of factual question indicators
+        factual_indicators = [
+            "what is", "when did", "where is", "who is", 
+            "which", "how many", "list the", "tell me about",
+            "explain", "define"
+        ]
+        has_factual = any(indicator in query_lower for indicator in factual_indicators)
+        
+        # Feature 3: Presence of time indicators (recency)
+        time_indicators = [
+            "recent", "latest", "current", "today", "now",
+            "this week", "this month", "this year",
+            "2023", "2024", "2025"  # Add current years
+        ]
+        has_time = any(indicator in query_lower for indicator in time_indicators)
+        
+        # Feature 4: Entity recognition (simplified)
+        # Check for capitalized terms that may indicate named entities
+        words = query.split()
+        capitalized_words = [w for w in words if w[0:1].isupper()]
+        entity_ratio = len(capitalized_words) / max(1, len(words))
+        
+        # Feature 5: Question complexity based on interrogative words
+        complex_indicators = [
+            "why", "how does", "explain", "compare", "contrast",
+            "what if", "analyze", "evaluate", "synthesize"
+        ]
+        complexity_score = sum(indicator in query_lower for indicator in complex_indicators) / 3.0
+        complexity_score = min(1.0, complexity_score)
+        
+        # Return features
+        return {
+            "length": norm_length,
+            "has_factual": float(has_factual),
+            "has_time": float(has_time),
+            "entity_ratio": entity_ratio,
+            "complexity": complexity_score
+        }
+        
+    def _heuristic_decision(self, features: Dict[str, float]) -> Tuple[bool, float]:
+        """
+        Make decision based on heuristic features
+        
+        Args:
+            features: Feature dictionary
+            
+        Returns:
+            Tuple of (needs_patches, confidence)
+        """
+        # Weights for different features
+        weights = {
+            "length": 0.1,
+            "has_factual": 0.3,
+            "has_time": 0.4,  # Highest weight for time indicators
+            "entity_ratio": 0.1,
+            "complexity": -0.1  # Negative weight - complex reasoning queries may not need patches
+        }
+        
+        # Calculate weighted score
+        score = sum(features[f] * weights[f] for f in features)
+        
+        # Normalize to 0-1 range
+        score = max(0.0, min(1.0, score))
+        
+        # Decision based on threshold
+        needs_patches = score >= self.gate_threshold
+        
+        return needs_patches, score
+        
+    def _model_decision(self, query: str) -> Tuple[bool, float]:
+        """
+        Ask the model to decide if the query needs up-to-date knowledge
+        
+        Args:
+            query: Query string
+            
+        Returns:
+            Tuple of (needs_patches, confidence)
+        """
+        if not self.gate_model:
+            # Fall back to heuristic if model not available
+            features = self._heuristic_features(query)
+            return self._heuristic_decision(features)
+        
+        # Prompt for model
+        prompt = f"""<s>[INST] Please analyze this question and determine if it requires the most up-to-date knowledge to answer correctly. 
+Respond with only a single word: 'YES' if up-to-date knowledge is needed, or 'NO' if it can be answered with general knowledge.
+
+Question: "{query}"
+
+Requires up-to-date knowledge? [/INST]"""
+
+        # Generate completion
+        completion = self.gate_model.create_completion(
+            prompt=prompt,
+            max_tokens=5,
+            temperature=0.1,  # Low temperature for consistent results
+            stop=["</s>", "\n"]
+        )
+        
+        # Extract response
+        response_text = completion.get("choices", [{}])[0].get("text", "").strip().upper()
+        
+        # Calculate confidence from logprobs if available
+        confidence = 0.7  # Default confidence
+        
+        # Decision based on response
+        needs_patches = "YES" in response_text
+        
+        return needs_patches, confidence
+    
+    def should_use_patches(self, query: str, use_model: bool = True) -> Dict[str, Any]:
+        """
+        Determine if the query requires up-to-date knowledge patches
+        
+        Args:
+            query: Query string
+            use_model: Whether to use model for decision (vs pure heuristics)
+            
+        Returns:
+            Decision dictionary with keys:
+              - needs_patches: Boolean decision
+              - confidence: Confidence score (0.0-1.0)
+              - method: Decision method used
+              - features: Feature values if heuristic method used
+        """
+        # Check cache first
+        query_hash = self._query_hash(query)
+        if query_hash in self.decision_cache.get("queries", {}):
+            cached = self.decision_cache["queries"][query_hash]
+            cached["from_cache"] = True
+            return cached
+        
+        # Extract features
+        features = self._heuristic_features(query)
+        
+        # Make decision
+        if use_model and self.gate_model:
+            needs_patches, confidence = self._model_decision(query)
+            method = "model"
+        else:
+            needs_patches, confidence = self._heuristic_decision(features)
+            method = "heuristic"
+        
+        # Create decision
+        decision = {
+            "needs_patches": needs_patches,
+            "confidence": confidence,
+            "method": method,
+            "features": features,
+            "from_cache": False
+        }
+        
+        # Cache decision
+        self.decision_cache.setdefault("queries", {})[query_hash] = decision
+        self._save_cache()
+        
+        return decision
+
+def main():
+    """Test difficulty gate"""
+    # Find model path
+    model_dir = "models/slow"
+    model_files = [f for f in os.listdir(model_dir) if f.endswith(".gguf")]
+    
+    if not model_files:
+        print(f"No GGUF models found in {model_dir}")
+        return
+    
+    model_path = os.path.join(model_dir, model_files[0])
+    print(f"Using model: {model_path}")
+    
+    # Initialize gate
+    gate = DifficultyGate(model_path=model_path)
+    
+    # Test queries
+    test_queries = [
+        "What is the capital of France?",
+        "Who is the current president of the United States?",
+        "Explain the theory of relativity",
+        "What are the latest developments in the conflict in Ukraine?",
+        "Who won the most recent Super Bowl?",
+        "How do I write a for loop in Python?"
+    ]
+    
+    for query in test_queries:
+        # Test heuristic decision
+        decision = gate.should_use_patches(query, use_model=False)
+        print(f"\nQuery: {query}")
+        print(f"Heuristic Decision: {decision['needs_patches']} (Confidence: {decision['confidence']:.2f})")
+        print(f"Features: {decision['features']}")
+        
+        # Test model decision if model is available
+        if gate.gate_model:
+            decision = gate.should_use_patches(query, use_model=True)
+            print(f"Model Decision: {decision['needs_patches']} (Confidence: {decision['confidence']:.2f})")
+
+if __name__ == "__main__":
+    main()