def test_non_frequency_vocabulary_merge(self, use_tf_compat_v1):
    """This test compares vocabularies produced with and without cache."""
    if not use_tf_compat_v1:
      tft_unit.skip_if_not_tf2('Tensorflow 2.x required.')
    mi_vocab_name = 'mutual_information_vocab'
    adjusted_mi_vocab_name = 'adjusted_mutual_information_vocab'
    weighted_frequency_vocab_name = 'weighted_frequency_vocab'

    def preprocessing_fn(inputs):
      _ = tft.vocabulary(
          inputs['s'],
          labels=inputs['label'],
          store_frequency=True,
          vocab_filename=mi_vocab_name,
          min_diff_from_avg=0.1,
          use_adjusted_mutual_info=False,
          name='with_mi')

      _ = tft.vocabulary(
          inputs['s'],
          labels=inputs['label'],
          store_frequency=True,
          vocab_filename=adjusted_mi_vocab_name,
          min_diff_from_avg=1.0,
          use_adjusted_mutual_info=True,
          name='with_adjusted_mi')

      _ = tft.vocabulary(
          inputs['s'],
          weights=inputs['weight'],
          store_frequency=True,
          vocab_filename=weighted_frequency_vocab_name,
          use_adjusted_mutual_info=False,
          name='with_weight')
      return inputs

    span_0_key = analyzer_cache.DatasetKey('span-0')
    span_1_key = analyzer_cache.DatasetKey('span-1')

    input_data = [
        dict(s='a', weight=1, label=1),
        dict(s='a', weight=0.5, label=1),
        dict(s='b', weight=0.75, label=1),
        dict(s='b', weight=1, label=0),
    ]
    feature_spec = {
        's': tf.io.FixedLenFeature([], tf.string),
        'label': tf.io.FixedLenFeature([], tf.int64),
        'weight': tf.io.FixedLenFeature([], tf.float32),
    }
    input_data_dict = {
        span_0_key: input_data,
        span_1_key: input_data,
    }
    transform_fn_with_cache_dir = os.path.join(self.base_test_dir,
                                               'transform_fn_with_cache')

    expected_accumulators = {
        _make_cache_key(b'VocabularyAccumulate[with_mi]'): [
            _encode_vocabulary_accumulator(b'a',
                                           b'[2, [0.0, 1.0], [0.0, 0.0], 1.0]'),
            _encode_vocabulary_accumulator(b'b',
                                           b'[2, [0.5, 0.5], [0.0, 0.0], 1.0]'),
            _encode_vocabulary_accumulator(
                b'global_y_count_sentinel',
                b'[4, [0.25, 0.75], [0.0, 0.0], 1.0]'),
        ],
        _make_cache_key(b'VocabularyAccumulate[with_adjusted_mi]'): [
            _encode_vocabulary_accumulator(b'a',
                                           b'[2, [0.0, 1.0], [0.0, 0.0], 1.0]'),
            _encode_vocabulary_accumulator(b'b',
                                           b'[2, [0.5, 0.5], [0.0, 0.0], 1.0]'),
            _encode_vocabulary_accumulator(
                b'global_y_count_sentinel',
                b'[4, [0.25, 0.75], [0.0, 0.0], 1.0]'),
        ],
        _make_cache_key(b'VocabularyAccumulate[with_weight]'): [
            _encode_vocabulary_accumulator(b'a', b'1.5'),
            _encode_vocabulary_accumulator(b'b', b'1.75')
        ],
    }
    expected_cache = {
        span: expected_accumulators for span in [span_0_key, span_1_key]
    }

    run_result = self._run_pipeline(
        feature_spec,
        input_data_dict,
        preprocessing_fn,
        transform_fn_output_dir=transform_fn_with_cache_dir,
        expected_cache=expected_cache,
        use_tf_compat_v1=use_tf_compat_v1)

    metrics = run_result.metrics
    # 4 from analysis on each of the input spans.
    self.assertMetricsCounterEqual(metrics, 'num_instances', 8)
    self.assertMetricsCounterEqual(metrics, 'cache_entries_decoded', 0)
    self.assertMetricsCounterEqual(metrics, 'cache_entries_encoded', 6)
    self.assertMetricsCounterEqual(
        metrics, 'saved_models_created', _SINGLE_PHASE_NUM_SAVED_MODELS
    )
    self.assertMetricsCounterEqual(
        metrics, 'analysis_input_bytes_from_cache', 0
    )

    with self._TestPipeline() as p:
      with tft_beam.Context():
        flat_data = p | 'CreateInputData' >> beam.Create(input_data * 2)

        input_metadata = dataset_metadata.DatasetMetadata.from_feature_spec(
            feature_spec)
        transform_fn_no_cache = ((flat_data, input_metadata)
                                 | tft_beam.AnalyzeDataset(preprocessing_fn))

        transform_fn_no_cache_dir = os.path.join(self.base_test_dir,
                                                 'transform_fn_no_cache')
        _ = transform_fn_no_cache | tft_beam.WriteTransformFn(
            transform_fn_no_cache_dir)

    # 4 from analysis on each of the input spans.
    self.assertMetricsCounterEqual(p.metrics, 'num_instances', 8)
    self.assertMetricsCounterEqual(p.metrics, 'cache_entries_decoded', 0)
    self.assertMetricsCounterEqual(p.metrics, 'cache_entries_encoded', 0)
    self.assertMetricsCounterEqual(p.metrics, 'saved_models_created',
                                   _SINGLE_PHASE_NUM_SAVED_MODELS)
    self.assertMetricsCounterEqual(p.metrics, 'analysis_input_bytes_from_cache',
                                   0)

    tft_output_cache = tft.TFTransformOutput(transform_fn_with_cache_dir)
    tft_output_no_cache = tft.TFTransformOutput(transform_fn_no_cache_dir)

    for vocab_filename in (mi_vocab_name, adjusted_mi_vocab_name,
                           weighted_frequency_vocab_name):
      cache_path = tft_output_cache.vocabulary_file_by_name(vocab_filename)
      no_cache_path = tft_output_no_cache.vocabulary_file_by_name(
          vocab_filename)
      with tf.io.gfile.GFile(cache_path, 'rb') as f1, tf.io.gfile.GFile(
          no_cache_path, 'rb') as f2:
        self.assertEqual(
            f1.readlines(), f2.readlines(),
            'vocab with cache != vocab without cache for: {}'.format(
                vocab_filename))
This test compares vocabularies produced with and without cache.
tensorflow_transform/beam/cached_impl_test.py