#
# 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',
'ApplyBuckets',
'Bucketize',
'TFIDF',
'TFTOperation',
'ScaleByMinMax',
'NGrams',
'BagOfWords',
]
# 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__)
@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]@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]@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
@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
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
@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
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]@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 vocabu index and
the tfidf weight. The vocabu 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]@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]@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]@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)')
def count_unique_words(
data: tf.SparseTensor, output_vocab_name: Optional[str]) -> None:
tft.count_per_key(data, key_vocabulary_filename=output_vocab_name)
@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)
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
@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)
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