Skip to content

Compose

Hybrid model building from extracted components.

PyTorchModelBuilder

PyTorchModelBuilder

Bases: BaseModelBuilder

Builder for PyTorch hybrid models.

Combines extracted components from different models into new architectures with dimension adapters.

Source code in src/model_garage/compose/builder.py 
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
class PyTorchModelBuilder(BaseModelBuilder):
    """
    Builder for PyTorch hybrid models.

    Combines extracted components from different models
    into new architectures with dimension adapters.
    """

    def __init__(self, name: str):
        super().__init__(name, framework="pytorch")
        self.execution_graph = {}

    def add_component(self, name: str, component: Any, metadata: Dict[str, Any]) -> None:
        """
        Add a component to the model.

        Args:
            name: Name for this component in the architecture
            component: The component (PyTorch module or ExtractedComponent)
            metadata: Component metadata with input/output dimensions
        """
        if name in self.components:
            logger.warning(f"Component {name} already exists, overwriting")
        self.components[name] = (component, metadata)
        logger.info(
            f"Added {name}: {metadata.get('input_dim')}→{metadata.get('output_dim')}"
        )

    def add_adapter(self, source: str, target: str, adapter: Any) -> None:
        """
        Add a dimension adapter between components.

        Args:
            source: Name of source component
            target: Name of target component
            adapter: Linear adapter module
        """
        if source not in self.components:
            raise ValueError(f"Source component {source} does not exist")
        if target not in self.components:
            raise ValueError(f"Target component {target} does not exist")

        _, source_meta = self.components[source]
        _, target_meta = self.components[target]

        adapter_meta = {
            "input_dim": source_meta.get("output_dim"),
            "output_dim": target_meta.get("input_dim"),
        }

        self.adapters[(source, target)] = (adapter, adapter_meta)

        if source not in self.execution_graph:
            self.execution_graph[source] = []
        self.execution_graph[source].append(target)

    def define_forward_pass(self, execution_order: List[str]) -> None:
        """Define the execution order for the forward pass."""
        for name in execution_order:
            if name not in self.components:
                raise ValueError(f"Component {name} not found")
        self.architecture["execution_order"] = execution_order

    def build(self) -> Dict[str, Any]:
        """
        Build the model from components.

        Returns:
            Dict containing the assembled model specification.
        """
        is_valid, errors = self.validate_architecture()
        if not is_valid:
            raise ValueError(f"Invalid architecture: {'; '.join(errors)}")

        logger.info(f"Building {self.name}")

        return {
            "name": self.name,
            "components": {n: c[0] for n, c in self.components.items()},
            "adapters": {
                f"{s}->{t}": a[0] for (s, t), a in self.adapters.items()
            },
            "execution_order": self.architecture.get("execution_order", []),
        }

add_component 

add_component(name, component, metadata)

Add a component to the model.

Parameters:

Name Type Description Default
name str

Name for this component in the architecture

 required
component Any

The component (PyTorch module or ExtractedComponent)

 required
metadata Dict[str, Any]

Component metadata with input/output dimensions

 required
Source code in src/model_garage/compose/builder.py 
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
def add_component(self, name: str, component: Any, metadata: Dict[str, Any]) -> None:
    """
    Add a component to the model.

    Args:
        name: Name for this component in the architecture
        component: The component (PyTorch module or ExtractedComponent)
        metadata: Component metadata with input/output dimensions
    """
    if name in self.components:
        logger.warning(f"Component {name} already exists, overwriting")
    self.components[name] = (component, metadata)
    logger.info(
        f"Added {name}: {metadata.get('input_dim')}→{metadata.get('output_dim')}"
    )

add_adapter 

add_adapter(source, target, adapter)

Add a dimension adapter between components.

Parameters:

Name Type Description Default
source str

Name of source component

 required
target str

Name of target component

 required
adapter Any

Linear adapter module

 required
Source code in src/model_garage/compose/builder.py 
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
def add_adapter(self, source: str, target: str, adapter: Any) -> None:
    """
    Add a dimension adapter between components.

    Args:
        source: Name of source component
        target: Name of target component
        adapter: Linear adapter module
    """
    if source not in self.components:
        raise ValueError(f"Source component {source} does not exist")
    if target not in self.components:
        raise ValueError(f"Target component {target} does not exist")

    _, source_meta = self.components[source]
    _, target_meta = self.components[target]

    adapter_meta = {
        "input_dim": source_meta.get("output_dim"),
        "output_dim": target_meta.get("input_dim"),
    }

    self.adapters[(source, target)] = (adapter, adapter_meta)

    if source not in self.execution_graph:
        self.execution_graph[source] = []
    self.execution_graph[source].append(target)

define_forward_pass 

define_forward_pass(execution_order)

Define the execution order for the forward pass.

Source code in src/model_garage/compose/builder.py 
71
72
73
74
75
76
def define_forward_pass(self, execution_order: List[str]) -> None:
    """Define the execution order for the forward pass."""
    for name in execution_order:
        if name not in self.components:
            raise ValueError(f"Component {name} not found")
    self.architecture["execution_order"] = execution_order

build 

build()

Build the model from components.

Returns:

Type Description
Dict[str, Any]

Dict containing the assembled model specification.
Source code in src/model_garage/compose/builder.py 
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
def build(self) -> Dict[str, Any]:
    """
    Build the model from components.

    Returns:
        Dict containing the assembled model specification.
    """
    is_valid, errors = self.validate_architecture()
    if not is_valid:
        raise ValueError(f"Invalid architecture: {'; '.join(errors)}")

    logger.info(f"Building {self.name}")

    return {
        "name": self.name,
        "components": {n: c[0] for n, c in self.components.items()},
        "adapters": {
            f"{s}->{t}": a[0] for (s, t), a in self.adapters.items()
        },
        "execution_order": self.architecture.get("execution_order", []),
    }

BaseModelBuilder

BaseModelBuilder

Bases: ABC

Base class for model builders.

Defines the interface for combining extracted components into new hybrid architectures.

Source code in src/model_garage/compose/base.py 
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
class BaseModelBuilder(ABC):
    """
    Base class for model builders.

    Defines the interface for combining extracted components
    into new hybrid architectures.
    """

    def __init__(self, name: str, framework: str = "pytorch"):
        self.name = name
        self.framework = framework
        self.components = {}
        self.adapters = {}
        self.architecture = {}

    @abstractmethod
    def add_component(self, name: str, component: Any, metadata: Dict[str, Any]) -> None:
        """Add a component to the model."""
        pass

    @abstractmethod
    def add_adapter(self, source: str, target: str, adapter: Any) -> None:
        """Add an adapter between components to match dimensions."""
        pass

    @abstractmethod
    def define_forward_pass(self, execution_order: List[str]) -> None:
        """Define the execution order for the forward pass."""
        pass

    @abstractmethod
    def build(self) -> Any:
        """Build the complete model from components."""
        pass

    def validate_architecture(self) -> Tuple[bool, List[str]]:
        """Validate the model architecture."""
        errors = []

        if "execution_order" not in self.architecture:
            errors.append("No execution order defined")
            return False, errors

        for name in self.architecture["execution_order"]:
            if name not in self.components:
                errors.append(f"Component {name} in execution order doesn't exist")

        for (source, target) in self.adapters:
            if source not in self.components:
                errors.append(f"Source component {source} for adapter doesn't exist")
            if target not in self.components:
                errors.append(f"Target component {target} for adapter doesn't exist")

        return len(errors) == 0, errors

    def save_architecture(self, output_dir: str) -> str:
        """Save the model architecture specification."""
        serializable = {
            "name": self.name,
            "framework": self.framework,
            "components": {
                name: {
                    "type": meta.get("type", "unknown"),
                    "input_dim": meta.get("input_dim"),
                    "output_dim": meta.get("output_dim"),
                    "source": meta.get("source", "custom"),
                }
                for name, (_, meta) in self.components.items()
            },
            "execution_order": self.architecture.get("execution_order", []),
        }

        os.makedirs(output_dir, exist_ok=True)
        path = os.path.join(output_dir, f"{self.name}_architecture.json")

        with open(path, "w") as f:
            json.dump(serializable, f, indent=2)

        return path

add_component abstractmethod 

add_component(name, component, metadata)

Add a component to the model.

Source code in src/model_garage/compose/base.py 
31
32
33
34
@abstractmethod
def add_component(self, name: str, component: Any, metadata: Dict[str, Any]) -> None:
    """Add a component to the model."""
    pass

add_adapter abstractmethod 

add_adapter(source, target, adapter)

Add an adapter between components to match dimensions.

Source code in src/model_garage/compose/base.py 
36
37
38
39
@abstractmethod
def add_adapter(self, source: str, target: str, adapter: Any) -> None:
    """Add an adapter between components to match dimensions."""
    pass

define_forward_pass abstractmethod 

define_forward_pass(execution_order)

Define the execution order for the forward pass.

Source code in src/model_garage/compose/base.py 
41
42
43
44
@abstractmethod
def define_forward_pass(self, execution_order: List[str]) -> None:
    """Define the execution order for the forward pass."""
    pass

build abstractmethod 

build()

Build the complete model from components.

Source code in src/model_garage/compose/base.py 
46
47
48
49
@abstractmethod
def build(self) -> Any:
    """Build the complete model from components."""
    pass

validate_architecture 

validate_architecture()

Validate the model architecture.

Source code in src/model_garage/compose/base.py 
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
def validate_architecture(self) -> Tuple[bool, List[str]]:
    """Validate the model architecture."""
    errors = []

    if "execution_order" not in self.architecture:
        errors.append("No execution order defined")
        return False, errors

    for name in self.architecture["execution_order"]:
        if name not in self.components:
            errors.append(f"Component {name} in execution order doesn't exist")

    for (source, target) in self.adapters:
        if source not in self.components:
            errors.append(f"Source component {source} for adapter doesn't exist")
        if target not in self.components:
            errors.append(f"Target component {target} for adapter doesn't exist")

    return len(errors) == 0, errors

save_architecture 

save_architecture(output_dir)

Save the model architecture specification.

Source code in src/model_garage/compose/base.py 
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
def save_architecture(self, output_dir: str) -> str:
    """Save the model architecture specification."""
    serializable = {
        "name": self.name,
        "framework": self.framework,
        "components": {
            name: {
                "type": meta.get("type", "unknown"),
                "input_dim": meta.get("input_dim"),
                "output_dim": meta.get("output_dim"),
                "source": meta.get("source", "custom"),
            }
            for name, (_, meta) in self.components.items()
        },
        "execution_order": self.architecture.get("execution_order", []),
    }

    os.makedirs(output_dir, exist_ok=True)
    path = os.path.join(output_dir, f"{self.name}_architecture.json")

    with open(path, "w") as f:
        json.dump(serializable, f, indent=2)

    return path