# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
"""isort:skip_file"""

from .dictionary import Dictionary, TruncatedDictionary

from .base_dataset import BaseDataset
from .base_wrapper_dataset import BaseWrapperDataset
from .append_token_dataset import AppendTokenDataset
from .cm3_dataset import CausalMaskedDocumentToSequenceDataset
from .concat_dataset import ConcatDataset
from .id_dataset import IdDataset
from .indexed_dataset import (
    IndexedCachedDataset,
    IndexedDataset,
    IndexedRawTextDataset,
    MMapIndexedDataset,
)
from .jsonl_dataset import JsonlDataset
from .list_dataset import ListDataset
from .lm_context_window_dataset import LMContextWindowDataset
from .monolingual_dataset import MonolingualDataset
from .nested_dictionary_dataset import NestedDictionaryDataset
from .numel_dataset import NumelDataset
from .pad_dataset import LeftPadDataset, PadDataset, RightPadDataset
from .partitioned_streaming_dataset import PartitionedStreamingDataset
from .prepend_token_dataset import PrependTokenDataset
from .resampling_dataset import ResamplingDataset
from .sort_dataset import SortDataset
from .streaming_shuffle_dataset import StreamingShuffleDataset
from .streaming_token_block_dataset import StreamingTokenBlockDataset
from .streaming_src_tgt_dataset import StreamingSrcTgtDataset
from .strip_token_dataset import StripTokenDataset
from .token_block_dataset import TokenBlockDataset
from .pad_dataset import MultiplePadDataset
from .shorten_dataset import TruncateDataset

from .iterators import (
    CountingIterator,
    EpochBatchIterator,
    GroupedIterator,
    ShardedIterator,
)

__all__ = [
    "AppendTokenDataset",
    "BaseWrapperDataset",
    "CausalMaskedDocumentToSequenceDataset",
    "ConcatDataset",
    "CountingIterator",
    "Dictionary",
    "EpochBatchIterator",
    "BaseDataset",
    "GroupedIterator",
    "IdDataset",
    "IndexedCachedDataset",
    "IndexedDataset",
    "IndexedRawTextDataset",
    "JsonlDataset",
    "LeftPadDataset",
    "ListDataset",
    "LMContextWindowDataset",
    "MMapIndexedDataset",
    "MonolingualDataset",
    "MultiplePadDataset",
    "NestedDictionaryDataset",
    "NumelDataset",
    "PadDataset",
    "PartitionedStreamingDataset",
    "PrependTokenDataset",
    "ResamplingDataset",
    "RightPadDataset",
    "ShardedIterator",
    "SortDataset",
    "StreamingShuffleDataset",
    "StreamingTokenBlockDataset",
    "StreamingSrcTgtDataset",
    "StripTokenDataset",
    "TokenBlockDataset",
    "TruncateDataset",
    "TruncatedDictionary",
]
