ProkBERT DataCollator

class prokbert.ProkBERTDataCollator.ProkBERTDataCollator(tokenizer: PreTrainedTokenizerBase, mlm: bool = True, whole_word_mask: bool = False, mlm_probability: float | None = 0.15, mask_replace_prob: float = 0.8, random_replace_prob: float = 0.1, pad_to_multiple_of: Optional[int] = None, return_tensors: str = 'pt', seed: Optional[int] = None, mask_to_left: int = 0, mask_to_right: int = 0, replace_prob: float = 0.8, random_prob: float = 0.1)

Bases: DataCollatorForLanguageModeling

Data collator for overlapping k-mers.

Args:

mask_to_left (int): Number of tokens masked to the left of a given token. Default is 0. mask_to_right (int): Number of tokens masked to the right of a given token. Default is 0. replace_prob (float): Probability of replacing a token. Default is 0.8. random_prob (float): Probability of changing a token to a random token. Default is 0.1.

set_mask_neighborhood_params(mask_to_left: int = 0, mask_to_right: int = 0)

Set the number of tokens that should be masked to the left and right of a given masked token.

Args:

mask_to_left (int): Number of tokens to be masked to the left. Default is 0. mask_to_right (int): Number of tokens to be masked to the right. Default is 0.

Raises:

AssertionError: If mask_to_left or mask_to_right are not non-negative integers.

torch_mask_tokens(inputs: Tensor, special_tokens_mask: Optional[Tensor] = None) → Tuple[Tensor, Tensor]

Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original.

Args:

inputs (torch.Tensor): Input tensor to be masked. special_tokens_mask (Optional[torch.Tensor]): Tensor representing special tokens (e.g., [CLS], [SEP]). Default is None.

Returns:

Tuple[torch.Tensor, torch.Tensor]: Tuple of input tensor with masked values and labels tensor.