Dataset collators

Callable objects passed to pytorch dataloaders to prepare process samples in one batch.

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