mi-crow 0.1.1.post12__py3-none-any.whl

amber/hooks/detector.py

@@ -0,0 +1,95 @@
+from __future__ import annotations
+
+import abc
+from typing import Any, TYPE_CHECKING, Dict
+
+import torch
+
+from amber.hooks.hook import Hook, HookType, HOOK_FUNCTION_INPUT, HOOK_FUNCTION_OUTPUT
+from amber.store.store import Store
+
+if TYPE_CHECKING:
+    pass
+
+
+class Detector(Hook):
+    """
+    Abstract base class for detector hooks that collect metadata during inference.
+    
+    Detectors can accumulate data across batches and optionally save it to a Store.
+    They are designed to observe and record information without modifying activations.
+    """
+
+    def __init__(
+            self,
+            hook_type: HookType | str = HookType.FORWARD,
+            hook_id: str | None = None,
+            store: Store | None = None,
+            layer_signature: str | int | None = None
+    ):
+        """
+        Initialize a detector hook.
+        
+        Args:
+            hook_type: Type of hook (HookType.FORWARD or HookType.PRE_FORWARD)
+            hook_id: Unique identifier
+            store: Optional Store for saving metadata
+            layer_signature: Layer to attach to (optional, for compatibility)
+        """
+        super().__init__(layer_signature=layer_signature, hook_type=hook_type, hook_id=hook_id)
+        self.store = store
+        self.metadata: Dict[str, Any] = {}
+        self.tensor_metadata: Dict[str, torch.Tensor] = {}
+
+    def _hook_fn(
+            self,
+            module: torch.nn.Module,
+            input: HOOK_FUNCTION_INPUT,
+            output: HOOK_FUNCTION_OUTPUT
+    ) -> None:
+        """
+        Internal hook function that collects metadata.
+        
+        This calls process_activations to allow subclasses to implement
+        their specific detection logic.
+        
+        Args:
+            module: The PyTorch module being hooked
+            input: Tuple of input tensors to the module
+            output: Output tensor(s) from the module
+            
+        Raises:
+            Exception: If process_activations raises an exception
+        """
+        if not self._enabled:
+            return None
+        try:
+            self.process_activations(module, input, output)
+        except Exception as e:
+            raise RuntimeError(
+                f"Error in detector {self.id} process_activations: {e}"
+            ) from e
+        return None
+
+    @abc.abstractmethod
+    def process_activations(
+            self,
+            module: torch.nn.Module,
+            input: HOOK_FUNCTION_INPUT,
+            output: HOOK_FUNCTION_OUTPUT
+    ) -> None:
+        """
+        Process activations from the hooked layer.
+        
+        This is where detector-specific logic goes (e.g., tracking top activations,
+        computing statistics, etc.).
+        
+        Args:
+            module: The PyTorch module being hooked
+            input: Tuple of input tensors to the module
+            output: Output tensor(s) from the module
+            
+        Raises:
+            Exception: Subclasses may raise exceptions for invalid inputs or processing errors
+        """
+        raise NotImplementedError("process_activations must be implemented by subclasses")

amber/hooks/hook.py

@@ -0,0 +1,218 @@
+from __future__ import annotations
+
+import abc
+import uuid
+from enum import Enum
+from typing import Callable, TypeAlias, Sequence, Optional, TYPE_CHECKING
+
+import torch
+from torch import nn, Tensor
+from torch.types import _TensorOrTensors
+
+if TYPE_CHECKING:
+    from amber.language_model.context import LanguageModelContext
+
+
+class HookType(str, Enum):
+    """Type of hook to register on a layer."""
+    FORWARD = "forward"
+    PRE_FORWARD = "pre_forward"
+
+
+HOOK_FUNCTION_INPUT: TypeAlias = Sequence[Tensor]
+HOOK_FUNCTION_OUTPUT: TypeAlias = _TensorOrTensors | None
+
+
+class HookError(Exception):
+    """Exception raised when a hook encounters an error during execution."""
+    
+    def __init__(self, hook_id: str, hook_type: str, original_error: Exception):
+        """
+        Initialize HookError.
+        
+        Args:
+            hook_id: Unique identifier of the hook that raised the error
+            hook_type: Type of hook (e.g., "forward", "pre_forward")
+            original_error: The original exception that was raised
+        """
+        self.hook_id = hook_id
+        self.hook_type = hook_type
+        self.original_error = original_error
+        message = f"Hook {hook_id} (type={hook_type}) raised exception: {original_error}"
+        super().__init__(message)
+
+
+class Hook(abc.ABC):
+    """
+    Abstract base class for hooks that can be registered on language model layers.
+    
+    Hooks provide a way to intercept and process activations during model inference.
+    They expose PyTorch-compatible callables via get_torch_hook() while providing
+    additional functionality like enable/disable and unique identification.
+    """
+
+    def __init__(
+            self,
+            layer_signature: str | int | None = None,
+            hook_type: HookType | str = HookType.FORWARD,
+            hook_id: str | None = None
+    ):
+        """
+        Initialize a hook.
+        
+        Args:
+            layer_signature: Layer name or index to attach hook to
+            hook_type: Type of hook - HookType.FORWARD or HookType.PRE_FORWARD
+            hook_id: Unique identifier (auto-generated if not provided)
+            
+        Raises:
+            ValueError: If hook_type string is invalid
+        """
+        self.layer_signature = layer_signature
+        self.hook_type = self._normalize_hook_type(hook_type)
+        self.id = hook_id if hook_id is not None else str(uuid.uuid4())
+        self._enabled = True
+        self._torch_hook_handle = None
+        self._context: Optional["LanguageModelContext"] = None
+
+    def _normalize_hook_type(self, hook_type: HookType | str) -> HookType:
+        """Normalize hook_type to HookType enum.
+        
+        Args:
+            hook_type: HookType enum or string value
+            
+        Returns:
+            HookType enum value
+            
+        Raises:
+            ValueError: If hook_type string is not a valid HookType value
+        """
+        if isinstance(hook_type, HookType):
+            return hook_type
+        
+        if isinstance(hook_type, str):
+            try:
+                return HookType(hook_type)
+            except ValueError:
+                valid_values = [ht.value for ht in HookType]
+                raise ValueError(
+                    f"Invalid hook_type string '{hook_type}'. "
+                    f"Must be one of: {valid_values}"
+                ) from None
+        
+        raise ValueError(
+            f"hook_type must be HookType enum or string, got: {type(hook_type)}"
+        )
+
+    def _create_pre_forward_wrapper(self) -> Callable:
+        """Create a pre-forward hook wrapper function.
+        
+        Returns:
+            Wrapper function for pre-forward hooks
+        """
+        def pre_forward_wrapper(module: nn.Module, input: HOOK_FUNCTION_INPUT) -> None | HOOK_FUNCTION_INPUT:
+            if not self._enabled:
+                return None
+            try:
+                result = self._hook_fn(module, input, None)
+                return result if result is not None else None
+            except Exception as e:
+                raise HookError(self.id, self.hook_type.value, e) from e
+
+        return pre_forward_wrapper
+
+    def _create_forward_wrapper(self) -> Callable:
+        """Create a forward hook wrapper function.
+        
+        Returns:
+            Wrapper function for forward hooks
+        """
+        def forward_wrapper(module: nn.Module, input: HOOK_FUNCTION_INPUT, output: HOOK_FUNCTION_OUTPUT) -> None:
+            if not self._enabled:
+                return None
+            try:
+                self._hook_fn(module, input, output)
+                return None
+            except Exception as e:
+                raise HookError(self.id, self.hook_type.value, e) from e
+
+        return forward_wrapper
+
+    @property
+    def enabled(self) -> bool:
+        """Whether this hook is currently enabled."""
+        return self._enabled
+
+    def enable(self) -> None:
+        """Enable this hook."""
+        self._enabled = True
+
+    def disable(self) -> None:
+        """Disable this hook."""
+        self._enabled = False
+
+    @property
+    def context(self) -> Optional["LanguageModelContext"]:
+        """Get the LanguageModelContext associated with this hook."""
+        return self._context
+
+    def set_context(self, context: "LanguageModelContext") -> None:
+        """Set the LanguageModelContext for this hook.
+        
+        Args:
+            context: The LanguageModelContext instance
+        """
+        self._context = context
+
+    def _is_both_controller_and_detector(self) -> bool:
+        """
+        Check if this hook instance inherits from both Controller and Detector.
+        
+        Uses MRO (Method Resolution Order) to check for both class names
+        without requiring imports, avoiding circular dependencies.
+        
+        Returns:
+            True if the instance inherits from both Controller and Detector, False otherwise
+        """
+        mro_class_names = [cls.__name__ for cls in type(self).__mro__]
+        return 'Controller' in mro_class_names and 'Detector' in mro_class_names
+
+    def get_torch_hook(self) -> Callable:
+        """
+        Return a PyTorch-compatible hook function.
+        
+        The returned callable will check the enabled flag before executing
+        and call the abstract _hook_fn method.
+        
+        Returns:
+            A callable compatible with PyTorch's register_forward_hook or
+            register_forward_pre_hook APIs.
+        """
+        if self.hook_type == HookType.PRE_FORWARD:
+            return self._create_pre_forward_wrapper()
+        else:
+            return self._create_forward_wrapper()
+
+    @abc.abstractmethod
+    def _hook_fn(
+            self,
+            module: torch.nn.Module,
+            input: HOOK_FUNCTION_INPUT,
+            output: HOOK_FUNCTION_OUTPUT
+    ) -> None | HOOK_FUNCTION_INPUT:
+        """
+        Internal hook function to be implemented by subclasses.
+        
+        Args:
+            module: The PyTorch module being hooked
+            input: Tuple of input tensors to the module
+            output: Output tensor(s) from the module (None for pre_forward hooks)
+            
+        Returns:
+            For pre_forward hooks: modified inputs (tuple) or None to keep original
+            For forward hooks: None (forward hooks cannot modify output in PyTorch)
+            
+        Raises:
+            Exception: Subclasses may raise exceptions which will be caught by the wrapper
+        """
+        raise NotImplementedError("_hook_fn must be implemented by subclasses")

amber/hooks/implementations/function_controller.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+from typing import Callable, TYPE_CHECKING
+import torch
+
+from amber.hooks.controller import Controller
+from amber.hooks.hook import HookType
+
+if TYPE_CHECKING:
+    from torch import nn
+
+
+class FunctionController(Controller):
+    """
+    A controller that applies a user-provided function to tensors during inference.
+    
+    This controller allows users to pass any function and apply it to activations.
+    The function will be applied to:
+    - Single tensors directly
+    - All tensors in tuples/lists (default behavior)
+    
+    Example:
+        >>> # Scale activations by 2
+        >>> controller = FunctionController(
+        ...     layer_signature="layer_0",
+        ...     function=lambda x: x * 2.0
+        ... )
+    """
+    
+    def __init__(
+        self,
+        layer_signature: str | int,
+        function: Callable[[torch.Tensor], torch.Tensor],
+        hook_type: HookType | str = HookType.FORWARD,
+        hook_id: str | None = None,
+    ):
+        """
+        Initialize a function controller.
+        
+        Args:
+            layer_signature: Layer to attach to
+            function: Function to apply to tensors. Must take a torch.Tensor and return a torch.Tensor
+            hook_type: Type of hook (HookType.FORWARD or HookType.PRE_FORWARD)
+            hook_id: Unique identifier
+            
+        Raises:
+            ValueError: If function is None or not callable
+        """
+        if function is None:
+            raise ValueError("function cannot be None")
+        
+        if not callable(function):
+            raise ValueError(f"function must be callable, got: {type(function)}")
+        
+        super().__init__(hook_type=hook_type, hook_id=hook_id, layer_signature=layer_signature)
+        self.function = function
+    
+    def modify_activations(
+        self,
+        module: "nn.Module",
+        inputs: torch.Tensor | None,
+        output: torch.Tensor | None
+    ) -> torch.Tensor | None:
+        """
+        Apply the user-provided function to activations.
+        
+        Args:
+            module: The PyTorch module being hooked
+            inputs: Input tensor (None for forward hooks)
+            output: Output tensor (None for pre_forward hooks)
+            
+        Returns:
+            Modified tensor with function applied, or None if target tensor is None
+            
+        Raises:
+            RuntimeError: If function raises an exception when applied to tensor
+        """
+        target = output if self.hook_type == HookType.FORWARD else inputs
+        
+        if target is None or not isinstance(target, torch.Tensor):
+            return target
+        
+        try:
+            result = self.function(target)
+            if not isinstance(result, torch.Tensor):
+                raise TypeError(
+                    f"Function must return a torch.Tensor, got: {type(result)}"
+                )
+            return result
+        except Exception as e:
+            raise RuntimeError(
+                f"Error applying function in FunctionController {self.id}: {e}"
+            ) from e

amber/hooks/implementations/layer_activation_detector.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+import torch
+
+from amber.hooks.detector import Detector
+from amber.hooks.hook import HookType, HOOK_FUNCTION_INPUT, HOOK_FUNCTION_OUTPUT
+from amber.hooks.utils import extract_tensor_from_output
+
+if TYPE_CHECKING:
+    from torch import nn
+
+
+class LayerActivationDetector(Detector):
+    """
+    Detector hook that captures and saves activations during inference.
+    
+    This detector extracts activations from layer outputs and stores them
+    for later use (e.g., saving to disk, further analysis).
+    """
+
+    def __init__(
+            self,
+            layer_signature: str | int,
+            hook_id: str | None = None
+    ):
+        """
+        Initialize the activation saver detector.
+        
+        Args:
+            layer_signature: Layer to capture activations from
+            hook_id: Unique identifier for this hook
+            
+        Raises:
+            ValueError: If layer_signature is None
+        """
+        if layer_signature is None:
+            raise ValueError("layer_signature cannot be None for LayerActivationDetector")
+
+        super().__init__(
+            hook_type=HookType.FORWARD,
+            hook_id=hook_id,
+            store=None,
+            layer_signature=layer_signature
+        )
+
+    def process_activations(
+            self,
+            module: torch.nn.Module,
+            input: HOOK_FUNCTION_INPUT,
+            output: HOOK_FUNCTION_OUTPUT
+    ) -> None:
+        """
+        Extract and store activations from output.
+        
+        Handles various output types:
+        - Plain tensors
+        - Tuples/lists of tensors (takes first tensor)
+        - Objects with last_hidden_state attribute (e.g., HuggingFace outputs)
+        
+        Args:
+            module: The PyTorch module being hooked
+            input: Tuple of input tensors to the module
+            output: Output tensor(s) from the module
+            
+        Raises:
+            RuntimeError: If tensor extraction or storage fails
+        """
+        try:
+            tensor = extract_tensor_from_output(output)
+
+            if tensor is not None:
+                tensor_cpu = tensor.detach().to("cpu")
+                # Store current batch's tensor (overwrites previous)
+                self.tensor_metadata['activations'] = tensor_cpu
+                # Store activations shape to metadata
+                self.metadata['activations_shape'] = tuple(tensor_cpu.shape)
+        except Exception as e:
+            raise RuntimeError(
+                f"Error extracting activations in LayerActivationDetector {self.id}: {e}"
+            ) from e
+
+    def get_captured(self) -> torch.Tensor | None:
+        """
+        Get the captured activations from the current batch.
+        
+        Returns:
+            The captured activation tensor from the current batch or None if no activations captured yet
+        """
+        return self.tensor_metadata.get('activations')
+
+    def clear_captured(self) -> None:
+        """Clear captured activations for current batch."""
+        self.tensor_metadata.pop('activations', None)
+        self.metadata.pop('activations_shape', None)
+