def transform_data(working_dir: str):
  """Transform the data and write out as a TFRecord of Example protos.

  Read in the data from the positive and negative examples on disk, and
  transform it using a preprocessing pipeline that removes punctuation,
  tokenizes and maps tokens to int64 values indices.

  Args:
    working_dir: Directory to read shuffled data from and write transformed data
        and metadata to.
  """

  with beam.Pipeline() as pipeline:
    with tft_beam.Context(
        temp_dir=os.path.join(working_dir, TRANSFORM_TEMP_DIR)):
      tfxio_train_data = tfxio.TFExampleRecord(
          file_pattern=os.path.join(working_dir,
                                    SHUFFLED_TRAIN_DATA_FILEBASE + '*'),
          schema=SCHEMA)
      train_data = (
          pipeline | 'TFXIORead[Train]' >> tfxio_train_data.BeamSource())

      tfxio_test_data = tfxio.TFExampleRecord(
          file_pattern=os.path.join(working_dir,
                                    SHUFFLED_TEST_DATA_FILEBASE + '*'),
          schema=SCHEMA)
      test_data = (pipeline | 'TFXIORead[Test]' >> tfxio_test_data.BeamSource())

      def preprocessing_fn(inputs):
        """Preprocess input columns into transformed columns."""
        review = inputs[REVIEW_KEY]

        # Here tf.compat.v1.string_split behaves differently from
        # tf.strings.split.
        review_tokens = tf.compat.v1.string_split(review, DELIMITERS)
        review_indices = tft.compute_and_apply_vocabulary(
            review_tokens, top_k=VOCAB_SIZE)
        # Add one for the oov bucket created by compute_and_apply_vocabulary.
        review_bow_indices, review_weight = tft.tfidf(review_indices,
                                                      VOCAB_SIZE + 1)
        return {
            REVIEW_KEY: review_bow_indices,
            REVIEW_WEIGHT_KEY: review_weight,
            LABEL_KEY: tf.one_hot(inputs[LABEL_KEY], 2),
        }

      # The TFXIO output format is chosen for improved performance.
      transformed_train_data, transform_fn = (
          (train_data, tfxio_train_data.TensorAdapterConfig())
          | 'AnalyzeAndTransform' >> tft_beam.AnalyzeAndTransformDataset(
              preprocessing_fn, output_record_batches=True))

      transformed_test_data = (
          ((test_data, tfxio_test_data.TensorAdapterConfig()), transform_fn)
          |
          'Transform' >> tft_beam.TransformDataset(output_record_batches=True))

      # Extract transformed RecordBatches, encode and write them to the given
      # directory.
      _ = (
          transformed_train_data
          | 'EncodeTrainData' >> tft_beam.EncodeTransformedDataset()
          | 'WriteTrainData' >> beam.io.WriteToTFRecord(
              os.path.join(working_dir, TRANSFORMED_TRAIN_DATA_FILEBASE)))

      _ = (
          transformed_test_data
          | 'EncodeTestData' >> tft_beam.EncodeTransformedDataset()
          | 'WriteTestData' >> beam.io.WriteToTFRecord(
              os.path.join(working_dir, TRANSFORMED_TEST_DATA_FILEBASE)))

      # Will write a SavedModel and metadata to two subdirectories of
      # working_dir, given by tft.TRANSFORM_FN_DIR and
      # tft.TRANSFORMED_METADATA_DIR respectively.
      _ = (
          transform_fn
          | 'WriteTransformFn' >>
          tft_beam.WriteTransformFn(working_dir))
Transform the data and write out as a TFRecord of Example protos.

  Read in the data from the positive and negative examples on disk, and
  transform it using a preprocessing pipeline that removes punctuation,
  tokenizes and maps tokens to int64 values indices.

  Args:
    working_dir: Directory to read shuffled data from and write transformed data
        and metadata to.
  
examples/sentiment_example_v2.py