def testVocabulary(self,
                     x_data,
                     x_feature_spec,
                     label_data=None,
                     label_feature_spec=None,
                     weight_data=None,
                     weight_feature_spec=None,
                     expected_vocab_file_contents=None,
                     **kwargs):
    """Test tft.Vocabulary with various inputs."""

    input_data = [{'x': x} for x in x_data]
    input_feature_spec = {'x': x_feature_spec}

    if label_data is not None:
      for idx, label in enumerate(label_data):
        input_data[idx]['label'] = label
      input_feature_spec['label'] = label_feature_spec

    if weight_data is not None:
      for idx, weight in enumerate(weight_data):
        input_data[idx]['weights'] = weight
      input_feature_spec['weights'] = weight_feature_spec

    input_metadata = tft.DatasetMetadata.from_feature_spec(input_feature_spec)

    def preprocessing_fn(inputs):
      x = inputs['x']
      labels = inputs.get('label')
      weights = inputs.get('weights')
      # Note even though the return value is not used, calling tft.vocabulary
      # will generate the vocabulary as a side effect, and since we have named
      # this vocabulary it can be looked up using public APIs.
      tft.vocabulary(
          x,
          labels=labels,
          weights=weights,
          vocab_filename='my_vocab',
          file_format=self._VocabFormat(),
          **kwargs)
      return inputs

    self.assertAnalyzeAndTransformResults(
        input_data,
        input_metadata,
        preprocessing_fn,
        input_data,  # expected output data is same as input data
        input_metadata,  # expected output metadata is same as input metadata
        expected_vocab_file_contents={'my_vocab': expected_vocab_file_contents})
Test tft.Vocabulary with various inputs.
tensorflow_transform/beam/vocabulary_integration_test.py