""" 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()