<commit_before># Copyright 2017 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""An implementation of tf.Transform using Beam.

The beam implementation takes a user defined preprocessing function (see
../api.py for how to defined a preprocessing function) and implements it as a
Beam PTransform.

The AnalyzeDataset takes the user's preprocessing function and converts into
a TensorFlow function that can be applied to each row of a dataset.  For
example if the user's preprocessing function describes normalizing a column by
subtracting its mean, the tensorflow function will contain the mean of the
column as a constant, and will subtract this value from each value of the
column.  We refer to the result of AnalyzeDataset as a "transform function".

Since AnalyzeDataset is implemented with beam, it accepts a PCollection that
represents the dataset (see below for the exact format) and returns a singleton
PCollection containing the transform function (as a serialized TF graph).

The TransformDataset PTransform takes a dataset and a transform function, and
returns the transformed dataset where the transform function is applied to each
row of the original dataset.

There is also an AnalyzeAndTransformDataset PTransform that applies
AnalyzeDataset and TransformDataset to the same input dataset, possibly with
optimizations.

Typical usage of these functions is shown below.

def preprocessing_fn(inputs):
  ...

with beam.Pipeline(...) as p:
  with beam_impl.Context(temp_dir=my_temp_dir):
    input = p | beam_impl.read_examples(..., schema)
    transformed, transform_fn = ((input, schema)
        | beam_impl.AnalyzeAndTransformDataset(preprocessing_fn))
    transformed | beam_impl.write_examples_and_metadata(
        examples_path, metadata_path)
    transform_fn | beam_impl.write_transform_fn(transform_fn_path)

Implementation note: TensorFlow code (including our code) makes frequent use of
the default graph.  We want to avoid adding to the default graph, or including
the default graph in our own SavedModel's.  This means that wherever we call
TensorFlow code (or our code that uses the default graph) we should create a
graph and mark it as the default.  This is achieved by identifying the
entrypoints into our code where this happens and creating a
"with ... .as_default()" block.  There are four places this happens.

1) In AnalyzeDataset.expand() which is typically called from the main thread
2) In _GraphStateCommon.__init__ which is called from the worker running
   _RunMetaGraphDoFn
3) In _replace_tensors_with_constant_values, which is called in a beam.Map.
4) In extract_scalar_constants, which is called in a beam.Map.
"""
# TODO(KesterTong): Document data format.
# TODO(KesterTong): Refactor and rename now that "TransformFn" is the path to a
# SavedModel, not an in-memory object.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import copy
import datetime

import apache_beam as beam

from apache_beam.runners.portability import fn_api_runner
from apache_beam.transforms import util
from apache_beam.typehints import Any
from apache_beam.typehints import Dict
from apache_beam.typehints import Iterable
from apache_beam.typehints import List
from apache_beam.typehints import Tuple
from apache_beam.typehints import Union

import numpy as np
import pyarrow as pa
import six
import tensorflow as tf
from tensorflow_transform import common
from tensorflow_transform import graph_tools
from tensorflow_transform import impl_helper
from tensorflow_transform import nodes
from tensorflow_transform import schema_inference
from tensorflow_transform.beam import analysis_graph_builder
from tensorflow_transform.beam import analyzer_cache
from tensorflow_transform.beam import beam_nodes
from tensorflow_transform.beam import common as beam_common
from tensorflow_transform.beam import context
from tensorflow_transform.beam import deep_copy
from tensorflow_transform.beam.tft_beam_io import beam_metadata_io
from tensorflow_transform.coders import example_proto_coder
from tensorflow_transform.saved import saved_transform_io
from tensorflow_transform.saved import saved_transform_io_v2
from tensorflow_transform.tf_metadata import dataset_metadata
from tensorflow_transform.tf_metadata import schema_utils
from tfx_bsl.beam import shared
from tfx_bsl.tfxio import tf_example_record
from tfx_bsl.tfxio.tensor_adapter import TensorAdapter
# TODO(https://issues.apache.org/jira/browse/SPARK-22674): Switch to
# `tfx_namedtuple.namedtuple` or `typing.NamedTuple` once the Spark issue is
# resolved.
from tfx_bsl.types import tfx_namedtuple

Context = context.Context

_CREATE_SAVED_MODEL_COUNTER_NAME = 'saved_models_created'

# For some runners, we rely on Beam to manage concurrency, i.e. we expect it to
# run one session per CPU--so we don't want to proliferate TF threads.
# Nonetheless we provide 4 threads per session for TF ops, 2 inter-
# and 2 intra-thread.  In many cases only 2 of these will be runnable
# at any given time.  This approach oversubscribes a bit to make sure
# the CPUs are really saturated.
_FIXED_PARALLELISM_TF_CONFIG = tf.compat.v1.ConfigProto(
    # TODO(b/36091595): use_per_session_threads is deprecated, but the
    # replacement session_inter_op_thread_pool is experimental; using
    # the former for now.
    use_per_session_threads=True,
    inter_op_parallelism_threads=2,
    intra_op_parallelism_threads=2)

_DEFAULT_TENSORFLOW_CONFIG_BY_BEAM_RUNNER_TYPE = {
    # TODO(katsiapis): Perhaps remove this entry once b/69922446 and b/30837990
    # are resolved.
    beam.runners.DataflowRunner: _FIXED_PARALLELISM_TF_CONFIG,

    beam.runners.DirectRunner: _FIXED_PARALLELISM_TF_CONFIG,
    fn_api_runner.FnApiRunner: _FIXED_PARALLELISM_TF_CONFIG,
}

# TODO(b/123325923): Fix the key type here to agree with the actual keys.
_DATASET_ELEMENT_TYPE = Dict[Any,  # Any -> six.text_type?
                             Union[beam_common.PRIMITIVE_TYPE,
                                   # Arbitrarily-nested lists are allowed.
                                   List[Any], np.generic, np.ndarray]]

# TODO(b/68154497): pylint: disable=no-value-for-parameter


# TODO(b/64956765): Remove this once either the keepalive issue (b/30837990), or
# the mentioned bug above is resolved.
# TODO(zoyahav): Make this a PTransform.
def _clear_shared_state_after_barrier(pipeline, input_barrier):
  """Clears any shared state from within a pipeline context.

  This will only be cleared once input_barrier becomes available.

  Args:
    pipeline: A `beam.Pipeline` object.
    input_barrier: A `PCollection` which the pipeline should wait for.

  Returns:
    An empty `PCollection`.
  """
  empty_pcoll = input_barrier | 'MakeCheapBarrier' >> beam.FlatMap(
      lambda x: None)
  return (pipeline
          | 'PrepareToClearSharedKeepAlives' >> beam.Create([None])
          | 'WaitAndClearSharedKeepAlives' >> beam.Map(
              lambda x, empty_side_input: shared.Shared().acquire(lambda: None),
              beam.pvalue.AsIter(empty_pcoll)))


@beam.ptransform_fn
@beam.typehints.with_input_types(_DATASET_ELEMENT_TYPE)
@beam.typehints.with_output_types(List[_DATASET_ELEMENT_TYPE])
def _BatchElements(pcoll):  # pylint: disable=invalid-name
  """Batches elements either automatically or to the given batch_size."""
  desired_batch_size = Context.get_desired_batch_size()
  kwargs = dict(
      min_batch_size=desired_batch_size, max_batch_size=desired_batch_size
  ) if desired_batch_size is not None else {}
  return pcoll | 'BatchElements' >> util.BatchElements(**kwargs)


# TODO(b/36223892): Verify that these type hints work and make needed fixes.
@beam.typehints.with_input_types(
    Union[List[_DATASET_ELEMENT_TYPE], pa.RecordBatch], str)
@beam.typehints.with_output_types(
    Dict[str, Union[np.ndarray, tf.compat.v1.SparseTensorValue]])
class _RunMetaGraphDoFn(beam.DoFn):
  """Maps a PCollection of dicts to a PCollection of dicts via a TF graph.

  The TF graph may contain more inputs than the schema provided. In that case,
  a subset of the inputs will be fed, which may cause an error if the excluded
  inputs are required to produce the included outputs.
  """

  class _GraphStateCommon(object):
    """A container for a shared graph state."""

    def __init__(self, saved_model_dir, input_tensor_keys, output_tensor_keys,
                 callable_get_outputs):
      self.saved_model_dir = saved_model_dir
      self.inputs_tensor_keys = input_tensor_keys
      self.outputs_tensor_keys = output_tensor_keys
      self.callable_get_outputs = callable_get_outputs

  # Thread-safe.
  class _GraphStateCompatV1(_GraphStateCommon):
    """A container for a shared TF1 graph state."""
<commit_msg>Remove _BatchElements which is dead code from impl.py<commit_after># Copyright 2017 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""An implementation of tf.Transform using Beam.

The beam implementation takes a user defined preprocessing function (see
../api.py for how to defined a preprocessing function) and implements it as a
Beam PTransform.

The AnalyzeDataset takes the user's preprocessing function and converts into
a TensorFlow function that can be applied to each row of a dataset.  For
example if the user's preprocessing function describes normalizing a column by
subtracting its mean, the tensorflow function will contain the mean of the
column as a constant, and will subtract this value from each value of the
column.  We refer to the result of AnalyzeDataset as a "transform function".

Since AnalyzeDataset is implemented with beam, it accepts a PCollection that
represents the dataset (see below for the exact format) and returns a singleton
PCollection containing the transform function (as a serialized TF graph).

The TransformDataset PTransform takes a dataset and a transform function, and
returns the transformed dataset where the transform function is applied to each
row of the original dataset.

There is also an AnalyzeAndTransformDataset PTransform that applies
AnalyzeDataset and TransformDataset to the same input dataset, possibly with
optimizations.

Typical usage of these functions is shown below.

def preprocessing_fn(inputs):
  ...

with beam.Pipeline(...) as p:
  with beam_impl.Context(temp_dir=my_temp_dir):
    input = p | beam_impl.read_examples(..., schema)
    transformed, transform_fn = ((input, schema)
        | beam_impl.AnalyzeAndTransformDataset(preprocessing_fn))
    transformed | beam_impl.write_examples_and_metadata(
        examples_path, metadata_path)
    transform_fn | beam_impl.write_transform_fn(transform_fn_path)

Implementation note: TensorFlow code (including our code) makes frequent use of
the default graph.  We want to avoid adding to the default graph, or including
the default graph in our own SavedModel's.  This means that wherever we call
TensorFlow code (or our code that uses the default graph) we should create a
graph and mark it as the default.  This is achieved by identifying the
entrypoints into our code where this happens and creating a
"with ... .as_default()" block.  There are four places this happens.

1) In AnalyzeDataset.expand() which is typically called from the main thread
2) In _GraphStateCommon.__init__ which is called from the worker running
   _RunMetaGraphDoFn
3) In _replace_tensors_with_constant_values, which is called in a beam.Map.
4) In extract_scalar_constants, which is called in a beam.Map.
"""
# TODO(KesterTong): Document data format.
# TODO(KesterTong): Refactor and rename now that "TransformFn" is the path to a
# SavedModel, not an in-memory object.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import copy
import datetime

import apache_beam as beam

from apache_beam.runners.portability import fn_api_runner
from apache_beam.typehints import Any
from apache_beam.typehints import Dict
from apache_beam.typehints import Iterable
from apache_beam.typehints import List
from apache_beam.typehints import Tuple
from apache_beam.typehints import Union

import numpy as np
import pyarrow as pa
import six
import tensorflow as tf
from tensorflow_transform import common
from tensorflow_transform import graph_tools
from tensorflow_transform import impl_helper
from tensorflow_transform import nodes
from tensorflow_transform import schema_inference
from tensorflow_transform.beam import analysis_graph_builder
from tensorflow_transform.beam import analyzer_cache
from tensorflow_transform.beam import beam_nodes
from tensorflow_transform.beam import common as beam_common
from tensorflow_transform.beam import context
from tensorflow_transform.beam import deep_copy
from tensorflow_transform.beam.tft_beam_io import beam_metadata_io
from tensorflow_transform.coders import example_proto_coder
from tensorflow_transform.saved import saved_transform_io
from tensorflow_transform.saved import saved_transform_io_v2
from tensorflow_transform.tf_metadata import dataset_metadata
from tensorflow_transform.tf_metadata import schema_utils
from tfx_bsl.beam import shared
from tfx_bsl.tfxio import tf_example_record
from tfx_bsl.tfxio.tensor_adapter import TensorAdapter
# TODO(https://issues.apache.org/jira/browse/SPARK-22674): Switch to
# `tfx_namedtuple.namedtuple` or `typing.NamedTuple` once the Spark issue is
# resolved.
from tfx_bsl.types import tfx_namedtuple

Context = context.Context

_CREATE_SAVED_MODEL_COUNTER_NAME = 'saved_models_created'

# For some runners, we rely on Beam to manage concurrency, i.e. we expect it to
# run one session per CPU--so we don't want to proliferate TF threads.
# Nonetheless we provide 4 threads per session for TF ops, 2 inter-
# and 2 intra-thread.  In many cases only 2 of these will be runnable
# at any given time.  This approach oversubscribes a bit to make sure
# the CPUs are really saturated.
_FIXED_PARALLELISM_TF_CONFIG = tf.compat.v1.ConfigProto(
    # TODO(b/36091595): use_per_session_threads is deprecated, but the
    # replacement session_inter_op_thread_pool is experimental; using
    # the former for now.
    use_per_session_threads=True,
    inter_op_parallelism_threads=2,
    intra_op_parallelism_threads=2)

_DEFAULT_TENSORFLOW_CONFIG_BY_BEAM_RUNNER_TYPE = {
    # TODO(katsiapis): Perhaps remove this entry once b/69922446 and b/30837990
    # are resolved.
    beam.runners.DataflowRunner: _FIXED_PARALLELISM_TF_CONFIG,

    beam.runners.DirectRunner: _FIXED_PARALLELISM_TF_CONFIG,
    fn_api_runner.FnApiRunner: _FIXED_PARALLELISM_TF_CONFIG,
}

# TODO(b/123325923): Fix the key type here to agree with the actual keys.
_DATASET_ELEMENT_TYPE = Dict[Any,  # Any -> six.text_type?
                             Union[beam_common.PRIMITIVE_TYPE,
                                   # Arbitrarily-nested lists are allowed.
                                   List[Any], np.generic, np.ndarray]]

# TODO(b/68154497): pylint: disable=no-value-for-parameter


# TODO(b/64956765): Remove this once either the keepalive issue (b/30837990), or
# the mentioned bug above is resolved.
# TODO(zoyahav): Make this a PTransform.
def _clear_shared_state_after_barrier(pipeline, input_barrier):
  """Clears any shared state from within a pipeline context.

  This will only be cleared once input_barrier becomes available.

  Args:
    pipeline: A `beam.Pipeline` object.
    input_barrier: A `PCollection` which the pipeline should wait for.

  Returns:
    An empty `PCollection`.
  """
  empty_pcoll = input_barrier | 'MakeCheapBarrier' >> beam.FlatMap(
      lambda x: None)
  return (pipeline
          | 'PrepareToClearSharedKeepAlives' >> beam.Create([None])
          | 'WaitAndClearSharedKeepAlives' >> beam.Map(
              lambda x, empty_side_input: shared.Shared().acquire(lambda: None),
              beam.pvalue.AsIter(empty_pcoll)))


# TODO(b/36223892): Verify that these type hints work and make needed fixes.
@beam.typehints.with_input_types(
    Union[List[_DATASET_ELEMENT_TYPE], pa.RecordBatch], str)
@beam.typehints.with_output_types(
    Dict[str, Union[np.ndarray, tf.compat.v1.SparseTensorValue]])
class _RunMetaGraphDoFn(beam.DoFn):
  """Maps a PCollection of dicts to a PCollection of dicts via a TF graph.

  The TF graph may contain more inputs than the schema provided. In that case,
  a subset of the inputs will be fed, which may cause an error if the excluded
  inputs are required to produce the included outputs.
  """

  class _GraphStateCommon(object):
    """A container for a shared graph state."""

    def __init__(self, saved_model_dir, input_tensor_keys, output_tensor_keys,
                 callable_get_outputs):
      self.saved_model_dir = saved_model_dir
      self.inputs_tensor_keys = input_tensor_keys
      self.outputs_tensor_keys = output_tensor_keys
      self.callable_get_outputs = callable_get_outputs

  # Thread-safe.
  class _GraphStateCompatV1(_GraphStateCommon):
    """A container for a shared TF1 graph state."""

{'fraction_non_alphanumeric': 0.05827829272353888, 'fraction_numerical': 0.010323583282455459, 'mean_word_length': 4.63238512035011, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 19, 'https://': 2, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '502810', 'n_tokens_mistral': 5297, 'n_tokens_neox': 5073, 'n_words': 1905}