def METHOD_NAME(  # pylint: disable=invalid-name
    extracts: beam.pvalue.PCollection,
    model_agnostic_config: agnostic_predict.ModelAgnosticConfig,
    desired_batch_size: Optional[int] = None) -> beam.pvalue.PCollection:
  """A PTransform that generates features, predictions, labels.
  Args:
    extracts: PCollection of Extracts containing a serialized example to be fed
      to the model.
    model_agnostic_config: A config specifying how to extract
      FeaturesPredictionsLabels from the input input Extracts.
    desired_batch_size: Optional batch size for batching in Aggregate.
  Returns:
    PCollection of Extracts, where the extracts contains the features,
    predictions, labels retrieved.
  """
  batch_args = {}
  if desired_batch_size:
    batch_args = dict(
        min_batch_size=desired_batch_size, max_batch_size=desired_batch_size)
  return (extracts
          | 'Batch' >> beam.BatchElements(**batch_args)
          | 'ModelAgnosticExtract' >> beam.ParDo(
              _ModelAgnosticExtractDoFn(
                  model_agnostic_config=model_agnostic_config)))