LOTD - Lord Of The Datasets

Full API overview.

lotd

LOTD - Lord of the Datasets Efficient NLP dataset preprocessing library.

ChatTokenizer

Applies chat template and tokenizes messages in a dataset.

__call__ 

__call__(prompts: Union[List[str], List[List[str]]], responses: Union[List[str], List[List[str]]] = [], system: Union[List[str], None] = None) -> Dict[str, List[List[int]]]

Tokenizes chats.

Parameters:
  • prompts (Union[List[str], List[List[str]]]) –

    a list where each item is either a single prompt or multiple prompts in a dialog.

  • responses (Union[List[str], List[List[str]]], default: [] ) –

    a list where each item is either a single response or multiple responses in a dialog.
Returns:
  • Dict[str, List[List[int]]] –

    dict with input_ids and prompt_mask for all text in a batch.

__init__ 

__init__(tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast], max_length: Union[int, None] = None) -> None

Initializes the dataset chat tokenizer.

Parameters:
  • tokenizer (Union[PreTrainedTokenizer, PreTrainedTokenizerFast]) –

    pre-trained transformers tokenizer.

  • max_length (optional, default: None ) –

    maximum sequence length after truncation.

LengthFilter

Filters datasets by max and min length of input_ids.

__call__ 

__call__(input_ids: List[List[int]]) -> List[bool]

Processes a batch of input_ids.

__init__ 

__init__(min_length: int = 0, max_length: Union[int, None] = None) -> None

Initializes the collator.

Parameters:
  • min_length (int, default: 0 ) –

    lowest number of tokens.

  • max_length (Union[int, None], default: None ) –

    highest number of tokens.

PadCollator

Collator which converts inputs into pytorch tensors and applies padding.

__call__ 

__call__(batch: List[Dict[str, Any]]) -> Dict[str, torch.Tensor]

Collates inputs.

Parameters:
  • batch (List[Dict[str, Any]]) –

    list of dicts where each key is a feature.
Returns:
  • Dict[str, Tensor] –

    a dict of padded tensors of form {"input_ids": ..., "attention_mask": ..., "prompt_mask": ...}.

__init__ 

__init__(pad_id: int, pre: Union[Callable, None] = None, post: Union[Callable, None] = None, padding_side: Literal['right', 'left'] = 'right', ignore_columns: List[str] = []) -> None

Initializes the collator.

Parameters:
  • pad_id (int) –

    pad token id used for padding.

  • pre (Union[Callable, None], default: None ) –

    callable that accepts list of dicts where every key is a dataset feature, modifies it and returns it back.

  • post (Union[Callable, None], default: None ) –

    callable that accepts a dict of tensors of a form {"input_ids": ..., "attention_mask": ..., "prompt_mask": ...}, modifies it and returns it back.

  • padding_side (Literal['right', 'left'], default: 'right' ) –

    can be right or left.

  • ignore_columns (List[str], default: [] ) –

    columns to be excluded.

TextTokenizer

Applies template and tokenizes texts in a dataset.

__call__ 

__call__(texts: List[str]) -> Dict[str, List[List[int]]]

Tokenizes a batch of text samples.

Parameters:
  • texts (List[str]) –

    list of strings.
Returns:
  • Dict[str, List[List[int]]] –

    dict with input_ids and prompt_mask for all text in a batch.

__init__ 

__init__(tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast], template: str = '[CLS]{{text}}[SEP]', max_length: Union[int, None] = None) -> None

Initializes the dataset tokenizer.

Parameters:
  • tokenizer (Union[PreTrainedTokenizer, PreTrainedTokenizerFast]) –

    pre-trained transformers tokenizer.

  • template (str, default: '[CLS]{{text}}[SEP]' ) –

    jinja2 template containing {{text}} keyword.

  • max_length (optional, default: None ) –

    maximum sequence length after truncation.

format_chat 

format_chat(prompts: Union[str, List[str]] = [], responses: Union[str, List[str]] = [], system: Union[str, None] = None) -> List[Dict[str, str]]

Format list of prompts, responses and system messages into a chat.

Parameters:
  • prompts (Union[str, List[str]], default: [] ) –

    a single prompt or a list of prompts provided by a user in a dialog.

  • responses (Union[str, List[str]], default: [] ) –

    a single response or a list of responses producrd by a model in a dialog.

  • system (Union[str, None], default: None ) –

    system prompt used for this chat.
Returns:
  • List[Dict[str, str]] –

    a list of messages of form [{"role": "...", "content": "..."}, ...].

generate_chat_template 

generate_chat_template(last_only: bool = False, wrapper: str = '{{template}}', system_tags: Tuple[str, str] = ('<|system|>', '<|end|>'), user_tags: Tuple[str, str] = ('<|user|>', '<|end|>'), assistant_tags: Tuple[str, str] = ('<|assistant|>', '<|end|>')) -> str

Generate chat template for instruction tuning.

Parameters:
  • last_only (bool, default: False ) –

    only enable generation for last response. If true, assistant_masks ignores previous responses. Default: False.

  • wrapper (str, default: '{{template}}' ) –

    a wrapper for adding text before and after the template. Should contain {{template}}.

  • system_tags (Tuple[str, str], default: ('<|system|>', '<|end|>') ) –

    a pair of tags to mark the beginning and the end of a system role in a dialog.

  • user_tags (Tuple[str, str], default: ('<|user|>', '<|end|>') ) –

    see system_tags.

  • assistant_tags (Tuple[str, str], default: ('<|assistant|>', '<|end|>') ) –

    see system_tags.
Returns:
  • str –

    a jinja2 template string.

get_loaders 

get_loaders(dataset: Dataset, collate_fn: Callable = lambda x: x, batch_size: int = 16, train_size: float = 0.8, val_size: float = 0.1, num_workers: int = 15, seed: int = 42) -> Tuple[DataLoader, DataLoader, DataLoader]

Shortcut to generate pytorch dataloaders (train/val/test) from hf dataset.

Parameters:
  • dataset (Dataset) –

    HF dataset.

  • collate_fn (Callable, default: lambda x: x ) –

    function used for dataset collation.

  • batch_size (int, default: 16 ) –

    batch_size for dataloaders.

  • train_size (float, default: 0.8 ) –

    train split size. see split_dataset.

  • val_size (float, default: 0.1 ) –

    validation split size. see split_dataset.

  • num_workers (int, default: 15 ) –

    number of pytorch dataloader workers.

  • seed (int, default: 42 ) –

    random seed used for splitting.
Returns:
  • Tuple[DataLoader, DataLoader, DataLoader] –

    a tuple with train, validation and test pytorch dataloaders.

Splits dataset and assigns collators automatically.

load_cached 

load_cached(cache_path: str, process_fn: Callable) -> Dataset

Try loading processed dataset from cache. Processes dataset and saves it to cache if pre-cached dataset is not found.

Parameters:
  • cache_path (str) –

    path to load/save dataset.

  • process_fn (Callable) –

    function that will return a new processed dataset if cache is not found.
Returns:
  • Dataset –

    a pre-processed HF dataset.

split_dataset 

split_dataset(dataset: Dataset, train_size: float = 0.8, val_size: float = 0.1, seed: int = 42) -> Tuple[Dataset, Dataset, Dataset]

Split HF dataset into train, validation and test.

Parameters:
  • dataset (Dataset) –

    HF dataset.

  • train_size (float, default: 0.8 ) –

    train ratio from 0 to 1.

  • val_size (float, default: 0.1 ) –

    validation ratio from 0 to 1.

  • seed (int, default: 42 ) –

    seed used for splitting.
Returns:
  • Tuple[Dataset, Dataset, Dataset] –

    a tuple of 3 datasets for train, validation and test.

train_size and val_size are taken from total and their sum should not be more than 1.

Test size would be equal to 1 - train_size - val_size.

strip_features 

strip_features(dataset: Dataset, keep_features: List[str] = ['input_ids', 'prompt_mask']) -> Dataset

Remove all features from dataset except specified ones.

Parameters:
  • dataset (Dataset) –

    HF dataset to purify.

  • keep_features (List[str], default: ['input_ids', 'prompt_mask'] ) –

    list of feature names to keep.
Returns:
  • Dataset –

    a new dataset with only specified features.

Useful for reducing memory usage and cleaning up datasets.

Datasets

lotd.datasets

alpaca 

alpaca(tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast], cache_path: Union[str, None] = None, max_length: int = 512, batch_size: int = 16, seed: int = 42, pre: Union[Callable, None] = None, post: Union[Callable, None] = None) -> Tuple[DataLoader, DataLoader, DataLoader]

Downloads and Pre-processes Alpaca dataset for instruction fine-tuning.

Parameters:
  • tokenizer (Union[PreTrainedTokenizer, PreTrainedTokenizerFast]) –

    transformers tokenizer with chat_template parameter set.

  • cache_path (Union[str, None], default: None ) –

    path to load/save dataset.

  • max_length (int, default: 512 ) –

    maximum sequence length for filter.

  • batch_size (int, default: 16 ) –

    dataloaders batch size.

  • seed (int, default: 42 ) –

    splitting and shuffling seed.

  • pre (Union[Callable, None], default: None ) –

    see PadCollator.

  • post (Union[Callable, None], default: None ) –

    see PadCollator.
Returns:
  • Tuple[DataLoader, DataLoader, DataLoader] –

    a tuple of train, validation and test dataloaders.

build_dataset 

build_dataset(preprocess_func)

Decorator for building dataloaders. Processes or tries loading the dataset if cache_path is specified, creates new PadCollator, and returns dataloaders.

tinystories 

tinystories(tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast], template: str = '[CLS]{{tex}}[SEP]', cache_path: Union[str, None] = None, max_length: int = 512, batch_size: int = 16, seed: int = 42, pre: Union[Callable, None] = None, post: Union[Callable, None] = None) -> Tuple[DataLoader, DataLoader, DataLoader]

Downloads and Pre-processes TinyStories dataset for language modeling. train and validation are merged, shuffled and re-split with 80%/10%/10% ratio.

Parameters:
  • tokenizer (Union[PreTrainedTokenizer, PreTrainedTokenizerFast]) –

    transformers tokenizer.

  • template (str, default: '[CLS]{{tex}}[SEP]' ) –

    prompt template with {{text}} placeholder.

  • cache_path (Union[str, None], default: None ) –

    path to load/save dataset.

  • max_length (int, default: 512 ) –

    maximum sequence length for truncation.

  • batch_size (int, default: 16 ) –

    dataloaders batch size.

  • seed (int, default: 42 ) –

    splitting and shuffling seed.

  • pre (Union[Callable, None], default: None ) –

    see PadCollator.

  • post (Union[Callable, None], default: None ) –

    see PadCollator.
Returns:
  • Tuple[DataLoader, DataLoader, DataLoader] –

    a tuple of train, validation and test dataloaders.