def METHOD_NAME(
    text_and_mask: Tuple[str, str],
    tokenizer: AutoTokenizer) -> Tuple[str, Dict[str, torch.Tensor]]:
  text, masked_text = text_and_mask
  tokenized_sentence = tokenizer.encode_plus(masked_text, return_tensors="pt")
  # Workaround to manually remove batch dim until we have the feature to
  # add optional batching flag.
  # TODO(https://github.com/apache/beam/issues/21863): Remove once optional
  # batching flag added
  return text, {
      k: torch.squeeze(v)
      for k, v in dict(tokenized_sentence).items()
  }