architecture-design

What this skill does

# 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.dataset.target_size[cfg.dataset.task]

    def forward(self, x, labels=None, **kwargs):
        if self.training:
            # Training logic
            pass
        else:
            # Inference logic
            pass

        return {"loss": loss, "labels": labels, "logits": logits}
```

### Adding Data Augmentation

When adding augmentation:

1. Create file in `src/data_module/augmentation/`
2. Implement transformation function
3. Register with factory if needed

## Code Style Guidelines

For comprehensive style guidelines, refer to `references/code_style.md`.

**Key principles:**
- Always use type hints for function signatures
- Follow import order: standard library → third-party → local
- Module `__init__.py` files contain factory/registry logic
- Model classes must be config-driven

## Configuration Management

The project uses Hydra for configuration management:

- Config files in `run/conf/` organize by module
- Each stage (training, analysis) has its own config structure
- Use YAML files for all configuration

## When Working on This Project

### Before Modifying Code

1. Read the relevant module's factory/registry pattern
2. Check existing implementations for consistency
3. Follow the established directory structure
4. Use registration decorators for new components

### Adding New Features

1. Determine which module the feature belongs to
2. Check if similar functionality exists
3. Follow factory/registry pattern if creating new component types
4. Add configuration files if needed
5. Update documentation

### Code Review Checklist

- [ ] Uses factory/registry pattern appropriately
- [ ] Follows module directory structure
- [ ] Has proper type annotations
- [ ] Imports are correctly ordered
- [ ] Registration decorator is used
- [ ] Configuration files are added if needed

## Additional Resources

### Reference Files

For detailed information, consult:
- **`references/structure.md`** - Detailed directory structure with file descriptions
- **`references/fact