src/model/encoder/__init__.py

from typing import Optional, Union

from .encoder import Encoder
from .visualization.encoder_visualizer import EncoderVisualizer
from .stylos import EncoderStylos, EncoderStylosCfg

ENCODERS = {
    "stylos": (EncoderStylos, None),
}

EncoderCfg = Union[EncoderStylosCfg]


def get_encoder(cfg: EncoderCfg) -> tuple[Encoder, Optional[EncoderVisualizer]]:
    encoder, visualizer = ENCODERS[cfg.name]
    encoder = encoder(cfg)
    if visualizer is not None:
        visualizer = visualizer(cfg.visualizer, encoder)
    return encoder, visualizer