Create safety_wrapper.py

safety_wrapper.py

@@ -0,0 +1,433 @@
+"""
+Safety wrapper for Helion-V2.0-Thinking
+Implements content filtering, rate limiting, and safety checks
+"""
+
+import torch
+from transformers import AutoModelForCausalLM, AutoProcessor
+from typing import Optional, List, Dict, Any, Union
+from PIL import Image
+import json
+import re
+import time
+from collections import defaultdict, deque
+from datetime import datetime, timedelta
+import hashlib
+
+
+class SafetyViolation(Exception):
+    """Exception raised when safety policies are violated"""
+    pass
+
+
+class ContentFilter:
+    """Content filtering for inputs and outputs"""
+    
+    # Harmful patterns to detect
+    HARMFUL_PATTERNS = [
+        r'(?i)how\s+to\s+(hack|crack|break\s+into)',
+        r'(?i)make\s+(explosive|bomb|weapon)',
+        r'(?i)(kill|murder|harm)\s+(myself|someone|people)',
+        r'(?i)credit\s+card\s+number',
+        r'(?i)social\s+security\s+number',
+        r'(?i)(steal|fraud|scam)\s+',
+        r'(?i)illegal\s+(drugs|substances)',
+        r'(?i)child\s+(abuse|exploitation)',
+    ]
+    
+    # PII patterns
+    PII_PATTERNS = {
+        'email': r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b',
+        'phone': r'\b\d{3}[-.]?\d{3}[-.]?\d{4}\b',
+        'ssn': r'\b\d{3}-\d{2}-\d{4}\b',
+        'credit_card': r'\b\d{4}[\s-]?\d{4}[\s-]?\d{4}[\s-]?\d{4}\b',
+    }
+    
+    # Toxic keywords
+    TOXIC_KEYWORDS = [
+        'hate', 'violence', 'threat', 'abuse', 'harass',
+        'discriminate', 'racist', 'sexist', 'offensive'
+    ]
+    
+    @staticmethod
+    def check_harmful_content(text: str) -> tuple[bool, Optional[str]]:
+        """
+        Check if text contains harmful content
+        
+        Returns:
+            (is_harmful, reason)
+        """
+        for pattern in ContentFilter.HARMFUL_PATTERNS:
+            if re.search(pattern, text):
+                return True, f"Matched harmful pattern: {pattern}"
+        
+        return False, None
+    
+    @staticmethod
+    def check_pii(text: str) -> tuple[bool, List[str]]:
+        """
+        Check for personally identifiable information
+        
+        Returns:
+            (contains_pii, pii_types)
+        """
+        found_pii = []
+        for pii_type, pattern in ContentFilter.PII_PATTERNS.items():
+            if re.search(pattern, text):
+                found_pii.append(pii_type)
+        
+        return len(found_pii) > 0, found_pii
+    
+    @staticmethod
+    def check_toxicity(text: str) -> tuple[float, List[str]]:
+        """
+        Simple toxicity check based on keywords
+        
+        Returns:
+            (toxicity_score, matched_keywords)
+        """
+        text_lower = text.lower()
+        matched = [kw for kw in ContentFilter.TOXIC_KEYWORDS if kw in text_lower]
+        score = len(matched) / len(ContentFilter.TOXIC_KEYWORDS)
+        
+        return score, matched
+    
+    @staticmethod
+    def redact_pii(text: str) -> str:
+        """Redact PII from text"""
+        redacted = text
+        
+        for pii_type, pattern in ContentFilter.PII_PATTERNS.items():
+            redacted = re.sub(pattern, f'[REDACTED_{pii_type.upper()}]', redacted)
+        
+        return redacted
+
+
+class RateLimiter:
+    """Rate limiting for API usage"""
+    
+    def __init__(
+        self,
+        requests_per_minute: int = 60,
+        tokens_per_minute: int = 90000,
+        concurrent_limit: int = 10
+    ):
+        self.requests_per_minute = requests_per_minute
+        self.tokens_per_minute = tokens_per_minute
+        self.concurrent_limit = concurrent_limit
+        
+        self.request_times = defaultdict(deque)
+        self.token_counts = defaultdict(deque)
+        self.active_requests = defaultdict(int)
+    
+    def check_rate_limit(self, user_id: str, estimated_tokens: int = 0) -> tuple[bool, Optional[str]]:
+        """
+        Check if request is within rate limits
+        
+        Returns:
+            (allowed, reason_if_denied)
+        """
+        now = datetime.now()
+        minute_ago = now - timedelta(minutes=1)
+        
+        # Clean old entries
+        while self.request_times[user_id] and self.request_times[user_id][0] < minute_ago:
+            self.request_times[user_id].popleft()
+        
+        while self.token_counts[user_id] and self.token_counts[user_id][0][0] < minute_ago:
+            self.token_counts[user_id].popleft()
+        
+        # Check requests per minute
+        if len(self.request_times[user_id]) >= self.requests_per_minute:
+            return False, f"Rate limit exceeded: {self.requests_per_minute} requests per minute"
+        
+        # Check tokens per minute
+        total_tokens = sum(t[1] for t in self.token_counts[user_id])
+        if total_tokens + estimated_tokens > self.tokens_per_minute:
+            return False, f"Token limit exceeded: {self.tokens_per_minute} tokens per minute"
+        
+        # Check concurrent requests
+        if self.active_requests[user_id] >= self.concurrent_limit:
+            return False, f"Concurrent request limit exceeded: {self.concurrent_limit}"
+        
+        return True, None
+    
+    def record_request(self, user_id: str, tokens: int = 0):
+        """Record a request"""
+        now = datetime.now()
+        self.request_times[user_id].append(now)
+        self.token_counts[user_id].append((now, tokens))
+        self.active_requests[user_id] += 1
+    
+    def release_request(self, user_id: str):
+        """Release an active request slot"""
+        if self.active_requests[user_id] > 0:
+            self.active_requests[user_id] -= 1
+
+
+class SafeHelionWrapper:
+    """Safety wrapper for Helion-V2.0-Thinking"""
+    
+    def __init__(
+        self,
+        model_name: str = "DeepXR/Helion-V2.0-Thinking",
+        safety_config_path: Optional[str] = None,
+        enable_safety: bool = True,
+        enable_rate_limiting: bool = True,
+        device: str = "auto"
+    ):
+        """
+        Initialize safe wrapper
+        
+        Args:
+            model_name: Model name or path
+            safety_config_path: Path to safety_config.json
+            enable_safety: Enable safety checks
+            enable_rate_limiting: Enable rate limiting
+            device: Device for model
+        """
+        print(f"Loading model with safety wrapper: {model_name}")
+        
+        # Load safety config
+        if safety_config_path:
+            with open(safety_config_path, 'r') as f:
+                self.safety_config = json.load(f)
+        else:
+            self.safety_config = self._default_safety_config()
+        
+        self.enable_safety = enable_safety
+        self.enable_rate_limiting = enable_rate_limiting
+        
+        # Initialize components
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            torch_dtype=torch.bfloat16,
+            device_map=device,
+            trust_remote_code=True
+        )
+        self.processor = AutoProcessor.from_pretrained(model_name)
+        self.model.eval()
+        
+        self.content_filter = ContentFilter()
+        self.rate_limiter = RateLimiter(
+            requests_per_minute=self.safety_config['safety_settings']['rate_limiting']['requests_per_minute'],
+            tokens_per_minute=self.safety_config['safety_settings']['rate_limiting']['tokens_per_minute'],
+            concurrent_requests=self.safety_config['safety_settings']['rate_limiting']['concurrent_requests']
+        )
+        
+        # Violation tracking
+        self.violation_log = []
+        
+        print("Safety wrapper initialized successfully")
+    
+    def _default_safety_config(self) -> Dict[str, Any]:
+        """Default safety configuration"""
+        return {
+            "safety_settings": {
+                "rate_limiting": {
+                    "requests_per_minute": 60,
+                    "tokens_per_minute": 90000,
+                    "concurrent_requests": 10
+                },
+                "content_filtering": {
+                    "profanity_filter": {"enabled": True},
+                    "pii_detection": {"enabled": True},
+                    "toxicity_detection": {"enabled": True, "threshold": 0.7}
+                }
+            }
+        }
+    
+    def _validate_input(
+        self,
+        prompt: str,
+        images: Optional[List[Image.Image]] = None,
+        user_id: str = "default"
+    ):
+        """Validate input against safety policies"""
+        if not self.enable_safety:
+            return
+        
+        # Check harmful content
+        is_harmful, reason = self.content_filter.check_harmful_content(prompt)
+        if is_harmful:
+            self._log_violation(user_id, "harmful_content", reason)
+            raise SafetyViolation(f"Input rejected: {reason}")
+        
+        # Check PII
+        if self.safety_config['safety_settings']['content_filtering']['pii_detection']['enabled']:
+            has_pii, pii_types = self.content_filter.check_pii(prompt)
+            if has_pii:
+                self._log_violation(user_id, "pii_detected", f"Types: {pii_types}")
+                print(f"Warning: PII detected in input: {pii_types}")
+        
+        # Check toxicity
+        if self.safety_config['safety_settings']['content_filtering']['toxicity_detection']['enabled']:
+            toxicity_score, keywords = self.content_filter.check_toxicity(prompt)
+            threshold = self.safety_config['safety_settings']['content_filtering']['toxicity_detection']['threshold']
+            
+            if toxicity_score > threshold:
+                self._log_violation(user_id, "high_toxicity", f"Score: {toxicity_score}")
+                raise SafetyViolation(f"Input rejected: High toxicity score ({toxicity_score:.2f})")
+        
+        # Validate images
+        if images:
+            max_images = self.safety_config.get('guardrails', {}).get('input_validation', {}).get('max_images_per_request', 10)
+            if len(images) > max_images:
+                raise SafetyViolation(f"Too many images: {len(images)} (max: {max_images})")
+    
+    def _validate_output(self, output: str, user_id: str = "default"):
+        """Validate output against safety policies"""
+        if not self.enable_safety:
+            return output
+        
+        # Check for harmful content in output
+        is_harmful, reason = self.content_filter.check_harmful_content(output)
+        if is_harmful:
+            self._log_violation(user_id, "harmful_output", reason)
+            return "I cannot provide that information as it may be harmful."
+        
+        # Redact PII if found
+        if self.safety_config['safety_settings']['content_filtering']['pii_detection']['enabled']:
+            output = self.content_filter.redact_pii(output)
+        
+        return output
+    
+    def _log_violation(self, user_id: str, violation_type: str, details: str):
+        """Log safety violation"""
+        self.violation_log.append({
+            "timestamp": datetime.now().isoformat(),
+            "user_id": hashlib.sha256(user_id.encode()).hexdigest()[:16],
+            "type": violation_type,
+            "details": details
+        })
+        
+        # Keep only last 1000 violations
+        if len(self.violation_log) > 1000:
+            self.violation_log = self.violation_log[-1000:]
+    
+    def generate(
+        self,
+        prompt: str,
+        images: Optional[List[Image.Image]] = None,
+        user_id: str = "default",
+        max_new_tokens: int = 512,
+        temperature: float = 0.7,
+        **kwargs
+    ) -> str:
+        """
+        Safe generation with input/output filtering
+        
+        Args:
+            prompt: Input prompt
+            images: Optional list of images
+            user_id: User identifier for rate limiting
+            max_new_tokens: Maximum tokens to generate
+            temperature: Sampling temperature
+            **kwargs: Additional generation parameters
+        
+        Returns:
+            Generated text (filtered)
+        """
+        # Check rate limit
+        if self.enable_rate_limiting:
+            allowed, reason = self.rate_limiter.check_rate_limit(user_id, max_new_tokens)
+            if not allowed:
+                raise SafetyViolation(reason)
+            
+            self.rate_limiter.record_request(user_id, max_new_tokens)
+        
+        try:
+            # Validate input
+            self._validate_input(prompt, images, user_id)
+            
+            # Generate
+            if images:
+                inputs = self.processor(text=prompt, images=images, return_tensors="pt").to(self.model.device)
+            else:
+                inputs = self.processor(text=prompt, return_tensors="pt").to(self.model.device)
+            
+            with torch.no_grad():
+                outputs = self.model.generate(
+                    **inputs,
+                    max_new_tokens=max_new_tokens,
+                    temperature=temperature,
+                    **kwargs
+                )
+            
+            response = self.processor.decode(outputs[0], skip_special_tokens=True)
+            
+            # Remove prompt from response
+            if response.startswith(prompt):
+                response = response[len(prompt):].strip()
+            
+            # Validate output
+            response = self._validate_output(response, user_id)
+            
+            return response
+        
+        finally:
+            if self.enable_rate_limiting:
+                self.rate_limiter.release_request(user_id)
+    
+    def get_violation_stats(self) -> Dict[str, Any]:
+        """Get violation statistics"""
+        if not self.violation_log:
+            return {"total_violations": 0}
+        
+        violation_types = defaultdict(int)
+        for log in self.violation_log:
+            violation_types[log['type']] += 1
+        
+        return {
+            "total_violations": len(self.violation_log),
+            "by_type": dict(violation_types),
+            "recent_violations": self.violation_log[-10:]
+        }
+    
+    def export_violation_log(self, filename: str = "violations.json"):
+        """Export violation log to file"""
+        with open(filename, 'w') as f:
+            json.dump(self.violation_log, f, indent=2)
+        print(f"Violation log exported to {filename}")
+
+
+def main():
+    """Example usage of safe wrapper"""
+    # Initialize safe wrapper
+    wrapper = SafeHelionWrapper(
+        model_name="DeepXR/Helion-V2.0-Thinking",
+        enable_safety=True,
+        enable_rate_limiting=True
+    )
+    
+    # Test cases
+    test_prompts = [
+        "Explain how photosynthesis works.",  # Safe
+        "Write a poem about nature.",  # Safe
+        "How do I hack into an email account?",  # Should be blocked
+    ]
+    
+    print("\n" + "="*60)
+    print("Testing Safety Wrapper")
+    print("="*60 + "\n")
+    
+    for prompt in test_prompts:
+        print(f"Prompt: {prompt}")
+        try:
+            response = wrapper.generate(prompt, max_new_tokens=128)
+            print(f"Response: {response}\n")
+        except SafetyViolation as e:
+            print(f"BLOCKED: {e}\n")
+        except Exception as e:
+            print(f"ERROR: {e}\n")
+    
+    # Print violation stats
+    print("="*60)
+    print("Violation Statistics")
+    print("="*60)
+    stats = wrapper.get_violation_stats()
+    print(json.dumps(stats, indent=2))
+
+
+if __name__ == "__main__":
+    main()