architecture-design

# 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:

```python
# 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:

```python
# 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:

```python
# 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__`

```python
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`

```python
from src.model_module.brain_decoder import register_model

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

        # ALL parameters from cfg
        self.hidden_dim = cfg.model.hidden_dim
        self.output_dim = cfg.da