feat(models): model config dataclass and architecture registry #18

New Issue

forbes · 2026-02-04T20:22:22Z

forbes commented

2026-02-04 20:22:22 +00:00

Summary

Implement solver/models/config.py with config parsing and an encoder registry pattern.

Components

ModelConfig dataclass

Parsed from YAML configs (configs/model/baseline.yaml, configs/model/gat.yaml)
Fields: name, architecture, encoder (EncoderConfig), node_features_dim, edge_features_dim, heads (dict of HeadConfig)

HeadConfig dataclass

Fields: enabled, hidden_dim, num_classes (optional)

EncoderConfig dataclass

Fields: num_layers, hidden_dim, dropout, num_heads (optional), residual (optional)

Registry

@register_encoder(name) decorator for registering encoder classes
build_model(config: ModelConfig) -> nn.Module factory function that resolves architecture name to encoder class

Files

solver/models/config.py
tests/models/test_config.py

Acceptance criteria

Parses both baseline.yaml and gat.yaml into ModelConfig
build_model raises ValueError on unknown architecture
Config round-trips correctly (to_dict / from_dict)
Passes ruff, mypy, pytest

## Summary Implement `solver/models/config.py` with config parsing and an encoder registry pattern. ## Components ### ModelConfig dataclass - Parsed from YAML configs (`configs/model/baseline.yaml`, `configs/model/gat.yaml`) - Fields: `name`, `architecture`, `encoder` (EncoderConfig), `node_features_dim`, `edge_features_dim`, `heads` (dict of HeadConfig) ### HeadConfig dataclass - Fields: `enabled`, `hidden_dim`, `num_classes` (optional) ### EncoderConfig dataclass - Fields: `num_layers`, `hidden_dim`, `dropout`, `num_heads` (optional), `residual` (optional) ### Registry - `@register_encoder(name)` decorator for registering encoder classes - `build_model(config: ModelConfig) -> nn.Module` factory function that resolves architecture name to encoder class ## Files - `solver/models/config.py` - `tests/models/test_config.py` ## Acceptance criteria - [ ] Parses both `baseline.yaml` and `gat.yaml` into `ModelConfig` - [ ] `build_model` raises `ValueError` on unknown architecture - [ ] Config round-trips correctly (to_dict / from_dict) - [ ] Passes ruff, mypy, pytest

forbes referenced this issue

2026-02-04 20:22:33 +00:00

feat(models): GIN encoder for assembly graphs #19

forbes referenced this issue

2026-02-04 20:22:41 +00:00

feat(models): GAT encoder with multi-head attention and residual connections #20

forbes referenced this issue

2026-02-04 20:22:53 +00:00

feat(models): multi-task prediction heads for constraint analysis #21

forbes referenced this issue

2026-02-04 20:23:08 +00:00

feat(models): assembled multi-task model and loss function #22