model_config.py

from __future__ import annotations

from typing import Protocol


class AdditionModelConfig(Protocol):
    d_model: int
    n_heads: int
    d_mlp: int
    n_layers: int
    max_seq_len: int
    init_mode: str
    residual_alpha: float