agent-learning-kit 0.1.0__py3-none-any.whl

fi/evals/guardrails/backends/local_base.py

@@ -0,0 +1,479 @@
+"""
+Base class for Local Model Backends.
+
+Provides common functionality for running HuggingFace models locally,
+either via VLLM server or direct transformers loading.
+"""
+
+import os
+import time
+from abc import abstractmethod
+from typing import Any, Dict, List, Optional
+
+from fi.evals.guardrails.backends.base import BaseBackend
+from fi.evals.guardrails.backends.vllm_client import VLLMClient, get_vllm_url
+from fi.evals.guardrails.config import GuardrailModel, RailType
+from fi.evals.guardrails.types import GuardrailResult
+
+
+class LocalModelBackend(BaseBackend):
+    """
+    Base class for local HuggingFace model backends.
+
+    Supports two modes:
+    1. VLLM Server Mode: Use external VLLM server for inference
+    2. Direct Mode: Load model directly using transformers
+
+    Subclasses must implement:
+    - _format_prompt(): Model-specific prompt formatting
+    - _parse_response(): Model-specific response parsing
+    - HF_MODEL_NAME: HuggingFace model identifier
+
+    Usage:
+        class WildGuardBackend(LocalModelBackend):
+            HF_MODEL_NAME = "allenai/wildguard"
+
+            def _format_prompt(self, content, context=None):
+                return f"Human user:\\n{content}\\n..."
+
+            def _parse_response(self, response):
+                # Parse "harmful_request: yes/no" format
+                ...
+    """
+
+    # Subclasses should override these
+    HF_MODEL_NAME: str = ""
+    MAX_NEW_TOKENS: int = 128
+    TEMPERATURE: float = 0.1
+
+    def __init__(
+        self,
+        model: GuardrailModel,
+        vllm_url: Optional[str] = None,
+        device: str = "auto",
+        hf_token: Optional[str] = None,
+        load_in_8bit: bool = False,
+        load_in_4bit: bool = False,
+    ):
+        """
+        Initialize local model backend.
+
+        Args:
+            model: The guardrail model enum value
+            vllm_url: VLLM server URL (auto-detected from env if not provided)
+            device: Device to use ("auto", "cuda", "mps", "cpu")
+            hf_token: HuggingFace token for gated models
+            load_in_8bit: Use 8-bit quantization (reduces memory)
+            load_in_4bit: Use 4-bit quantization (reduces memory further)
+        """
+        super().__init__(model)
+
+        self._vllm_url = vllm_url or get_vllm_url(model.value)
+        self._device = device
+        self._hf_token = hf_token or os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
+        self._load_in_8bit = load_in_8bit
+        self._load_in_4bit = load_in_4bit
+
+        # Lazy-loaded resources
+        self._vllm_client: Optional[VLLMClient] = None
+        self._transformers_model = None
+        self._tokenizer = None
+
+        # Check what's available
+        self._use_vllm = False
+        if self._vllm_url:
+            client = VLLMClient(self._vllm_url)
+            if client.health_check():
+                self._use_vllm = True
+                self._vllm_client = client
+                # Resolve which model on the server matches this backend
+                self._resolve_server_model()
+
+    def _resolve_server_model(self):
+        """Match this backend's model to an available model on the VLLM/ollama server."""
+        try:
+            available = self._vllm_client.get_models()
+        except Exception:
+            return
+
+        if not available:
+            return
+
+        # Our enum value e.g. "llamaguard-3-1b", "qwen3guard-0.6b"
+        our_key = self.model.value.lower()
+        # Normalize: strip hyphens/underscores/dots for fuzzy matching
+        our_norm = our_key.replace("-", "").replace("_", "").replace(".", "")
+        # Also get the HF model name for matching
+        hf_name = self._get_hf_model_name().lower()
+
+        for server_model in available:
+            sm = server_model.lower()
+            sm_norm = sm.replace("-", "").replace("_", "").replace(".", "").replace("/", "").replace(":", "")
+
+            # Exact match on enum value
+            if our_key in sm or sm.startswith(our_key):
+                self._vllm_client.model = server_model
+                return
+
+            # Normalized fuzzy match (e.g. "llamaguard31b" in "llamaguard31b")
+            if our_norm in sm_norm or sm_norm in our_norm:
+                self._vllm_client.model = server_model
+                return
+
+            # HF name match
+            if hf_name and (hf_name in sm or sm in hf_name):
+                self._vllm_client.model = server_model
+                return
+
+        # No match found — VLLMClient will fall back to server default
+
+    def _get_hf_model_name(self) -> str:
+        """Get the HuggingFace model name. Override in subclasses if needed."""
+        return self.HF_MODEL_NAME
+
+    def _load_transformers_model(self):
+        """Load the model using transformers library."""
+        if self._transformers_model is not None:
+            return
+
+        try:
+            import torch
+            from transformers import AutoModelForCausalLM, AutoTokenizer
+        except ImportError:
+            raise ImportError(
+                "transformers and torch packages required for local model backends. "
+                "Install with: pip install torch transformers accelerate"
+            )
+
+        model_name = self._get_hf_model_name()
+
+        # Determine device
+        if self._device == "auto":
+            if torch.cuda.is_available():
+                device = "cuda"
+            elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+                device = "mps"
+            else:
+                device = "cpu"
+        else:
+            device = self._device
+
+        # Load tokenizer
+        self._tokenizer = AutoTokenizer.from_pretrained(
+            model_name,
+            token=self._hf_token,
+        )
+
+        # Determine dtype and load kwargs
+        if device in ("cuda", "mps"):
+            dtype = torch.float16
+        else:
+            dtype = torch.float32
+
+        load_kwargs = {
+            "torch_dtype": dtype,
+            "token": self._hf_token,
+            "low_cpu_mem_usage": True,
+        }
+
+        if self._load_in_8bit:
+            load_kwargs["load_in_8bit"] = True
+        elif self._load_in_4bit:
+            load_kwargs["load_in_4bit"] = True
+
+        # Load model
+        self._transformers_model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            **load_kwargs,
+        )
+
+        if not (self._load_in_8bit or self._load_in_4bit):
+            self._transformers_model.to(device)
+
+        self._transformers_model.eval()
+        self._transformers_device = device
+
+    def _generate_with_transformers(
+        self,
+        prompt: str,
+        max_new_tokens: Optional[int] = None,
+        temperature: Optional[float] = None,
+    ) -> str:
+        """
+        Generate text using local transformers model.
+
+        Args:
+            prompt: Input prompt
+            max_new_tokens: Max tokens to generate
+            temperature: Sampling temperature
+
+        Returns:
+            Generated text (response only, not including prompt)
+        """
+        import torch
+
+        self._load_transformers_model()
+
+        max_tokens = max_new_tokens or self.MAX_NEW_TOKENS
+        temp = temperature or self.TEMPERATURE
+
+        inputs = self._tokenizer(prompt, return_tensors="pt").to(self._transformers_device)
+        prompt_length = inputs.input_ids.shape[1]
+
+        with torch.no_grad():
+            outputs = self._transformers_model.generate(
+                **inputs,
+                max_new_tokens=max_tokens,
+                temperature=max(temp, 0.01),
+                do_sample=temp > 0,
+                pad_token_id=self._tokenizer.eos_token_id,
+            )
+
+        # Decode only the generated tokens (excluding prompt)
+        generated_tokens = outputs[0][prompt_length:]
+        response = self._tokenizer.decode(generated_tokens, skip_special_tokens=True)
+
+        return response
+
+    @abstractmethod
+    def _format_prompt(
+        self,
+        content: str,
+        rail_type: RailType,
+        context: Optional[str] = None,
+    ) -> str:
+        """
+        Format content into model-specific prompt.
+
+        Args:
+            content: Content to classify
+            rail_type: Type of rail (input, output, retrieval)
+            context: Optional context (e.g., for output rail)
+
+        Returns:
+            Formatted prompt string
+        """
+        pass
+
+    @abstractmethod
+    def _parse_response(
+        self,
+        response: str,
+        content: str,
+        rail_type: RailType,
+    ) -> List[GuardrailResult]:
+        """
+        Parse model response into GuardrailResult objects.
+
+        Args:
+            response: Raw model response text
+            content: Original content that was classified
+            rail_type: Type of rail
+
+        Returns:
+            List of GuardrailResult objects
+        """
+        pass
+
+    def classify(
+        self,
+        content: str,
+        rail_type: RailType,
+        context: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> List[GuardrailResult]:
+        """
+        Classify content using local model.
+
+        Args:
+            content: Content to classify
+            rail_type: Type of rail
+            context: Optional context
+            metadata: Optional metadata
+
+        Returns:
+            List of GuardrailResult objects
+        """
+        start_time = time.time()
+
+        # Handle empty content
+        if not content or not content.strip():
+            return [
+                GuardrailResult(
+                    passed=True,
+                    category="empty",
+                    score=0.0,
+                    model=self.model_name,
+                    reason="Empty or whitespace-only content",
+                    action="pass",
+                    latency_ms=(time.time() - start_time) * 1000,
+                )
+            ]
+
+        try:
+            # Format prompt
+            prompt = self._format_prompt(content, rail_type, context)
+
+            # Generate response
+            if self._use_vllm and self._vllm_client:
+                # Use chat endpoint — lets the server apply its own chat
+                # template (required for ollama, also works with VLLM).
+                # Send raw content so safety models use their built-in training
+                # rather than getting confused by our prompt template.
+                vllm_response = self._vllm_client.chat(
+                    messages=[{"role": "user", "content": content}],
+                    max_tokens=self.MAX_NEW_TOKENS,
+                    temperature=self.TEMPERATURE,
+                )
+                response_text = vllm_response.text
+            else:
+                response_text = self._generate_with_transformers(prompt)
+
+            elapsed_ms = (time.time() - start_time) * 1000
+
+            # Parse response
+            results = self._parse_response(response_text, content, rail_type)
+
+            # Update latency in results
+            for result in results:
+                result.latency_ms = elapsed_ms
+
+            return results
+
+        except Exception as e:
+            elapsed_ms = (time.time() - start_time) * 1000
+            return [
+                GuardrailResult(
+                    passed=False,
+                    category="error",
+                    score=0.0,
+                    model=self.model_name,
+                    reason=f"Local model error: {str(e)}",
+                    action="block",
+                    latency_ms=elapsed_ms,
+                )
+            ]
+
+    async def classify_async(
+        self,
+        content: str,
+        rail_type: RailType,
+        context: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> List[GuardrailResult]:
+        """
+        Async version of classify.
+
+        Uses VLLM async client if available, otherwise falls back to sync.
+
+        Args:
+            content: Content to classify
+            rail_type: Type of rail
+            context: Optional context
+            metadata: Optional metadata
+
+        Returns:
+            List of GuardrailResult objects
+        """
+        start_time = time.time()
+
+        # Handle empty content
+        if not content or not content.strip():
+            return [
+                GuardrailResult(
+                    passed=True,
+                    category="empty",
+                    score=0.0,
+                    model=self.model_name,
+                    reason="Empty or whitespace-only content",
+                    action="pass",
+                    latency_ms=(time.time() - start_time) * 1000,
+                )
+            ]
+
+        try:
+            # Format prompt
+            prompt = self._format_prompt(content, rail_type, context)
+
+            # Generate response
+            if self._use_vllm and self._vllm_client:
+                vllm_response = await self._vllm_client.chat_async(
+                    messages=[{"role": "user", "content": content}],
+                    max_tokens=self.MAX_NEW_TOKENS,
+                    temperature=self.TEMPERATURE,
+                )
+                response_text = vllm_response.text
+            else:
+                # Fall back to sync for transformers
+                import asyncio
+                response_text = await asyncio.get_event_loop().run_in_executor(
+                    None,
+                    lambda: self._generate_with_transformers(prompt)
+                )
+
+            elapsed_ms = (time.time() - start_time) * 1000
+
+            # Parse response
+            results = self._parse_response(response_text, content, rail_type)
+
+            # Update latency in results
+            for result in results:
+                result.latency_ms = elapsed_ms
+
+            return results
+
+        except Exception as e:
+            elapsed_ms = (time.time() - start_time) * 1000
+            return [
+                GuardrailResult(
+                    passed=False,
+                    category="error",
+                    score=0.0,
+                    model=self.model_name,
+                    reason=f"Local model error: {str(e)}",
+                    action="block",
+                    latency_ms=elapsed_ms,
+                )
+            ]
+
+    # Shared keyword-based category inference for local models that
+    # return binary safe/unsafe without fine-grained categories.
+    _CATEGORY_KEYWORDS = {
+        "violence": ["violence", "violent", "kill", "murder", "attack", "weapon", "harm", "bomb", "gun", "hurt"],
+        "self_harm": ["suicide", "self-harm", "self harm", "cut myself", "end my life", "kill myself"],
+        "hate_speech": ["hate", "racist", "sexist", "discrimination", "discriminate", "slur"],
+        "sexual_content": ["sexual", "explicit", "nude", "porn"],
+        "harassment": ["harass", "bully", "threaten", "intimidate"],
+        "illegal_activity": ["illegal", "drug", "hack", "steal", "fraud", "crime", "counterfeit"],
+        "jailbreak": ["jailbreak", "ignore", "bypass", "pretend", "roleplay", "dan"],
+        "prompt_injection": ["ignore previous", "new instructions", "system prompt", "override", "injection"],
+    }
+
+    def _infer_categories(self, content: str, response: str) -> List[str]:
+        """Infer safety categories from content + response via keyword matching."""
+        combined = (content + " " + response).lower()
+        categories = [
+            cat for cat, keywords in self._CATEGORY_KEYWORDS.items()
+            if any(kw in combined for kw in keywords)
+        ]
+        return categories if categories else ["harmful_content"]
+
+    def is_available(self) -> bool:
+        """
+        Check if this backend is available.
+
+        Returns:
+            True if VLLM server is up or transformers can load the model
+        """
+        if self._use_vllm:
+            return True
+
+        # Check if we can load via transformers
+        try:
+            import torch
+            has_gpu = torch.cuda.is_available() or (
+                hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
+            )
+            return has_gpu or self._device == "cpu"
+        except ImportError:
+            return False