src/taoTrain/data/__init__.py

"""Dataset implementations and loaders."""

# HuggingFace-based datasets
from .hf_base import BaseHFDataset
from .hf_pretrain import PretrainDataset
from .hf_sft import SFTDataset
from .hf_rl import RLDataset

# JSONL-based datasets (async-only)
from .jsonl_base import BaseJSONLDataset
from .pretrain_jsonl import PretrainJSONLDataset
from .sft_jsonl import SFTJSONLDataset
from .rl_jsonl import RLJSONLDataset
from .vlm_jsonl import VLMJSONLDataset, VLMSFTJSONLDataset

# Utilities
from .tokenizer import SentencePieceTokenizerWrapper
from .sft_utils import parse_sft_record, build_sft_sequence_tokens, apply_response_masking
from .loaders import get_dataloader
from .async_loader import AsyncBatchIterator
from .tokenization_queue import TokenizationQueue
from .factory import DatasetFactory

__all__ = [
    # HuggingFace datasets
    "BaseHFDataset",
    "PretrainDataset",
    "SFTDataset", 
    "RLDataset",
    # JSONL datasets
    "BaseJSONLDataset",
    "PretrainJSONLDataset",
    "SFTJSONLDataset",
    "RLJSONLDataset",
    "VLMJSONLDataset",
    "VLMSFTJSONLDataset",
    # Utilities
    "SentencePieceTokenizerWrapper",
    "parse_sft_record",
    "build_sft_sequence_tokens",
    "apply_response_masking",
    # Data loading
    "get_dataloader",
    "AsyncBatchIterator",
    "TokenizationQueue",
    "DatasetFactory",
]