code: def transform_model_gen(
      session,
      predict_saved_model,
      style_dataset,
      content_dataset,
      sample_count,
  ):
      """Create a generator function that emits content images & style bottlenecks.
  
      Args:
        session: tf.Session, the session that contains subgraph to load the training
          dataset.
        predict_saved_model: str, path to the style predict SavedModel.
        style_dataset: tf.data.Dataset that contains training style images.
        content_dataset: tf.data.Dataset that contains training style images.
        sample_count: int, number of sample to create.
  
      Returns:
        (str, str) A generator function to use as representative dataset for
        TFLiteConverter.
      """
  
      # Calculate style bottleneck in advance for representative dataset
      style_bottleneck_list = calculate_style_bottleneck(
          session,
          predict_saved_model,
          style_dataset,
          min_sample_count=sample_count,
      )
  
      def generator():
          """A generator to be used as representative_dataset for TFLiteConverter."""
          # Get ImageNet data to use as content_image representative dataset
          dataset = content_dataset.batch(1)
          iterator = tf.data.make_initializable_iterator(dataset)
          session.run(iterator.initializer)
          next_element = iterator.get_next()
  
          # Generate representative dataset
          for index in range(sample_count):
              content_image = session.run(next_element)
              style_bottleneck_input = np.expand_dims(
                  style_bottleneck_list[index], axis=0
              )
              yield [content_image, style_bottleneck_input]
  
      return generator
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import os
  import tempfile
  from magenta.models.arbitrary_image_stylization import arbitrary_image_stylization_build_mobilenet_model
  import numpy
  import tensorflow.compat.v1
label_desc: Use image node_id 10 for calling a global function with example usage: transform_model_gen(session, predict_saved_model, style_dataset, content_dataset, sample_count) and returns: generator