Source code for apache_beam.ml.transforms.tft

#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You 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.

"""
This module defines a set of data processing transforms that can be used
to perform common data transformations on a dataset. These transforms are
implemented using the TensorFlow Transform (TFT) library. The transforms
in this module are intended to be used in conjunction with the
MLTransform class, which provides a convenient interface for
applying a sequence of data processing transforms to a dataset.

See the documentation for MLTransform for more details.

Note: The data processing transforms defined in this module don't
perform the transformation immediately. Instead, it returns a
configured operation object, which encapsulates the details of the
transformation. The actual computation takes place later in the Apache Beam
pipeline, after all transformations are set up and the pipeline is run.
"""

# pytype: skip-file

import logging
from typing import Any
from typing import Dict
from typing import Iterable
from typing import List
from typing import Optional
from typing import Tuple
from typing import Union

import apache_beam as beam
import tensorflow as tf
import tensorflow_transform as tft
from apache_beam.ml.transforms.base import BaseOperation
from tensorflow_transform import common_types

__all__ = [
    'ComputeAndApplyVocabulary',
    'ScaleToZScore',
    'ScaleTo01',
    'ScaleToGaussian',
    'ApplyBuckets',
    'ApplyBucketsWithInterpolation',
    'Bucketize',
    'TFIDF',
    'TFTOperation',
    'ScaleByMinMax',
    'NGrams',
    'BagOfWords',
    'HashStrings',
    'DeduplicateTensorPerRow',
]

# Register the expected input types for each operation
# this will be used to determine schema for the tft.AnalyzeDataset
_EXPECTED_TYPES: Dict[str, Union[int, str, float]] = {}

_LOGGER = logging.getLogger(__name__)


def register_input_dtype(type):
  def wrapper(fn):
    _EXPECTED_TYPES[fn.__name__] = type
    return fn

  return wrapper


# TODO: https://github.com/apache/beam/pull/29016
# Add support for outputting artifacts to a text file in human readable form.
[docs]class TFTOperation(BaseOperation[common_types.TensorType, common_types.TensorType]): def __init__(self, columns: List[str]) -> None: """ Base Operation class for TFT data processing transformations. Processing logic for the transformation is defined in the apply_transform() method. If you have a custom transformation that is not supported by the existing transforms, you can extend this class and implement the apply_transform() method. Args: columns: List of column names to apply the transformation. """ super().__init__(columns) if not columns: raise RuntimeError( "Columns are not specified. Please specify the column for the " " op %s" % self.__class__.__name__)
[docs] def get_ptransform_for_processing(self, **kwargs) -> beam.PTransform: from apache_beam.ml.transforms.handlers import TFTProcessHandler params = {} artifact_location = kwargs.get('artifact_location') if not artifact_location: raise RuntimeError( "artifact_location is not specified. Please specify the " "artifact_location for the op %s" % self.__class__.__name__) artifact_mode = kwargs.get('artifact_mode') if artifact_mode: params['artifact_mode'] = artifact_mode return TFTProcessHandler(artifact_location=artifact_location, **params)
@tf.function def _split_string_with_delimiter(self, data, delimiter): """ only applicable to string columns. """ data = tf.sparse.to_dense(data) # this method acts differently compared to tf.strings.split # this will split the string based on multiple delimiters while # the latter will split the string based on a single delimiter. fn = lambda data: tf.compat.v1.string_split( data, delimiter, result_type='RaggedTensor') # tf.compat.v1.string_split works on a single string. Use tf.map_fn # to apply the function on each element of the input data. data = tf.map_fn( fn, data, fn_output_signature=tf.RaggedTensorSpec( tf.TensorShape([None, None]), tf.string)) data = data.values.to_sparse() # the columns of the sparse tensor are suffixed with $indices, $values # related to sparse tensor. Create a new sparse tensor by extracting # the indices, values and dense_shape from the original sparse tensor # to preserve the original column name. data = tf.sparse.SparseTensor( indices=data.indices, values=data.values, dense_shape=data.dense_shape) # for list of string, batch dimensions becomes inverted after tf.map_fn, # transpose the data to get the original shape. if tf.shape(data)[1] == 1: data = tf.sparse.transpose(data) return data
[docs]@register_input_dtype(str) class ComputeAndApplyVocabulary(TFTOperation): def __init__( self, columns: List[str], split_string_by_delimiter: Optional[str] = None, *, default_value: Any = -1, top_k: Optional[int] = None, frequency_threshold: Optional[int] = None, num_oov_buckets: int = 0, vocab_filename: Optional[str] = None, name: Optional[str] = None): """ This function computes the vocabulary for the given columns of incoming data. The transformation converts the input values to indices of the vocabulary. Args: columns: List of column names to apply the transformation. split_string_by_delimiter: (Optional) A string that specifies the delimiter to split strings. default_value: (Optional) The value to use for out-of-vocabulary values. top_k: (Optional) The number of most frequent tokens to keep. frequency_threshold: (Optional) Limit the generated vocabulary only to elements whose absolute frequency is >= to the supplied threshold. If set to None, the full vocabulary is generated. num_oov_buckets: Any lookup of an out-of-vocabulary token will return a bucket ID based on its hash if `num_oov_buckets` is greater than zero. Otherwise it is assigned the `default_value`. vocab_filename: The file name for the vocabulary file. The vocab file will be suffixed with the column name. NOTE in order to make your pipelines resilient to implementation details please set `vocab_filename` when you are using the vocab_filename on a downstream component. """ super().__init__(columns) self._default_value = default_value self._top_k = top_k self._frequency_threshold = frequency_threshold self._num_oov_buckets = num_oov_buckets self._vocab_filename = vocab_filename self._name = name self.split_string_by_delimiter = split_string_by_delimiter
[docs] def apply_transform( self, data: common_types.TensorType, output_column_name: str) -> Dict[str, common_types.TensorType]: if self.split_string_by_delimiter: data = self._split_string_with_delimiter( data, self.split_string_by_delimiter) vocab_filename = self._vocab_filename if vocab_filename: vocab_filename = vocab_filename + f'_{output_column_name}' return { output_column_name: tft.compute_and_apply_vocabulary( x=data, default_value=self._default_value, top_k=self._top_k, frequency_threshold=self._frequency_threshold, num_oov_buckets=self._num_oov_buckets, vocab_filename=vocab_filename, name=self._name) }
[docs]@register_input_dtype(float) class ScaleToZScore(TFTOperation): def __init__( self, columns: List[str], *, elementwise: bool = False, name: Optional[str] = None): """ This function performs a scaling transformation on the specified columns of the incoming data. It processes the input data such that it's normalized to have a mean of 0 and a variance of 1. The transformation achieves this by subtracting the mean from the input data and then dividing it by the square root of the variance. Args: columns: A list of column names to apply the transformation on. elementwise: If True, the transformation is applied elementwise. Otherwise, the transformation is applied on the entire column. name: A name for the operation (optional). scale_to_z_score also outputs additional artifacts. The artifacts are mean, which is the mean value in the column, and var, which is the variance in the column. The artifacts are stored in the column named with the suffix <original_col_name>_mean and <original_col_name>_var respectively. """ super().__init__(columns) self.elementwise = elementwise self.name = name
[docs] def apply_transform( self, data: common_types.TensorType, output_column_name: str) -> Dict[str, common_types.TensorType]: output_dict = { output_column_name: tft.scale_to_z_score( x=data, elementwise=self.elementwise, name=self.name) } return output_dict
[docs]@register_input_dtype(float) class ScaleTo01(TFTOperation): def __init__( self, columns: List[str], elementwise: bool = False, name: Optional[str] = None): """ This function applies a scaling transformation on the given columns of incoming data. The transformation scales the input values to the range [0, 1] by dividing each value by the maximum value in the column. Args: columns: A list of column names to apply the transformation on. elementwise: If True, the transformation is applied elementwise. Otherwise, the transformation is applied on the entire column. name: A name for the operation (optional). ScaleTo01 also outputs additional artifacts. The artifacts are max, which is the maximum value in the column, and min, which is the minimum value in the column. The artifacts are stored in the column named with the suffix <original_col_name>_min and <original_col_name>_max respectively. """ super().__init__(columns) self.elementwise = elementwise self.name = name
[docs] def apply_transform( self, data: common_types.TensorType, output_column_name: str) -> Dict[str, common_types.TensorType]: output = tft.scale_to_0_1( x=data, elementwise=self.elementwise, name=self.name) output_dict = {output_column_name: output} return output_dict
[docs]@register_input_dtype(float) class ScaleToGaussian(TFTOperation): def __init__( self, columns: List[str], elementwise: bool = False, name: Optional[str] = None): """ This operation scales the given input column values to an approximately normal distribution with mean 0 and variance of 1. The Gaussian transformation is only applied if the column has long tails; otherwise, the transformation is the same as normalizing to z scores. For more information, see: https://www.tensorflow.org/tfx/transform/api_docs/python/tft/scale_to_gaussian Args: columns: A list of column names to apply the transformation on. elementwise: If True, the transformation is applied elementwise. Otherwise, the transformation is applied on the entire column. name: A name for the operation (optional). """ super().__init__(columns) self.elementwise = elementwise self.name = name
[docs] def apply_transform( self, data: common_types.TensorType, output_column_name: str) -> Dict[str, common_types.TensorType]: output_dict = { output_column_name: tft.scale_to_gaussian( x=data, elementwise=self.elementwise, name=self.name) } return output_dict
[docs]@register_input_dtype(float) class ApplyBuckets(TFTOperation): def __init__( self, columns: List[str], bucket_boundaries: Iterable[Union[int, float]], name: Optional[str] = None): """ This functions is used to map the element to a positive index i for which `bucket_boundaries[i-1] <= element < bucket_boundaries[i]`, if it exists. If `input < bucket_boundaries[0]`, then element is mapped to 0. If `element >= bucket_boundaries[-1]`, then element is mapped to len(bucket_boundaries). NaNs are mapped to len(bucket_boundaries). Args: columns: A list of column names to apply the transformation on. bucket_boundaries: An iterable of ints or floats representing the bucket boundaries. Must be sorted in ascending order. name: (Optional) A string that specifies the name of the operation. """ super().__init__(columns) self.bucket_boundaries = [bucket_boundaries] self.name = name
[docs] def apply_transform( self, data: common_types.TensorType, output_column_name: str) -> Dict[str, common_types.TensorType]: output = { output_column_name: tft.apply_buckets( x=data, bucket_boundaries=self.bucket_boundaries, name=self.name) } return output
[docs]@register_input_dtype(float) class ApplyBucketsWithInterpolation(TFTOperation): def __init__( self, columns: List[str], bucket_boundaries: Iterable[Union[int, float]], name: Optional[str] = None): """Interpolates values within the provided buckets and then normalizes to [0, 1]. Input values are bucketized based on the provided boundaries such that the input is mapped to a positive index i for which `bucket_boundaries[i-1] <= element < bucket_boundaries[i]`, if it exists. The values are then normalized to the range [0,1] within the bucket, with NaN values being mapped to 0.5. For more information, see: https://www.tensorflow.org/tfx/transform/api_docs/python/tft/apply_buckets_with_interpolation Args: columns: A list of column names to apply the transformation on. bucket_boundaries: An iterable of ints or floats representing the bucket boundaries sorted in ascending order. name: (Optional) A string that specifies the name of the operation. """ super().__init__(columns) self.bucket_boundaries = [bucket_boundaries] self.name = name
[docs] def apply_transform( self, data: common_types.TensorType, output_column_name: str) -> Dict[str, common_types.TensorType]: output = { output_column_name: tft.apply_buckets_with_interpolation( x=data, bucket_boundaries=self.bucket_boundaries, name=self.name) } return output
[docs]@register_input_dtype(float) class Bucketize(TFTOperation): def __init__( self, columns: List[str], num_buckets: int, *, epsilon: Optional[float] = None, elementwise: bool = False, name: Optional[str] = None): """ This function applies a bucketizing transformation on the given columns of incoming data. The transformation splits the input data range into a set of consecutive bins/buckets, and converts the input values to bucket IDs (integers) where each ID corresponds to a particular bin. Args: columns: List of column names to apply the transformation. num_buckets: Number of buckets to be created. epsilon: (Optional) A float number that specifies the error tolerance when computing quantiles, so that we guarantee that any value x will have a quantile q such that x is in the interval [q - epsilon, q + epsilon] (or the symmetric interval for even num_buckets). Must be greater than 0.0. elementwise: (Optional) A boolean that specifies whether the quantiles should be computed on an element-wise basis. If False, the quantiles are computed globally. name: (Optional) A string that specifies the name of the operation. """ super().__init__(columns) self.num_buckets = num_buckets self.epsilon = epsilon self.elementwise = elementwise self.name = name
[docs] def apply_transform( self, data: common_types.TensorType, output_column_name: str) -> Dict[str, common_types.TensorType]: output = { output_column_name: tft.bucketize( x=data, num_buckets=self.num_buckets, epsilon=self.epsilon, elementwise=self.elementwise, name=self.name) } return output
[docs]@register_input_dtype(float) class TFIDF(TFTOperation): def __init__( self, columns: List[str], vocab_size: Optional[int] = None, smooth: bool = True, name: Optional[str] = None, ): """ This function applies a tf-idf transformation on the given columns of incoming data. TFIDF outputs two artifacts for each column: the vocabulary index and the tfidf weight. The vocabulary index is a mapping from the original vocabulary to the new vocabulary. The tfidf weight is a mapping from the original vocabulary to the tfidf score. Input passed to the TFIDF is not modified and used to calculate the required artifacts. Args: columns: List of column names to apply the transformation. vocab_size: (Optional) An integer that specifies the size of the vocabulary. Defaults to None. If vocab_size is None, then the size of the vocabulary is determined by `tft.get_num_buckets_for_transformed_feature`. smooth: (Optional) A boolean that specifies whether to apply smoothing to the tf-idf score. Defaults to True. name: (Optional) A string that specifies the name of the operation. """ super().__init__(columns) self.vocab_size = vocab_size self.smooth = smooth self.name = name self.tfidf_weight = None
[docs] def apply_transform( self, data: common_types.TensorType, output_column_name: str) -> common_types.TensorType: if self.vocab_size is None: try: _LOGGER.info( 'vocab_size is not specified. Trying to infer vocab_size ' 'from the input data using ' 'tft.get_num_buckets_for_transformed_feature.') vocab_size = tft.get_num_buckets_for_transformed_feature(data) except RuntimeError: raise RuntimeError( 'vocab_size is not specified. Tried to infer vocab_size from the ' 'input data using tft.get_num_buckets_for_transformed_feature, but ' 'failed. Please specify vocab_size explicitly.') else: vocab_size = self.vocab_size vocab_index, tfidf_weight = tft.tfidf( data, vocab_size, self.smooth, self.name ) output = { output_column_name + '_vocab_index': vocab_index, output_column_name + '_tfidf_weight': tfidf_weight } return output
[docs]@register_input_dtype(float) class ScaleByMinMax(TFTOperation): def __init__( self, columns: List[str], min_value: float = 0.0, max_value: float = 1.0, name: Optional[str] = None): """ This function applies a scaling transformation on the given columns of incoming data. The transformation scales the input values to the range [min_value, max_value]. Args: columns: A list of column names to apply the transformation on. min_value: The minimum value of the output range. max_value: The maximum value of the output range. name: A name for the operation (optional). """ super().__init__(columns) self.min_value = min_value self.max_value = max_value self.name = name if self.max_value <= self.min_value: raise ValueError('max_value must be greater than min_value')
[docs] def apply_transform( self, data: common_types.TensorType, output_column_name: str) -> common_types.TensorType: output = tft.scale_by_min_max( x=data, output_min=self.min_value, output_max=self.max_value) return {output_column_name: output}
[docs]@register_input_dtype(str) class NGrams(TFTOperation): def __init__( self, columns: List[str], split_string_by_delimiter: Optional[str] = None, *, ngram_range: Tuple[int, int] = (1, 1), ngrams_separator: Optional[str] = None, name: Optional[str] = None): """ An n-gram is a contiguous sequence of n items from a given sample of text or speech. This operation applies an n-gram transformation to specified columns of incoming data, splitting the input data into a set of consecutive n-grams. Args: columns: A list of column names to apply the transformation on. split_string_by_delimiter: (Optional) A string that specifies the delimiter to split the input strings before computing ngrams. ngram_range: A tuple of integers(inclusive) specifying the range of n-gram sizes. ngrams_separator: A string that will be inserted between each ngram. name: A name for the operation (optional). """ super().__init__(columns) self.ngram_range = ngram_range self.ngrams_separator = ngrams_separator self.name = name self.split_string_by_delimiter = split_string_by_delimiter if ngram_range != (1, 1) and not ngrams_separator: raise ValueError( 'ngrams_separator must be specified when ngram_range is not (1, 1)')
[docs] def apply_transform( self, data: common_types.TensorType, output_column_name: str) -> Dict[str, common_types.TensorType]: if self.split_string_by_delimiter: data = self._split_string_with_delimiter( data, self.split_string_by_delimiter) output = tft.ngrams(data, self.ngram_range, self.ngrams_separator) return {output_column_name: output}
[docs]@register_input_dtype(str) class BagOfWords(TFTOperation): def __init__( self, columns: List[str], split_string_by_delimiter: Optional[str] = None, *, ngram_range: Tuple[int, int] = (1, 1), ngrams_separator: Optional[str] = None, compute_word_count: bool = False, key_vocab_filename: Optional[str] = None, name: Optional[str] = None, ): """ Bag of words contains the unique words present in the input text. This operation applies a bag of words transformation to specified columns of incoming data. Also, the transformation accepts a Tuple of integers specifying the range of n-gram sizes. The transformation splits the input data into a set of consecutive n-grams if ngram_range is specified. The n-grams are then converted to a bag of words. Also, you can specify a seperator string that will be inserted between each ngram. Args: columns: A list of column names to apply the transformation on. split_string_by_delimiter: (Optional) A string that specifies the delimiter to split the input strings before computing ngrams. ngram_range: A tuple of integers(inclusive) specifying the range of n-gram sizes. seperator: A string that will be inserted between each ngram. compute_word_count: A boolean that specifies whether to compute the unique word count over the entire dataset. Defaults to False. key_vocab_filename: The file name for the key vocabulary file when compute_word_count is True. If empty, a file name will be chosen based on the current scope. If provided, the vocab file will be suffixed with the column name. name: A name for the operation (optional). Note that original order of the input may not be preserved. """ self.columns = columns self.ngram_range = ngram_range self.ngrams_separator = ngrams_separator self.name = name self.split_string_by_delimiter = split_string_by_delimiter self.key_vocab_filename = key_vocab_filename if compute_word_count: self.compute_word_count_fn = count_unique_words else: self.compute_word_count_fn = lambda *args, **kwargs: None if ngram_range != (1, 1) and not ngrams_separator: raise ValueError( 'ngrams_separator must be specified when ngram_range is not (1, 1)')
[docs] def apply_transform(self, data: tf.SparseTensor, output_col_name: str): if self.split_string_by_delimiter: data = self._split_string_with_delimiter( data, self.split_string_by_delimiter) output = tft.bag_of_words( data, self.ngram_range, self.ngrams_separator, self.name) # word counts are written to the file only if compute_word_count is True key_vocab_filename = self.key_vocab_filename if key_vocab_filename: key_vocab_filename = key_vocab_filename + f'_{output_col_name}' self.compute_word_count_fn(data, key_vocab_filename) return {output_col_name: output}
def count_unique_words( data: tf.SparseTensor, output_vocab_name: Optional[str]) -> None: tft.count_per_key(data, key_vocabulary_filename=output_vocab_name)
[docs]@register_input_dtype(str) class HashStrings(TFTOperation): def __init__( self, columns: List[str], hash_buckets: int, key: Optional[Tuple[int, int]] = None, name: Optional[str] = None): '''Hashes strings into the provided number of buckets. Args: columns: A list of the column names to apply the transformation on. hash_buckets: the number of buckets to hash the strings into. key: optional. An array of two Python `uint64`. If passed, output will be a deterministic function of `strings` and `key`. Note that hashing will be slower if this value is specified. name: optional. A name for this operation. Raises: ValueError if `hash_buckets` is not a positive and non-zero integer. ''' self.hash_buckets = hash_buckets self.key = key self.name = name if hash_buckets < 1: raise ValueError( 'number of hash buckets must be positive, got ', hash_buckets) super().__init__(columns)
[docs] def apply_transform( self, data: common_types.TensorType, output_col_name: str) -> Dict[str, common_types.TensorType]: output_dict = { output_col_name: tft.hash_strings( strings=data, hash_buckets=self.hash_buckets, key=self.key, name=self.name) } return output_dict
[docs]@register_input_dtype(str) class DeduplicateTensorPerRow(TFTOperation): def __init__(self, columns: List[str], name: Optional[str] = None): """ Deduplicates each row (0th dimension) of the provided tensor. Args: columns: A list of the columns to apply the transformation on. name: optional. A name for this operation. """ self.name = name super().__init__(columns)
[docs] def apply_transform( self, data: common_types.TensorType, output_col_name: str) -> Dict[str, common_types.TensorType]: output_dict = { output_col_name: tft.deduplicate_tensor_per_row( input_tensor=data, name=self.name) } return output_dict