Class DataModule

lit_gpt/datamodules/base.py:16–50

Base class for all data modules in LitGPT.

class DataModule(LightningDataModule):
    """Base class for all data modules in LitGPT."""

    """Whether to mask the prompt section from the label (with ``ignore_index``)."""
    mask_prompt: bool = True
    """The style to apply to instruction prompts. See `litgpt.prompts` for a list of available styles."""
    prompt_style: Union[str, PromptStyle] = 'chatml'
    """The index to use for elements to be ignored in the label."""
    ignore_index: int = -100
    """The random seed for shuffling the dataset."""
    seed: int = 42
    """How many DataLoader processes to use for loading."""
    num_workers = max(1, cpu_count() // 2) # Use half of available cores
    """The directory in which the downloaded dataset gets saved."""
    include_multiturn_conversations: bool = True
    """The repo from where the data is downloaded"""
    repo_id: str = ''

    @abstractmethod
    def connect(
        self,
        tokenizer: Optional[Tokenizer] = None,
        batch_size: int = 1,
        max_seq_length: Optional[int] = None,
    ) -> None:
        """All settings that can't be determined at the time of instantiation need to be passed through here
        before any dataloaders can be accessed.
        """

    def setup(self, stage: str = '') -> PreparedDataset:
        # Stub is to redefine the default signature, because the concept of 'stage' does not exist in LitGPT
        pass

    def __repr__(self) -> str:
        return f'{self.__class__.__name__}()'


def get_sft_collate_fn(