diffract-core 0.2.1__py3-none-any.whl

diffract/core/data/nn/extractors/base.py

@@ -0,0 +1,300 @@
+"""Base classes and utilities for parameter extraction.
+
+This module provides the foundational infrastructure for parameter extraction
+from neural network models. It implements a handler-based architecture that
+allows for extensible and framework-specific parameter processing.
+
+Key Components:
+    - ExtractorOverrides: Configuration for customizing parameter metadata
+    - ParameterOverrides: Per-parameter override specifications
+    - BaseParameterExtractor: Abstract base class for all extractors
+
+The handler-based approach allows different parameter types to be processed
+by specialized handlers while maintaining a consistent extraction interface.
+This design supports extensibility for new parameter types and frameworks.
+
+Example:
+    >>> overrides = ExtractorOverrides(
+    ...     model_id="custom_model",
+    ...     parameter_overrides={"layer1.weight": ParameterOverrides(ptype="CUSTOM")},
+    ... )
+    >>> extractor = ConcreteExtractor(model, overrides=overrides)
+    >>> parameters = extractor.extract_parameters(storage, cache)
+"""
+
+from __future__ import annotations
+
+import logging
+from abc import ABC, abstractmethod
+from contextlib import suppress
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+
+from diffract.core.constants import WEIGHTS_FIELD
+from diffract.core.data.nn.params.interface import IParameterRepository
+from diffract.core.data.nn.params.metadata import ParameterMetadata
+from diffract.core.data.nn.params.proxy import ParameterDataProxy
+from diffract.core.data.nn.params.schema import ParameterType
+from diffract.core.utils.hashing import get_unique_id
+
+from .interface import IParameterExtractor
+
+if TYPE_CHECKING:
+    from collections.abc import Generator
+
+    from .handlers.base import ParameterHandler, ParameterHandlerRegistry
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ExtractorOverrides:
+    """Configuration for customizing parameter extraction behavior.
+
+    Provides a mechanism to override default parameter metadata and
+    model identification during the extraction process. This allows
+    for custom naming, typing, and categorization of parameters.
+
+    Attributes:
+        model_id: Custom model identifier to override default.
+        parameter_overrides: Per-parameter override specifications.
+    """
+
+    @dataclass
+    class ParameterOverrides:
+        """Override specifications for individual parameters.
+
+        Allows customization of parameter metadata on a per-parameter
+        basis, including name changes, type reclassification, and extra metadata.
+
+        Attributes:
+            name: Custom parameter name to override default.
+            ptype: Custom parameter type to override default.
+            other_meta: Additional metadata to merge into ParameterMetadata.other_meta.
+        """
+
+        name: str | None = field(default=None)
+        ptype: str | ParameterType | None = field(default=None)
+        other_meta: dict[str, Any] | None = field(default=None)
+
+    model_id: str | None = field(default=None)
+    parameter_overrides: dict[str, ParameterOverrides] = field(default_factory=dict)
+
+
+ParameterOverrides = ExtractorOverrides.ParameterOverrides
+
+
+class BaseParameterExtractor(IParameterExtractor, ABC):
+    """Abstract base class for parameter extractors with handler-based processing.
+
+    Provides a common framework for extracting parameters from neural network
+    models using a pluggable handler system. Handles parameter iteration,
+    metadata construction, override application, and error management.
+
+    The extractor uses a registry of parameter handlers to process different
+    types of parameters. Each handler specializes in extracting weights and
+    metadata from specific parameter types (e.g., dense layers, convolutions).
+
+    Attributes:
+        model: The neural network model to extract parameters from.
+        overrides: Configuration for customizing extraction behavior.
+        skip_not_implemented_types: Whether to skip unsupported parameter types.
+        handler_registry: Registry of parameter handlers for processing.
+    """
+
+    def __init__(
+        self,
+        model: Any,
+        overrides: ExtractorOverrides | None = None,
+        skip_not_implemented_types: bool = True,
+        custom_handlers: list[ParameterHandler] | None = None,
+    ) -> None:
+        """Initialize the base parameter extractor.
+
+        Args:
+            model: Neural network model to extract parameters from.
+            overrides: Optional extraction behavior overrides.
+            skip_not_implemented_types: Whether to skip unsupported parameters.
+            custom_handlers: Optional list of custom parameter handlers.
+        """
+        self.model = model
+        self.overrides = overrides or ExtractorOverrides()
+        self.skip_not_implemented_types = skip_not_implemented_types
+
+        # Import here to avoid circular imports
+        from .handlers.base import ParameterHandlerRegistry
+
+        self.handler_registry: ParameterHandlerRegistry[Any] = ParameterHandlerRegistry[
+            Any
+        ]()
+        self._register_default_handlers()
+
+        if custom_handlers:
+            for handler in custom_handlers:
+                self.handler_registry.register_handler(handler)
+
+    @abstractmethod
+    def _register_default_handlers(self) -> None:
+        """Register default parameter handlers for this extractor.
+
+        Subclasses must implement this method to register the appropriate
+        handlers for their specific framework and parameter types.
+        """
+
+    def extract_parameters(
+        self,
+        parameter_repository: IParameterRepository,
+    ) -> None:
+        """Extract all parameters from the model using registered handlers.
+
+        Iterates through model parameters, processes each with appropriate
+        handlers, constructs metadata with overrides, and creates parameter
+        proxies for storage and caching.
+
+        Args:
+            parameter_repository: Repository that owns storage/cache and receives
+                extracted parameters.
+
+        Raises:
+            NotImplementedError: If parameter type unsupported and not skipping.
+            Exception: If handler processing fails and not skipping.
+        """
+        with parameter_repository:
+            model_id = get_unique_id()
+
+            for param_idx, (param_name, parameter) in enumerate(
+                self._iter_parameters()
+            ):
+                try:
+                    result = self.handler_registry.process_parameter(
+                        parameter, param_name
+                    )
+                except Exception:
+                    # Preserve error context; decide policy: skip or raise
+                    logger.exception('Handler failed for "%s"', param_name)
+                    if self.skip_not_implemented_types:
+                        continue
+                    raise
+
+                if result is None:
+                    if self.skip_not_implemented_types:
+                        logger.debug('Skipping unsupported parameter "%s"', param_name)
+                        continue
+                    msg = (
+                        f'No handler found for parameter "{param_name}" of type '
+                        f"{type(parameter)}"
+                    )
+                    raise NotImplementedError(msg)
+
+                ptype, weights, handler = result
+
+                param_meta = self._build_metadata(
+                    parameter, model_id, param_name, param_idx, ptype, handler
+                )
+                param_proxy = ParameterDataProxy.create_and_store(
+                    meta=param_meta,
+                    repository=parameter_repository,
+                )
+                param_proxy.set_field(WEIGHTS_FIELD, weights)
+
+                logger.debug(
+                    'Extracted parameter "%s" with type %s (handler=%s)',
+                    param_meta.name,
+                    param_meta.ptype,
+                    type(handler).__name__,
+                )
+
+                parameter_repository.append(param_proxy)
+
+    def _build_metadata(
+        self,
+        parameter: Any,
+        model_id: str,
+        param_name: str,
+        param_idx: int,
+        ptype: ParameterType,
+        handler: ParameterHandler,
+    ) -> ParameterMetadata:
+        """Build parameter metadata with handler input and overrides.
+
+        Constructs base metadata, enriches it with handler-provided additional
+        metadata, and applies any configured overrides.
+
+        Args:
+            parameter: The original parameter object.
+            model_id: Identifier of the model the parameter belongs to.
+            param_name: Name of the parameter.
+            param_idx: Index of parameter in model iteration order.
+            ptype: Parameter type determined by handler.
+            handler: Handler that processed this parameter.
+
+        Returns:
+            Complete parameter metadata with overrides applied.
+        """
+        metadata: dict[str, Any] = {
+            "name": param_name,
+            "ptype": ptype,
+            "model_id": model_id,
+            "other_meta": {"in_model_idx": param_idx},
+        }
+
+        # Best-effort metadata from the selected handler (no exceptions)
+        with suppress(Exception):
+            extra = handler.get_additional_metadata(parameter, param_name) or {}
+            metadata["other_meta"].update(extra)
+
+        self._apply_overrides(metadata)
+        return ParameterMetadata(**metadata)
+
+    def _apply_overrides(self, metadata: dict[str, Any]) -> None:
+        """Apply configured overrides to parameter metadata.
+
+        Modifies the metadata dictionary in-place based on configured
+        parameter-specific and global overrides. Preserves original
+        values in other_meta for traceability.
+
+        Args:
+            metadata: Parameter metadata dictionary to modify.
+        """
+        param_name = metadata["name"]
+
+        if param_name in self.overrides.parameter_overrides:
+            param_overrides = self.overrides.parameter_overrides[param_name]
+
+            if param_overrides.name:
+                metadata["name"] = param_overrides.name
+                metadata.setdefault("other_meta", {})["original_name"] = param_name
+
+            if param_overrides.ptype is not None:
+                metadata.setdefault("other_meta", {})["original_ptype"] = metadata[
+                    "ptype"
+                ].name
+                metadata["ptype"] = (
+                    ParameterType.from_string(param_overrides.ptype)
+                    if isinstance(param_overrides.ptype, str)
+                    else param_overrides.ptype
+                )
+
+            if param_overrides.other_meta:
+                metadata.setdefault("other_meta", {}).update(param_overrides.other_meta)
+
+        if self.overrides.model_id:
+            metadata.setdefault("other_meta", {})["original_model_id"] = metadata[
+                "model_id"
+            ]
+            metadata["model_id"] = self.overrides.model_id
+
+    @abstractmethod
+    def _iter_parameters(self) -> Generator[tuple[str, Any], None, None]:
+        """Iterate over model parameters yielding (name, parameter) tuples.
+
+        Subclasses must implement this method to provide framework-specific
+        parameter iteration logic.
+
+        Yields:
+            Tuples of (parameter_name, parameter_object) for each parameter.
+
+        Raises:
+            NotImplementedError: Must be implemented by subclasses.
+        """
+        raise NotImplementedError

diffract/core/data/nn/extractors/factory.py

@@ -0,0 +1,234 @@
+"""Factory pattern for creating framework-specific parameter extractors.
+
+This module implements the factory pattern for automatically detecting neural
+network model types and creating appropriate parameter extractors. It provides
+a unified interface that abstracts away framework-specific details while
+enabling extensible handler-based parameter processing.
+
+The factory automatically detects supported frameworks and model types,
+creating the most appropriate extractor implementation. It supports custom
+handlers for specialized parameter processing and provides comprehensive
+error handling for unsupported model types.
+
+Supported Frameworks:
+    - NumPy: dict[str, numpy.ndarray] mappings (no framework required)
+    - PyTorch: torch.nn.Module and state_dict objects
+    - Extensible: Easy to add support for additional frameworks
+
+Key Features:
+    - Automatic model type detection
+    - Framework-specific extractor creation
+    - Custom handler support for specialized processing
+    - Comprehensive error reporting with supported types
+    - Runtime framework availability checking
+
+Example:
+    >>> import torch
+    >>> model = torch.nn.Linear(10, 5)
+    >>> extractor = create_extractor(model)
+    >>> # With custom handlers:
+    >>> extractor = create_extractor(model, custom_handlers=[MyHandler()])
+    >>> parameters = extractor.extract_parameters(storage, cache)
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+import numpy as np
+
+import diffract.core.utils.imports as import_utils
+
+from .interface import IParameterExtractor
+
+torch = import_utils.get_module("torch")
+onnx = import_utils.get_module("onnx")
+tf = import_utils.get_module("tensorflow")
+flax = import_utils.get_module("flax")
+
+logger = logging.getLogger(__name__)
+
+
+def create_extractor(model: Any, *args: Any, **kwargs: Any) -> IParameterExtractor:
+    """Create appropriate parameter extractor based on model type.
+
+    Automatically detects the model type and framework, then creates the
+    most suitable parameter extractor. Supports custom handlers and
+    extractor overrides for specialized parameter processing.
+
+    Args:
+        model: Neural network model to extract parameters from.
+        *args: Additional positional arguments passed to extractor.
+        **kwargs: Additional keyword arguments passed to extractor.
+
+    Returns:
+        Parameter extractor instance appropriate for the model type.
+
+    Raises:
+        ImportError: If the model is not a NumPy dict and no supported
+            frameworks are available.
+        TypeError: If model type is not supported by any available framework.
+
+    Example:
+        >>> extractor = create_extractor(torch_model)
+        >>> extractor = create_extractor({"encoder.weight": np.random.rand(10, 5)})
+        >>> extractor = create_extractor(state_dict, custom_handlers=[handler])
+    """
+    if isinstance(model, dict) and all(
+        isinstance(name, str) and isinstance(value, np.ndarray)
+        for name, value in model.items()
+    ):
+        from .numpy import NumpyDictExtractor
+
+        logger.debug(
+            "Creating NumpyDictExtractor for array dict with %d parameters",
+            len(model),
+        )
+
+        return NumpyDictExtractor(*args, model=model, **kwargs)
+
+    if not get_supported_frameworks():
+        msg = (
+            "No supported deep learning frameworks available. "
+            "Please install PyTorch (torch) or other supported frameworks, "
+            "or pass a dict[str, numpy.ndarray] of weight matrices."
+        )
+        raise ImportError(msg)
+
+    if torch:
+        # PyTorch state_dict (dictionary of tensors)
+        if (
+            isinstance(model, dict)
+            and model
+            and all(isinstance(v, torch.Tensor) for v in model.values())
+        ):
+            from .torch import TorchStateDictExtractor
+
+            logger.debug(
+                "Creating TorchStateDictExtractor for state dict with %d parameters",
+                len(model),
+            )
+
+            return TorchStateDictExtractor(*args, model=model, **kwargs)
+
+        # PyTorch nn.Module (standard PyTorch models)
+        if isinstance(model, torch.nn.Module):
+            from .torch import TorchModuleExtractor
+
+            logger.debug(
+                "Creating TorchModuleExtractor for %s model", type(model).__name__
+            )
+
+            return TorchModuleExtractor(*args, model=model, **kwargs)
+
+    if onnx and isinstance(model, onnx.ModelProto):
+        # ONNX model
+        from .onnx import OnnxModelExtractor
+
+        logger.debug("Creating OnnxModelExtractor for ONNX ModelProto")
+        return OnnxModelExtractor(*args, model=model, **kwargs)
+
+    if tf and isinstance(model, (tf.keras.Model, tf.keras.layers.Layer)):
+        # TensorFlow / Keras model or layer
+        from .tensorflow import TensorFlowModelExtractor
+
+        logger.debug("Creating TensorFlowModelExtractor for %s", type(model).__name__)
+        return TensorFlowModelExtractor(*args, model=model, **kwargs)
+
+    if flax:
+        # Flax variables/params: mapping with "params" key or params mapping itself.
+        # We avoid importing JAX types here; rely on duck-typing for mappings.
+        try:
+            from collections.abc import Mapping
+        except ImportError:
+            mapping_type = None  # type: ignore[assignment]
+        else:
+            mapping_type = Mapping
+
+        if mapping_type is not None and isinstance(model, mapping_type):
+            from .flax import FlaxParamsExtractor
+
+            logger.debug("Creating FlaxParamsExtractor for params/variables mapping")
+            return FlaxParamsExtractor(*args, model=model, **kwargs)
+
+        if hasattr(model, "variables"):
+            from .flax import FlaxParamsExtractor
+
+            logger.debug("Creating FlaxParamsExtractor for bound Flax module")
+            return FlaxParamsExtractor(*args, model=model, **kwargs)
+
+    error_msg = f"Unsupported model type: {type(model).__name__}. "
+
+    if supported_types := get_supported_types():
+        error_msg += f"Supported types: {', '.join(supported_types)}"
+    else:
+        error_msg += "No frameworks available."
+
+    raise TypeError(error_msg)
+
+
+def get_supported_frameworks() -> list[str]:
+    """Get list of available deep learning frameworks.
+
+    Checks for the availability of supported deep learning frameworks
+    and returns their names. This is useful for diagnostics and
+    conditional feature availability.
+
+    Returns:
+        List of names of available frameworks.
+
+    Example:
+        >>> frameworks = get_supported_frameworks()
+        >>> print(f"Available frameworks: {', '.join(frameworks)}")
+    """
+    frameworks: list[str] = []
+    if torch:
+        frameworks.append("PyTorch")
+    if tf:
+        frameworks.append("TensorFlow")
+    if flax:
+        frameworks.append("Flax")
+    if onnx:
+        frameworks.append("ONNX")
+    return frameworks
+
+
+def get_supported_types() -> list[str]:
+    """Get list of supported model types across all frameworks.
+
+    Returns a list of human-readable descriptions of supported model
+    types. This is useful for error messages and documentation.
+
+    Returns:
+        List of supported model type descriptions.
+
+    Example:
+        >>> types = get_supported_types()
+        >>> print(f"Supported types: {', '.join(types)}")
+    """
+    types: list[str] = ["dict[str, numpy.ndarray] (NumPy weight matrices)"]
+    if torch:
+        types.extend(
+            [
+                "torch.nn.Module (PyTorch models)",
+                "dict[str, torch.Tensor] (PyTorch state_dict)",
+            ]
+        )
+    if tf:
+        types.extend(
+            [
+                "tf.keras.Model (TensorFlow/Keras models)",
+                "tf.keras.layers.Layer (TensorFlow/Keras layers)",
+            ]
+        )
+    if flax:
+        types.extend(
+            [
+                'Mapping[str, Any] (Flax params tree / variables with key "params")',
+                "Bound module with `.variables['params']` (Flax Linen bound module)",
+            ]
+        )
+    if onnx:
+        types.append("onnx.ModelProto (ONNX model)")
+    return types

diffract/core/data/nn/extractors/flax.py

@@ -0,0 +1,112 @@
+"""Flax/JAX parameter extractors.
+
+Provides a parameter extractor for Flax models by reading a variables/params tree.
+Currently supports only dense (fully-connected) kernel weights.
+
+Supported inputs:
+  - A mapping with key `"params"` (e.g. flax variables FrozenDict)
+  - A bound Flax Linen module with `.variables` containing `"params"`
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Mapping
+from typing import TYPE_CHECKING, Any, Self
+
+import diffract.core.utils.imports as import_utils
+
+from .base import BaseParameterExtractor
+
+if TYPE_CHECKING:
+    from collections.abc import Generator
+
+    from .base import ExtractorOverrides
+    from .handlers.base import ParameterHandler
+
+logger = logging.getLogger(__name__)
+
+
+def _flatten_mapping(
+    mapping: Mapping[str, Any],
+    *,
+    prefix: str = "",
+) -> Generator[tuple[str, Any], None, None]:
+    """Flatten nested mappings into dot-separated key paths."""
+    for k, v in mapping.items():
+        name = f"{prefix}.{k}" if prefix else str(k)
+        if isinstance(v, Mapping):
+            yield from _flatten_mapping(v, prefix=name)
+        else:
+            yield name, v
+
+
+if not (import_utils.is_available("flax") and import_utils.is_available("jax")):
+    logger.debug("Flax/JAX not available, disabling Flax extractors")
+
+    class FlaxParamsExtractor(BaseParameterExtractor):
+        """Stub implementation when Flax/JAX is not available."""
+
+        def __new__(cls, *_args: Any, **_kwargs: Any) -> Self:  # noqa: D102
+            msg = "Flax/JAX packages not available"
+            raise ImportError(msg)
+
+else:
+    from .handlers import FlaxDenseHandler
+
+    class FlaxParamsExtractor(BaseParameterExtractor):
+        """Flax params/variables extractor (DENSE only)."""
+
+        def __init__(
+            self,
+            model: Any,
+            overrides: ExtractorOverrides | None = None,
+            skip_not_implemented_types: bool = True,
+            custom_handlers: list[ParameterHandler] | None = None,
+        ) -> None:
+            super().__init__(
+                model, overrides, skip_not_implemented_types, custom_handlers
+            )
+
+        def _register_default_handlers(self) -> None:
+            self.handler_registry.register_handler_class(FlaxDenseHandler)
+            logger.debug(
+                "Registered %d default handlers for Flax params (DENSE only)",
+                len(self.handler_registry.list_handlers()),
+            )
+
+        def _iter_parameters(self) -> Generator[tuple[str, Any], None, None]:
+            """Iterate params leaves.
+
+            We extract only params named `*.kernel` (Dense convention). Other leaves
+            (bias, embeddings, etc.) are ignored by handlers.
+            """
+            obj = self.model
+
+            # Case 1: bound module with variables dict.
+            if hasattr(obj, "variables"):
+                variables = obj.variables
+                if isinstance(variables, Mapping) and "params" in variables:
+                    params = variables["params"]
+                else:
+                    params = None
+            else:
+                params = None
+
+            # Case 2: variables mapping or params mapping passed directly.
+            if params is None:
+                if isinstance(obj, Mapping) and "params" in obj:
+                    params = obj["params"]
+                elif isinstance(obj, Mapping):
+                    # Treat passed object as the params tree.
+                    params = obj
+
+            if not isinstance(params, Mapping):
+                msg = (
+                    "FlaxParamsExtractor expects a mapping of params (or variables "
+                    'mapping containing key "params"), got '
+                    f"{type(self.model)}"
+                )
+                raise TypeError(msg)
+
+            yield from _flatten_mapping(params)

diffract/core/data/nn/extractors/handlers/__init__.py

@@ -0,0 +1,39 @@
+"""Parameter handlers for extensible parameter processing.
+
+This subpackage implements a handler-based architecture for processing neural
+network parameters across frameworks. Handlers detect whether they can process
+an input parameter and, if so, convert it to standardized NumPy weights and
+provide optional metadata.
+
+Components:
+    - base: Abstract handler interface and lightweight registry
+    - numpy_handlers: NumPy handlers for dense parameters
+    - torch_handlers: PyTorch handlers for dense parameters
+
+Example:
+    >>> from diffract.core.data.nn.extractors.handlers import (
+    ...     ParameterHandlerRegistry,
+    ...     TorchDenseHandler,
+    ... )
+    >>> registry = ParameterHandlerRegistry()
+    >>> registry.register_handler_class(TorchDenseHandler)
+    >>> # Later: handler = registry.process_parameter(module, name)
+"""
+
+from .base import ParameterHandler, ParameterHandlerRegistry
+from .flax_handlers import FlaxDenseHandler
+from .numpy_handlers import NumpyDenseHandler
+from .onnx_handlers import OnnxDenseHandler
+from .tensorflow_handlers import TensorFlowDenseHandler
+from .torch_handlers import TorchDenseHandler, TorchStateDictDenseHandler
+
+__all__ = [
+    "FlaxDenseHandler",
+    "NumpyDenseHandler",
+    "OnnxDenseHandler",
+    "ParameterHandler",
+    "ParameterHandlerRegistry",
+    "TensorFlowDenseHandler",
+    "TorchDenseHandler",
+    "TorchStateDictDenseHandler",
+]