ColossalAI/applications/ColossalChat/coati/dataset/__init__.py

27 lines
814 B
Python
Executable File

from .conversation import Conversation, setup_conversation_template
from .loader import (
DataCollatorForPreferenceDataset,
DataCollatorForPromptDataset,
DataCollatorForSupervisedDataset,
StatefulDistributedSampler,
load_tokenized_dataset,
setup_distributed_dataloader,
)
from .tokenization_utils import supervised_tokenize_sft, tokenize_prompt_dataset, tokenize_rlhf
__all__ = [
"tokenize_prompt_dataset",
"DataCollatorForPromptDataset",
"is_rank_0",
"DataCollatorForPreferenceDataset",
"DataCollatorForSupervisedDataset",
"StatefulDistributedSampler",
"load_tokenized_dataset",
"setup_distributed_dataloader",
"supervised_tokenize_pretrain",
"supervised_tokenize_sft",
"tokenize_rlhf",
"setup_conversation_template",
"Conversation",
]