import hydra
import pytorch_lightning as pl
from omegaconf import DictConfig


def instantiate_datamodule(cfg: DictConfig) -> pl.LightningDataModule:
    if cfg.get("dataset"):
        collator = hydra.utils.instantiate(cfg.dataset)
        return hydra.utils.instantiate(cfg.data, collate_fn=collator)
    else:
        return hydra.utils.instantiate(cfg.data)
