# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.

from collections.abc import Iterator

from .interface import HuggingfaceDataSource, IDataset, Row


class HuggingFaceDataset(IDataset):
    def __init__(self, dataset_id: str, config: HuggingfaceDataSource):
        super().__init__(dataset_id)
        self.config = config

    def load(self) -> None:
        from datasets import load_dataset

        print("Loading dataset...", self.config)
        if self.config.subsets:
            self.datasets = []
            for subset in self.config.subsets:
                self.datasets.append(load_dataset(**self.config.args, name=subset))
        else:
            self.dataset = load_dataset(**self.config.args)

    def __iter__(self) -> Iterator[Row]:
        if self.config.subsets:
            for dataset, subset in zip(self.datasets, self.config.subsets, strict=False):
                for row in dataset:
                    yield {
                        "subset": subset,
                        **row,
                    }
        else:
            for row in self.dataset:
                yield row

    def __len__(self) -> int:
        if self.config.subsets:
            return sum(len(dataset) for dataset in self.datasets)
        else:
            return len(self.dataset)
