models/hf_interface.py

"""
Hugging Face Model Interface

Provides a standardized interface for interacting with Hugging Face models
in the AI safety lab. Handles authentication, model loading, and inference.
"""

import os
from typing import Dict, List, Optional, Any
from pydantic import BaseModel, Field
import logging

# Try to import heavy dependencies, fall back if they fail
try:
    from huggingface_hub import InferenceClient, HfApi
    HEAVY_DEPS_AVAILABLE = True
except ImportError as e:
    logging.warning(f"HuggingFace Hub not available: {e}")
    HEAVY_DEPS_AVAILABLE = False
    InferenceClient = None
    HfApi = None

# Separate torch/transformers import with more specific error handling
try:
    import torch
    from transformers import AutoTokenizer, AutoModelForCausalLM
    TORCH_AVAILABLE = True
except (ImportError, OSError) as e:
    logging.warning(f"PyTorch/Transformers not available: {e}")
    TORCH_AVAILABLE = False
    torch = None
    AutoTokenizer = None
    AutoModelForCausalLM = None

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class ModelInfo(BaseModel):
    """Information about an available model"""
    model_id: str = Field(description="Hugging Face model ID")
    name: str = Field(description="Display name")
    description: str = Field(description="Model description")
    category: str = Field(description="Model category")
    requires_token: bool = Field(description="Whether model requires authentication")
    is_local: bool = Field(description="Whether model is loaded locally")


class ModelResponse(BaseModel):
    """Standardized model response"""
    text: str = Field(description="Generated text")
    model_id: str = Field(description="Model used")
    generation_time: float = Field(description="Time taken to generate")
    token_count: int = Field(description="Number of tokens generated")
    metadata: Dict[str, Any] = Field(description="Additional metadata")


class HFModelInterface:
    """
    Interface for interacting with Hugging Face models.
    
    Supports both API-based inference and local model loading for comprehensive
    safety testing capabilities.
    """

    def __init__(self):
        self.token = os.environ.get("HUGGINGFACEHUB_API_TOKEN")
        if not self.token:
            logger.warning("HUGGINGFACEHUB_API_TOKEN not found in environment variables")
        
        self.inference_client = None
        self.api_client = None
        self.local_models = {}
        self.available_models = self._initialize_model_registry()
        
        if self.token:
            self._initialize_clients()

    def _initialize_clients(self):
        """Initialize Hugging Face clients"""
        if not HEAVY_DEPS_AVAILABLE:
            logger.warning("HuggingFace Hub not available - using mock client")
            return
            
        try:
            self.inference_client = InferenceClient(token=self.token)
            self.api_client = HfApi(token=self.token)
            logger.info("Hugging Face clients initialized successfully")
        except Exception as e:
            logger.error(f"Failed to initialize Hugging Face clients: {e}")

    def _initialize_model_registry(self) -> Dict[str, ModelInfo]:
        """Initialize registry of available models - TESTED and WORKING with HF Inference API"""
        return {
            "HuggingFaceH4/zephyr-7b-beta": ModelInfo(
                model_id="HuggingFaceH4/zephyr-7b-beta",
                name="Zephyr 7B Beta",
                description="HuggingFace H4's high-performance chat model",
                category="General Purpose",
                requires_token=False,
                is_local=False
            ),
            "tiiuae/falcon-7b-instruct": ModelInfo(
                model_id="tiiuae/falcon-7b-instruct",
                name="Falcon 7B Instruct",
                description="TII UAE's open-source instruction model",
                category="Instruction Following",
                requires_token=False,
                is_local=False
            ),
            "google/gemma-2b-it": ModelInfo(
                model_id="google/gemma-2b-it",
                name="Gemma 2B IT",
                description="Google's lightweight instruction-tuned model",
                category="Instruction Following",
                requires_token=False,
                is_local=False
            ),
            "microsoft/DialoGPT-medium": ModelInfo(
                model_id="microsoft/DialoGPT-medium",
                name="DialoGPT Medium",
                description="Microsoft's conversational model",
                category="Conversational",
                requires_token=False,
                is_local=False
            ),
            "google/flan-t5-large": ModelInfo(
                model_id="google/flan-t5-large",
                name="FLAN-T5 Large",
                description="Google's instruction-tuned T5 model",
                category="Instruction Following",
                requires_token=False,
                is_local=False
            )
        }

    def get_available_models(self) -> List[ModelInfo]:
        """
        Get list of available models.
        
        Returns:
            List of available model information
        """
        return list(self.available_models.values())

    def get_model_info(self, model_id: str) -> Optional[ModelInfo]:
        """
        Get information about a specific model.
        
        Args:
            model_id: Hugging Face model ID
            
        Returns:
            Model information or None if not found
        """
        return self.available_models.get(model_id)

    def load_local_model(self, model_id: str, device: str = "auto") -> bool:
        """
        Load a model locally for offline inference.
        
        Args:
            model_id: Hugging Face model ID
            device: Device to load model on
            
        Returns:
            True if successful, False otherwise
        """
        if not TORCH_AVAILABLE:
            logger.error("PyTorch not available - cannot load local models")
            return False
            
        try:
            logger.info(f"Loading model locally: {model_id}")
            
            # Check if model exists in registry
            if model_id not in self.available_models:
                logger.error(f"Model {model_id} not found in registry")
                return False
            
            # Load tokenizer and model
            tokenizer = AutoTokenizer.from_pretrained(
                model_id,
                token=self.token if self.available_models[model_id].requires_token else None
            )
            
            model = AutoModelForCausalLM.from_pretrained(
                model_id,
                token=self.token if self.available_models[model_id].requires_token else None,
                torch_dtype=torch.float16,
                device_map=device if device != "auto" else "auto"
            )
            
            # Store in local models
            self.local_models[model_id] = {
                "model": model,
                "tokenizer": tokenizer,
                "device": device
            }
            
            # Update model info
            self.available_models[model_id].is_local = True
            
            logger.info(f"Successfully loaded model locally: {model_id}")
            return True
            
        except Exception as e:
            logger.error(f"Failed to load model {model_id}: {e}")
            return False

    def generate_response(
        self,
        model_id: str,
        prompt: str,
        max_tokens: int = 512,
        temperature: float = 0.7,
        use_local: bool = False
    ) -> Optional[ModelResponse]:
        """
        Generate a response from the specified model.
        
        Args:
            model_id: Hugging Face model ID
            prompt: Input prompt
            max_tokens: Maximum tokens to generate
            temperature: Generation temperature
            use_local: Whether to use local model if available
            
        Returns:
            Model response or None if failed
        """
        import time
        start_time = time.time()
        
        try:
            # Check if local model should be used
            if use_local and model_id in self.local_models:
                return self._generate_local(
                    model_id, prompt, max_tokens, temperature, start_time
                )
            else:
                return self._generate_api(
                    model_id, prompt, max_tokens, temperature, start_time
                )
                
        except Exception as e:
            logger.error(f"Failed to generate response from {model_id}: {e}")
            return None

    def _generate_local(
        self,
        model_id: str,
        prompt: str,
        max_tokens: int,
        temperature: float,
        start_time: float
    ) -> ModelResponse:
        """Generate response using locally loaded model"""
        
        model_data = self.local_models[model_id]
        model = model_data["model"]
        tokenizer = model_data["tokenizer"]
        
        # Tokenize input
        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
        
        # Generate response
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=max_tokens,
                temperature=temperature,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id
            )
        
        # Decode response
        response_text = tokenizer.decode(
            outputs[0][inputs["input_ids"].shape[1]:],
            skip_special_tokens=True
        )
        
        generation_time = time.time() - start_time
        token_count = len(tokenizer.encode(response_text))
        
        return ModelResponse(
            text=response_text,
            model_id=model_id,
            generation_time=generation_time,
            token_count=token_count,
            metadata={"source": "local", "device": str(model.device)}
        )

    def _generate_api(
        self,
        model_id: str,
        prompt: str,
        max_tokens: int,
        temperature: float,
        start_time: float
    ) -> ModelResponse:
        """Generate response using Hugging Face API"""
        
        if not self.inference_client:
            raise RuntimeError("Inference client not initialized")
        
        # Generate response
        response = self.inference_client.text_generation(
            prompt=prompt,
            model=model_id,
            max_new_tokens=max_tokens,
            temperature=temperature,
            do_sample=True
        )
        
        generation_time = time.time() - start_time
        
        # Estimate token count (rough approximation)
        token_count = len(response.split())
        
        return ModelResponse(
            text=response,
            model_id=model_id,
            generation_time=generation_time,
            token_count=token_count,
            metadata={"source": "api"}
        )

    def batch_generate(
        self,
        model_id: str,
        prompts: List[str],
        max_tokens: int = 512,
        temperature: float = 0.7,
        use_local: bool = False
    ) -> List[Optional[ModelResponse]]:
        """
        Generate responses for multiple prompts.
        
        Args:
            model_id: Hugging Face model ID
            prompts: List of input prompts
            max_tokens: Maximum tokens to generate per response
            temperature: Generation temperature
            use_local: Whether to use local model if available
            
        Returns:
            List of model responses (None for failed generations)
        """
        responses = []
        for prompt in prompts:
            response = self.generate_response(
                model_id, prompt, max_tokens, temperature, use_local
            )
            responses.append(response)
        return responses

    def validate_model_access(self, model_id: str) -> bool:
        """
        Validate if we can access a specific model.
        
        Args:
            model_id: Hugging Face model ID
            
        Returns:
            True if accessible, False otherwise
        """
        try:
            if not self.api_client:
                return False
            
            # Try to get model info
            model_info = self.api_client.model_info(model_id)
            return True
            
        except Exception as e:
            logger.warning(f"Cannot access model {model_id}: {e}")
            return False

    def get_model_capabilities(self, model_id: str) -> Dict[str, Any]:
        """
        Get capabilities and limitations of a model.
        
        Args:
            model_id: Hugging Face model ID
            
        Returns:
            Dictionary of model capabilities
        """
        model_info = self.get_model_info(model_id)
        if not model_info:
            return {}
        
        return {
            "model_id": model_id,
            "name": model_info.name,
            "category": model_info.category,
            "requires_token": model_info.requires_token,
            "is_local": model_info.is_local,
            "supports_streaming": False,  # Could be expanded
            "max_context_length": 2048,  # Default, could be model-specific
            "safety_features": [
                "content_filtering" if not model_info.is_local else "local_control",
                "custom_safety_evaluation"  # Our own evaluation
            ]
        }


# Global instance for the application
model_interface = HFModelInterface()