# 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 abc import ABC, abstractmethod
from collections.abc import Iterator
from typing import Any, Literal

from pydantic import BaseModel

Row = dict[str, Any]
Dataset = list[Row]


class DatasetSource(BaseModel):
    type: str


class HuggingfaceDataSource(DatasetSource):
    type: Literal["huggingface"] = "huggingface"
    args: dict[str, Any]
    subsets: list[str] | None = None


class IDataset(ABC):
    def __init__(self, dataset_id: str):
        self.dataset_id = dataset_id

    @abstractmethod
    def load(self) -> Dataset:
        raise NotImplementedError()

    @abstractmethod
    def __iter__(self) -> Iterator[Row]:
        raise NotImplementedError()

    def __str__(self) -> str:
        return f"{self.__class__.__name__}({self.dataset_id})"
