/architecture-design

name: architecture-design
description: Use only when creating new registrable ML components that require Factory or Registry patterns.
version: 1.2.0

Architecture Design - ML Project Template

This skill defines the standard code architecture for machine learning projects based on the template structure. When modifying or extending code, follow these patterns to maintain consistency.

Overview

The project follows a modular, extensible architecture with clear separation of concerns. Each module (data, model, trainer, analysis) is independently organized using factory and registry patterns for maximum flexibility.

When to Use

Use this skill when:

• Creating a new Dataset class that needs `@register_dataset`
• Creating a new Model class that needs `@register_model`
• Creating a new module directory with `__init__.py` factory wiring
• Initializing a new ML project structure from scratch
• Adding new component types such as Augmentation, CollateFunction, or Metrics

When Not to Use

Do not use this skill when:

• Modifying existing functions or methods
• Fixing bugs in existing code
• Adding helper functions or utilities
• Refactoring without adding new registrable components
• Making simple code changes to a single file
• Modifying configuration files
• Reading or understanding existing code

Key indicator: if the task does not require a `@register_*` decorator or a Factory pattern, skip this skill.

Core Design Patterns

Factory Pattern

Each module uses a factory to create instances dynamically:

# Example from data_module/dataset/__init__.py
DATASET_FACTORY: Dict = {}

def DatasetFactory(data_name: str):
    dataset = DATASET_FACTORY.get(data_name, None)
    if dataset is None:
        print(f"{data_name} dataset is not implementation, use simple dataset")
        dataset = DATASET_FACTORY.get('simple')
    return dataset

For detailed guidance, refer to `references/factory_pattern.md`.

Registry Pattern

Components register themselves via decorators:

# Example from data_module/dataset/simple_dataset.py
@register_dataset("simple")
class SimpleDataset(Dataset):
    def __init__(self, data):
        self.data = data

For detailed guidance, refer to `references/registry_pattern.md`.

Auto-Import Pattern

Modules automatically discover and import submodules:

# Example from data_module/dataset/__init__.py
models_dir = os.path.dirname(__file__)
import_modules(models_dir, "src.data_module.dataset")

For detailed guidance, refer to `references/auto_import.md`.

Directory Structure

project/
├── run/
│   ├── pipeline/            # Main workflow scripts
│   │   ├── training/        # Training pipelines
│   │   ├── prepare_data/    # Data preparation pipelines
│   │   └── analysis/        # Analysis pipelines
│   └── conf/                # Hydra configuration files
│       ├── training/        # Training configs
│       ├── dataset/         # Dataset configs
│       ├── model/           # Model configs
│       ├── prepare_data/    # Data prep configs
│       └── analysis/        # Analysis configs
│
├── src/
│   ├── data_module/         # Data processing module
│   │   ├── dataset/         # Dataset implementations
│   │   ├── augmentation/    # Data augmentation
│   │   ├── collate_fn/      # Collate functions
│   │   ├── compute_metrics/ # Metrics computation
│   │   ├── prepare_data/    # Data preparation logic
│   │   ├── data_func/       # Data utility functions
│   │   └── utils.py         # Module-specific utilities
│   │
│   ├── model_module/        # Model implementations
│   │   ├── brain_decoder/   # Brain decoder models
│   │   └── model/           # Alternative model location
│   │
│   ├── trainer_module/      # Training logic
│   ├── analysis_module/     # Analysis and evaluation
│   ├── llm/                 # LLM-related code
│   └── utils/               # Shared utilities
│
├── data/
│   ├── raw/                 # Original, immutable data
│   ├── processed/           # Cleaned, transformed data
│   └── external/            # Third-party data
│
├── outputs/
│   ├── logs/                # Training and evaluation logs
│   ├── checkpoints/         # Model checkpoints
│   ├── tables/              # Result tables
│   └── figures/             # Plots and visualizations
│
├── pyproject.toml           # Project configuration
├── uv.lock                  # Dependency lock file
├── TODO.md                  # Task tracking
├── README.md                # Project documentation
└── .gitignore               # Git ignore rules

For detailed directory structure with file descriptions, refer to `references/structure.md`.

Module Organization

Creating a New Dataset

When adding a new dataset:

1. Create file in `src/data_module/dataset/` 2. Use `@register_dataset("name")` decorator 3. Inherit from `torch.utils.data.Dataset` 4. Implement `__init__`, `__len__`, `__getitem__`

from torch.utils.data import Dataset
from typing import Dict
import torch
from src.data_module.dataset import register_dataset

@register_dataset("custom")
class CustomDataset(Dataset):
    def __init__(self, data):
        self.data = data

    def __len__(self):
        return len(self.data)

    def __getitem__(self, i: int) -> Dict[str, torch.Tensor]:
        return self.data[i]

Creating a New Model

**CRITICAL: Models use config-driven pattern**

When adding a new model:

1. Create file in `src/model_module/model/` or appropriate module subdirectory 2. Use `@register_model('ModelName')` decorator 3. `__init__` accepts **ONLY** `cfg` parameter - all hyperparameters come from config 4. `forward()` returns dict: `{"loss": loss, "labels": labels, "logits": logits}` 5. Handle training vs inference modes using `self.training`

from src.model_module.brain_decoder import register_model

@register_model('MyModel')
class MyModel(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg