def _get_data_loader_config()

in src/fmeval/data_loaders/util.py [0:0]

• 29 lines of code
• 5 McCabe index (conditional complexity)

def _get_data_loader_config(data_source: DataSource, config: DataConfig) -> JsonDataLoaderConfig:
    """
    Returns a dataloader config based on the dataset MIME type specified in `config`.

    :param data_source: The dataset's DataSource object.
    :param config: Configures the returned dataloader config.
    :returns: A dataloader config object, created from `data_source` and `config`.
    """
    if config.dataset_mime_type == MIME_TYPE_JSON:
        if not isinstance(data_source, DataFile):
            raise EvalAlgorithmInternalError(
                f"JSON datasets must be stored in a single file. " f"Provided dataset has type {type(data_source)}."
            )
        return JsonDataLoaderConfig(
            parser=JsonParser(config),
            data_file=data_source,
            dataset_mime_type=MIME_TYPE_JSON,
            dataset_name=config.dataset_name,
        )
    elif config.dataset_mime_type == MIME_TYPE_JSONLINES:
        if not isinstance(data_source, DataFile):
            raise EvalAlgorithmInternalError(
                f"JSONLines datasets must be stored in a single file. "
                f"Provided dataset has type {type(data_source)}."
            )
        return JsonDataLoaderConfig(
            parser=JsonParser(config),
            data_file=data_source,
            dataset_mime_type=MIME_TYPE_JSONLINES,
            dataset_name=config.dataset_name,
        )
    else:  # pragma: no cover
        raise EvalAlgorithmInternalError(
            "Dataset MIME types other than JSON and JSON Lines are not supported. "
            f"MIME type detected from config is {config.dataset_mime_type}."
        )