Utils

A set of helper functions for processing HF datasets.

lotd.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.

lotd.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.

lotd.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.

lotd.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.