def _get_dfidf_experimental_preprocessing_fn(self,
                                               is_str_input: bool = False,
                                               smooth: bool = True,
                                               add_baseline: bool = True,
                                               vocab_size: Optional[int] = None,
                                               top_k: Optional[int] = None):
    """Returns proper preprocessing fn for df/idf under tft.experimental."""

    def preprocessing_fn(inputs):
      if is_str_input:
        inputs_as_ints = tft.compute_and_apply_vocabulary(
            tf.compat.v1.strings.split(inputs['a']), top_k=top_k)
      else:
        inputs_as_ints = inputs['a']

      if vocab_size is None:
        computed_vocab_size = tft.get_num_buckets_for_transformed_feature(
            inputs_as_ints)
      else:
        computed_vocab_size = vocab_size

      out_df_counts = tft.experimental.document_frequency(
          inputs_as_ints, computed_vocab_size)
      out_idf_weights = tft.experimental.idf(
          inputs_as_ints,
          computed_vocab_size,
          smooth=smooth,
          add_baseline=add_baseline)
      return {'df': out_df_counts, 'idf': out_idf_weights}

    return preprocessing_fn
Returns proper preprocessing fn for df/idf under tft.experimental.
tensorflow_transform/beam/impl_test.py