diff --git a/README.md b/README.md index c1fa7879..43e16b93 100644 --- a/README.md +++ b/README.md @@ -116,6 +116,7 @@ torch_transformers = kamae.get_compatible_transformers('torch') | DateParse | Parses a string date of format YYYY-MM-DD to extract a given date part. E.g. day of year. | [Link](src/kamae/keras/tensorflow/layers/date_parse.py) | TensorFlow-only | [Link](src/kamae/spark/transformers/date_parse.py) | | DateTimeToUnixTimestamp | Converts a UTC datetime string to unix timestamp. | [Link](src/kamae/keras/tensorflow/layers/date_time_to_unix_timestamp.py) | TensorFlow-only | [Link](src/kamae/spark/transformers/date_time_to_unix_timestamp.py) | | Divide | Divides a single feature by a constant or divides multiple features against each other. | [Link](src/kamae/keras/core/layers/divide.py) | Multi-backend | [Link](src/kamae/spark/transformers/divide.py) | +| EventNgramLookup | Tokenizes sequences of fixed-size discrete ID tuples into token ids via a fitted within-event n-gram vocabulary. | [Link](src/kamae/keras/tensorflow/layers/event_ngram_lookup.py) | TensorFlow-only | [Link](src/kamae/spark/estimators/event_ngram_lookup.py) | | Exp | Applies the exp(x) operation to the feature. | [Link](src/kamae/keras/core/layers/exp.py) | Multi-backend | [Link](src/kamae/spark/transformers/exp.py) | | Exponent | Applies the x^exponent to a single feature or x^y for multiple features. | [Link](src/kamae/keras/core/layers/exponent.py) | Multi-backend | [Link](src/kamae/spark/transformers/exponent.py) | | HashIndex | Transforms strings to indices via a hash table of predeterminded size. | [Link](src/kamae/keras/tensorflow/layers/hash_index.py) | TensorFlow-only | [Link](src/kamae/spark/transformers/hash_index.py) | diff --git a/src/kamae/keras/tensorflow/layers/__init__.py b/src/kamae/keras/tensorflow/layers/__init__.py index 81e81324..f5e51b50 100644 --- a/src/kamae/keras/tensorflow/layers/__init__.py +++ b/src/kamae/keras/tensorflow/layers/__init__.py @@ -28,6 +28,7 @@ from .date_diff import DateDiffLayer # noqa: F401 from .date_parse import DateParseLayer # noqa: F401 from .date_time_to_unix_timestamp import DateTimeToUnixTimestampLayer # noqa: F401 +from .event_ngram_lookup import EventNgramLookupLayer # noqa: F401 from .hash_index import HashIndexLayer # noqa: F401 from .if_statement import IfStatementLayer # noqa: F401 from .lambda_function import LambdaFunctionLayer # noqa: F401 @@ -66,6 +67,7 @@ "DateDiffLayer", "DateParseLayer", "DateTimeToUnixTimestampLayer", + "EventNgramLookupLayer", "HashIndexLayer", "IfStatementLayer", "LambdaFunctionLayer", diff --git a/src/kamae/keras/tensorflow/layers/event_ngram_lookup.py b/src/kamae/keras/tensorflow/layers/event_ngram_lookup.py new file mode 100644 index 00000000..b604b779 --- /dev/null +++ b/src/kamae/keras/tensorflow/layers/event_ngram_lookup.py @@ -0,0 +1,359 @@ +# Copyright [2024] Expedia, Inc. +# +# 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. + +""" +TensorFlow layer that tokenizes event-level discrete ID values via a lookup table. + +Companion layer to the ``EventNgramLookupEstimator`` / ``EventNgramLookupTransformer``. +Each input row holds a sequence of events, and each event is a fixed-size tuple of +discrete ID values (e.g. 4 levels ``L0, L1, L2, L3``). The layer splits each input into +per-event tuples and maps every tuple to its pre-computed top-k token list, so that +inference is a pure ``O(1)`` table lookup identical to the Spark ``_transform`` path. + +A single layer instance tokenizes ALL input columns (one per ``num_events_per_input`` +entry), so the fitted lookup table is embedded once regardless of the number of inputs. +It returns one token tensor per input, in input order. When a ``token_type_lookup`` is +supplied, it additionally returns one "type" tensor per input (each token id mapped to a +bitmask of the ID levels its n-gram spans), grouped after all the token tensors: +``[tokens_0, ..., tokens_{N-1}, types_0, ..., types_{N-1}]``. + +Each event tuple is bijectively packed into a single ``int64`` key and resolved by a +Keras ``IntegerLookup`` sublayer plus a gathered values tensor, fully vectorized and +with no string ops in the serving path. The table is persisted in ``get_config`` as +parallel ``lookup_keys`` / ``lookup_values`` lists, so the layer reloads standalone. + +Token id conventions (shared with the vocabulary): + - ``0`` = padding (````): emitted for all-zero / missing events, and for the + unused token slots of an event that matched fewer than ``top_k`` n-grams. + - ``1`` = unknown (````): a non-zero tuple absent from the table yields a + single ```` followed by padding. +""" + +from typing import Any, Dict, Iterable, List, Optional, Union + +import keras +import tensorflow as tf +from keras import KerasTensor + +import kamae +from kamae.keras.core.backend import TENSORFLOW_ONLY +from kamae.keras.core.base import BaseLayer + +PAD_TOKEN_ID = 0 +UNK_TOKEN_ID = 1 + + +def compute_key_bits(ids: Iterable[int], tuple_size: int) -> int: + """ + Computes the bits allotted per ID level when packing an event tuple into an int64. + + Each tuple key is bijectively packed into a single signed ``int64`` by allotting + ``key_bits`` bits per ID level, where ``key_bits`` is the bit length of the largest + ID. This is the single definition of which IDs the lookup table can hold: the + estimator calls it at fit time, so an unpackable table fails there rather than + later, when the Keras layer is built. + + :param ids: Every ID in the lookup table keys. + :param tuple_size: Number of ID levels per event tuple. + :raises ValueError: If an ID is negative, or the packed key would not fit in a + signed ``int64``. + :returns: Number of bits allotted per ID level. + """ + ids = [int(x) for x in ids] + min_id = min(ids, default=0) + if min_id < 0: + raise ValueError( + f"Discrete ID values must be non-negative, but the lookup table " + f"contains {min_id}. Each event tuple is packed into a single " + f"non-negative int64 key, which a negative ID cannot represent." + ) + max_id = max(ids, default=1) + key_bits = max(1, max_id.bit_length()) + if tuple_size * key_bits > 63: + raise ValueError( + f"Cannot pack event tuples into a signed int64: tuple_size " + f"({tuple_size}) * key_bits ({key_bits}, from a largest ID " + f"of {max_id}) = {tuple_size * key_bits} bits, which exceeds " + f"63. Reduce tuple_size or the ID cardinality." + ) + return key_bits + + +@tf.keras.utils.register_keras_serializable(package=kamae.__name__) +class EventNgramLookupLayer(BaseLayer): + """ + Tokenizes discrete ID values using a pre-computed tuple to top-k token lookup table. + + For each event position the output is the tuple's ``top_k`` token ids. Unknown + (non-zero) tuples map to ```` and all-zero / padding events map to ````, + matching the Spark transformer exactly. + + One layer tokenizes all input columns (``num_events_per_input`` gives the event + count of each). Each input may be rank-2 ``(batch, num_events * tuple_size)`` or + rank-3 ``(batch, list_size, num_events * tuple_size)``; the list dimension is + preserved. The id axis is padded with all-zero (padding) events or truncated to + ``num_events * tuple_size`` before being split into tuples. + + As in the Spark transformer, the id axis must hold a whole number of events, and a + width that is not a multiple of ``tuple_size`` raises. A known width is checked + when the layer is called; an unknown width is checked in the graph, per batch. + + Outputs one token tensor per input. With ``token_type_lookup`` set it also outputs + one type tensor per input (same shape as the tokens), grouped after the tokens. + """ + + supported_backends = TENSORFLOW_ONLY + jit_compatible = False + + def __init__( + self, + num_events_per_input: List[int], + top_k: int, + tuple_size: int, + lookup_keys: Optional[List[List[int]]] = None, + lookup_values: Optional[List[List[int]]] = None, + token_type_lookup: Optional[List[int]] = None, + name: Optional[str] = None, + input_dtype: Optional[str] = None, + output_dtype: Optional[str] = None, + **kwargs: Any, + ) -> None: + """ + Initializes the EventNgramLookupLayer layer. + + The lookup table is supplied as the parallel ``lookup_keys`` / ``lookup_values`` + lists, which is both what the transformer passes and what ``get_config`` + persists, so construction and reload take the same path. + + :param num_events_per_input: Number of events for each input column, in input + order (e.g. ``[10, 10, 1]``). Its length is the number of inputs the layer + tokenizes. + :param top_k: Number of tokens emitted per event tuple. + :param tuple_size: Number of ID levels per event tuple. + :param lookup_keys: Lookup table keys: one ``tuple_size``-long int list per + event tuple. Defaults to `None` (an empty table, so every tuple is unknown). + :param lookup_values: Lookup table values: one ``top_k``-long token list per + key, positionally aligned with ``lookup_keys``. Defaults to `None`. + :param token_type_lookup: Optional per-token-id list mapping each token to a + bitmask of the ID levels its n-gram spans. When given, the layer also emits + a type tensor per input. Defaults to `None` (tokens only). + :param name: The name of the layer. Defaults to `None`. + :param input_dtype: The dtype to cast the input to. Defaults to `None`. + :param output_dtype: The dtype to cast the output to. Defaults to `None`. + """ + super().__init__( + name=name, input_dtype=input_dtype, output_dtype=output_dtype, **kwargs + ) + self.num_events_per_input = [int(n) for n in num_events_per_input] + self.top_k = int(top_k) + self.tuple_size = int(tuple_size) + + self._lookup_keys = lookup_keys if lookup_keys is not None else [] + self._lookup_values = lookup_values if lookup_values is not None else [] + self._build_lookup_table(self._lookup_keys, self._lookup_values) + + # Optional per-token type lookup (bitmask of ID levels; 0 = pad/unk). + self._token_type_lookup = token_type_lookup + if token_type_lookup is not None: + self._type_tensor = tf.constant(token_type_lookup, dtype=tf.int32) + + def _build_lookup_table( + self, keys: List[List[int]], values: List[List[int]] + ) -> None: + """ + Builds the ``IntegerLookup`` sublayer and gathered values tensor. + + Each tuple key is packed into a single ``int64`` (see ``compute_key_bits``). + Packing lets one ``IntegerLookup`` index the whole table, which keeps string + ops out of the serving path. The lookup maps a miss/OOV to index 0 and the i-th + key to index ``i + 1``; ``_call`` shifts back by one to gather from the values + tensor. + + :param keys: Tuple keys as a list of int lists. + :param values: Token lists (one per key), each of length ``top_k``. + :raises ValueError: If an ID is negative, or the packed key would not fit in a + signed ``int64``. + """ + self.key_bits = compute_key_bits((x for k in keys for x in k), self.tuple_size) + self.key_powers = [1 << (j * self.key_bits) for j in range(self.tuple_size)] + packed_keys = [ + sum(int(x) * p for x, p in zip(k, self.key_powers)) for k in keys + ] + + # Keep at least one value row so the gather in _call stays valid even for a + # (degenerate) empty table. + self.values_tensor = ( + tf.constant(values, dtype=tf.int32) + if values + else tf.zeros([1, self.top_k], dtype=tf.int32) + ) + # An empty table uses one dummy key that real packed keys, always + # non-negative, never match, so every lookup resolves to OOV (index 0). It must + # be -2: IntegerLookup strips a leading run of its own special tokens, and -1 is + # its default oov_token, so a vocabulary of just [-1] is left empty. + self.key_to_index = keras.layers.IntegerLookup( + vocabulary=packed_keys if packed_keys else [-2], + num_oov_indices=1, + mask_token=None, + name=f"{self.name}_key_lookup", + ) + self.unk_pattern = tf.constant( + [UNK_TOKEN_ID] + [PAD_TOKEN_ID] * (self.top_k - 1), dtype=tf.int32 + ) + self.pad_pattern = tf.constant([PAD_TOKEN_ID] * self.top_k, dtype=tf.int32) + + @property + def compatible_dtypes(self) -> Optional[List[str]]: + """ + Returns the compatible dtypes of the layer. + + :returns: The compatible dtypes of the layer. + """ + return ["int32", "int64"] + + def _tokenize(self, inputs: KerasTensor, num_events: int) -> KerasTensor: + """ + Tokenizes one input tensor into its flat token ids. + + :param inputs: ID tensor, rank-2 ``(batch, num_events * tuple_size)`` or rank-3 + ``(batch, list_size, num_events * tuple_size)``. + :param num_events: Number of events this input is padded/truncated to. + :raises ValueError: If the input is neither rank-2 nor rank-3, or its known id + axis width is not a whole number of events. + :raises tf.errors.InvalidArgumentError: If an id axis of unknown width is not a + whole number of events. + :returns: Token id tensor of shape ``(batch, num_events * top_k)`` for a + rank-2 input, or ``(batch, list_size, num_events * top_k)`` for rank-3. + """ + output_length = num_events * self.top_k + input_rank = len(inputs.shape) + if input_rank not in (2, 3): + raise ValueError( + f"Expected rank-2 (batch, ids) or rank-3 (batch, list, ids) inputs, " + f"but got an input of rank {input_rank}." + ) + width = inputs.shape[-1] + if width is None: + tf.debugging.assert_equal( + tf.shape(inputs)[-1] % self.tuple_size, + 0, + message=( + f"The id axis must be a whole number of events: its width is not " + f"a multiple of tuple_size {self.tuple_size}." + ), + ) + elif width % self.tuple_size != 0: + raise ValueError( + f"The id axis must be a whole number of events: got width {width}, " + f"which is not a multiple of tuple_size {self.tuple_size}." + ) + if input_rank == 3: + batch_size = tf.shape(inputs)[0] + list_size = tf.shape(inputs)[1] + inputs_flat = tf.reshape(inputs, [-1, tf.shape(inputs)[2]]) + restore_list = True + else: + inputs_flat = inputs + restore_list = False + + # Pad/truncate the id axis to num_events * tuple_size, then split into tuples. + # Padding by the full length first covers short and long inputs alike. + expected_length = num_events * self.tuple_size + inputs_padded = tf.pad(inputs_flat, [[0, 0], [0, expected_length]])[ + :, :expected_length + ] + all_tuples = tf.reshape(inputs_padded, [-1, self.tuple_size]) + + # All-zero tuples are padding; everything else is looked up (miss -> UNK). + is_padding = tf.reduce_all(tf.equal(all_tuples, 0), axis=1) + + # Pack each tuple into one int64 key. IDs outside [0, max_id] cannot be + # represented, so they are clamped for packing and then forced to miss -- + # clamping alone would let an out-of-range ID alias onto a valid key. + max_id = (1 << self.key_bits) - 1 + ids_64 = tf.cast(all_tuples, tf.int64) + in_range = tf.reduce_all((ids_64 >= 0) & (ids_64 <= max_id), axis=1) + powers = tf.constant(self.key_powers, dtype=tf.int64) + tuple_keys = tf.reduce_sum(tf.clip_by_value(ids_64, 0, max_id) * powers, axis=1) + + # Shift back by one (IntegerLookup uses index 0 for OOV) to index values_tensor. + indices = self.key_to_index(tuple_keys) + found = (indices > 0) & in_range + gathered = tf.gather(self.values_tensor, tf.maximum(indices - 1, 0)) + + # tf.where broadcasts the (top_k,) unk/pad constants across all tuples. + tokens = tf.where(found[:, None], gathered, self.unk_pattern) + tokens = tf.where(is_padding[:, None], self.pad_pattern, tokens) + + if restore_list: + return tf.reshape(tokens, [batch_size, list_size, output_length]) + return tf.reshape(tokens, [-1, output_length]) + + def _call( + self, inputs: Union[KerasTensor, List[KerasTensor]], **kwargs: Any + ) -> Union[KerasTensor, List[KerasTensor]]: + """ + Tokenizes every input, and optionally derives their per-token types. + + :param inputs: List of ID tensors (one per ``num_events_per_input`` entry). A + single input may be passed as a bare tensor. + :raises ValueError: If the number of inputs does not match + ``num_events_per_input``. + :returns: One token tensor per input. If a ``token_type_lookup`` was supplied, + the per-input type tensors are appended after the token tensors. A single token + tensor (one input, no types) is returned unwrapped. + """ + if not isinstance(inputs, (list, tuple)): + inputs = [inputs] + if len(inputs) != len(self.num_events_per_input): + raise ValueError( + f"Expected one input per num_events_per_input entry " + f"({len(self.num_events_per_input)}), but got {len(inputs)} inputs." + ) + + tokens = [ + self._tokenize(inp, num_events) + for inp, num_events in zip(inputs, self.num_events_per_input) + ] + + if self._token_type_lookup is None: + return tokens[0] if len(tokens) == 1 else tokens + + # Type = per-token bitmask of contributing ID levels, gathered by token id. + types = [tf.gather(self._type_tensor, tok) for tok in tokens] + return tokens + types + + def get_config(self) -> Dict[str, Any]: + """ + Gets the configuration of the EventNgramLookupLayer layer. + Used for saving and loading the layer from a model. + + Persists the lookup table as parallel ``lookup_keys`` / ``lookup_values`` lists + (and the optional ``token_type_lookup``) so the layer is reconstructed by + ``__init__`` without a custom ``from_config``. + + :returns: Dictionary of the configuration of the layer. + """ + config = super().get_config() + config.update( + { + "num_events_per_input": self.num_events_per_input, + "top_k": self.top_k, + "tuple_size": self.tuple_size, + "lookup_keys": self._lookup_keys, + "lookup_values": self._lookup_values, + "token_type_lookup": self._token_type_lookup, + } + ) + return config diff --git a/src/kamae/spark/estimators/__init__.py b/src/kamae/spark/estimators/__init__.py index 80589091..bdfcfe92 100644 --- a/src/kamae/spark/estimators/__init__.py +++ b/src/kamae/spark/estimators/__init__.py @@ -14,6 +14,7 @@ from .base import BaseEstimator # noqa: F401 from .conditional_standard_scale import ConditionalStandardScaleEstimator # noqa: F401 +from .event_ngram_lookup import EventNgramLookupEstimator # noqa: F401 from .impute import ImputeEstimator # noqa: F401 from .min_max_scale import MinMaxScaleEstimator # noqa: F401 from .one_hot_encode import OneHotEncodeEstimator # noqa: F401 diff --git a/src/kamae/spark/estimators/event_ngram_lookup.py b/src/kamae/spark/estimators/event_ngram_lookup.py new file mode 100644 index 00000000..c0908ddb --- /dev/null +++ b/src/kamae/spark/estimators/event_ngram_lookup.py @@ -0,0 +1,299 @@ +# Copyright [2024] Expedia, Inc. +# +# 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. + +""" +EventNgramLookupEstimator: learns an event-level n-gram vocabulary and lookup table. + +Discrete ID values arrive as sequences of fixed-size events (e.g. 4 ID levels +``L0, L1, L2, L3``). The estimator counts every within-event n-gram (all +combinations of length 1..``tupleSize``, so skipped ID levels are included) across +the corpus, keeps the ``vocabSize`` most frequent above ``minNgramFreq``, and +pre-computes, for every observed event tuple, its ``topK`` token ids. The fitted +lookup table is handed to the ``EventNgramLookupTransformer`` for O(1) tokenization. +Tuples not observed during fitting map to ```` at transform time. +""" + +# pylint: disable=unused-argument +# pylint: disable=invalid-name +# pylint: disable=too-many-ancestors +# pylint: disable=no-member +from typing import List, Optional, Tuple + +from pyspark import keyword_only +from pyspark.ml.param import Param, Params, TypeConverters +from pyspark.sql import DataFrame +from pyspark.sql.types import DataType, IntegerType, LongType + +from kamae.keras.core.backend import TENSORFLOW_ONLY +from kamae.keras.tensorflow.layers.event_ngram_lookup import compute_key_bits +from kamae.spark.params import ( + EventNgramLookupParams, + MultiInputMultiOutputParams, + SampleFractionParams, +) +from kamae.spark.transformers import EventNgramLookupTransformer +from kamae.spark.utils import ( + EventNgramVocabulary, + build_tuple_lookup_table, + build_vocabulary, + collect_ngrams_from_dataframe, + validate_event_column_lengths, +) + +from .base import BaseEstimator + + +class EventNgramLookupEstimatorParams(Params): + """ + Mixin class containing the fit-time parameters of the EventNgramLookupEstimator. + + These govern vocabulary selection only, so they are not carried by the fitted + transformer and live here rather than in the shared params module. + """ + + vocabSize = Param( + Params._dummy(), + "vocabSize", + "Target vocabulary size when training, including the reserved pad and unk " + "tokens.", + typeConverter=TypeConverters.toInt, + ) + + minNgramFreq = Param( + Params._dummy(), + "minNgramFreq", + "Minimum corpus frequency for an n-gram to enter the vocabulary.", + typeConverter=TypeConverters.toInt, + ) + + def setVocabSize(self, value: int) -> "EventNgramLookupEstimatorParams": + """ + Sets the vocabSize parameter. + + :param value: Target vocabulary size when training. + :raises ValueError: If the vocabulary size does not leave room for at least + one learned n-gram alongside the two reserved tokens. + :returns: Instance of class mixed in. + """ + if value < 3: + raise ValueError( + f"vocabSize must be at least 3 to hold the reserved pad and unk " + f"tokens plus one n-gram. Got {value}" + ) + return self._set(vocabSize=value) + + def getVocabSize(self) -> int: + """ + Gets the vocabSize parameter. + + :returns: Target vocabulary size when training. + """ + return self.getOrDefault(self.vocabSize) + + def setMinNgramFreq(self, value: int) -> "EventNgramLookupEstimatorParams": + """ + Sets the minNgramFreq parameter. + + :param value: Minimum n-gram frequency threshold. + :raises ValueError: If the threshold is not a positive integer. + :returns: Instance of class mixed in. + """ + if value < 1: + raise ValueError(f"minNgramFreq must be a positive integer. Got {value}") + return self._set(minNgramFreq=value) + + def getMinNgramFreq(self) -> int: + """ + Gets the minNgramFreq parameter. + + :returns: Minimum n-gram frequency threshold. + """ + return self.getOrDefault(self.minNgramFreq) + + +class EventNgramLookupEstimator( + BaseEstimator, + MultiInputMultiOutputParams, + EventNgramLookupParams, + EventNgramLookupEstimatorParams, + SampleFractionParams, +): + """ + Estimator that learns an event-level n-gram vocabulary and its lookup table. + + Fitting counts within-event n-grams across all input columns, selects the most + frequent ones into a shared vocabulary, and pre-computes each observed event + tuple's top-k token ids. It returns an ``EventNgramLookupTransformer`` carrying + that lookup table. Set ``sampleFraction`` to train on a random subset of rows. + """ + + supported_backends = TENSORFLOW_ONLY + jit_compatible = False + + @keyword_only + def __init__( + self, + inputCols: Optional[List[str]] = None, + outputCols: Optional[List[str]] = None, + inputDtype: Optional[str] = None, + outputDtype: Optional[str] = None, + layerName: Optional[str] = None, + numEventsPerInput: Optional[List[int]] = None, + tupleSize: int = 4, + topK: int = 10, + vocabSize: int = 50000, + minNgramFreq: int = 10, + sampleFraction: Optional[float] = None, + useFitSample: bool = False, + includeTokenTypes: bool = False, + ) -> None: + """ + Initializes the EventNgramLookupEstimator. + + :param inputCols: Input column names holding discrete ID values. + :param outputCols: Output column names for the token arrays. + :param inputDtype: Data type to cast the input columns to before fitting. + :param outputDtype: Data type the transformer casts its output columns to. + :param layerName: Name of the layer. Used as the name of the Keras layer in the + keras model. If not set, we use the uid of the Spark estimator. + :param numEventsPerInput: Number of events per input column, e.g. [10, 10, 1]. + :param tupleSize: Number of discrete ID values (ID levels) per event tuple. + :param topK: Number of tokens emitted per event tuple. + :param vocabSize: Target vocabulary size (including the reserved pad and unk + tokens). + :param minNgramFreq: Minimum corpus frequency for an n-gram to be kept. + :param sampleFraction: Optional fraction of rows to sample before fitting + (exclusive 0.0-1.0). None (default) uses all rows. Note that sampling biases + n-gram frequencies against the minNgramFreq threshold, so the learned + vocabulary is approximate. + :param useFitSample: If True, fit on the enclosing pipeline's shared sample + when fitSampleFraction is set. Default False. Leaving this False is + recommended: this estimator is a vocabulary builder thresholded on exact + corpus frequencies, not a sample-robust statistic like a mean or quantile. + :param includeTokenTypes: If `True`, the transformer also emits a parallel + ``_types`` column per input giving each token's ID-level bitmask. Defaults + to `False`. + :returns: None - class instantiated. + """ + super().__init__() + self._setDefault( + numEventsPerInput=None, + tupleSize=4, + topK=10, + vocabSize=50000, + minNgramFreq=10, + sampleFraction=None, + useFitSample=False, + includeTokenTypes=False, + ) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @property + def compatible_dtypes(self) -> Optional[List[DataType]]: + """ + List of compatible data types for the estimator input columns. + If the computation can be performed on any data type, return None. + + :returns: List of compatible data types for the estimator. + """ + return [IntegerType(), LongType()] + + def get_layer_inputs_outputs(self) -> Tuple[List[str], List[str]]: + """ + Gets the input and output column names, including the token-type columns. + + Overrides the base method because, with ``includeTokenTypes`` set, the fitted + transformer also produces a ``_types`` column per output column. Declaring + them here keeps the pipeline graph accurate at fit time, so that a downstream + estimator reading a type column is fitted after this stage has been applied. + + :returns: Tuple of the input column names and the output column names. + """ + inputs, token_cols = super().get_layer_inputs_outputs() + return inputs, token_cols + self.getTokenTypeCols(token_cols) + + def _fit(self, dataset: DataFrame) -> EventNgramLookupTransformer: + """ + Trains the n-gram vocabulary and returns the fitted transformer. + + :param dataset: Input DataFrame with the discrete-ID columns. + :returns: An ``EventNgramLookupTransformer`` carrying the fitted lookup table. + :raises ValueError: If ``inputCols``, ``outputCols`` and ``numEventsPerInput`` + do not all have the same length, or the observed IDs cannot be packed by the + Keras layer (negative, or too large for ``tupleSize``). A row whose ids are not + a whole number of events also fails the fit, raised on the Spark executors. + """ + input_cols = self.getInputCols() + output_cols = self.getOutputCols() + num_events_per_input = self.getNumEventsPerInput() + validate_event_column_lengths(input_cols, output_cols, num_events_per_input) + + tuple_size = self.getTupleSize() + top_k = self.getTopK() + + # Count within-event n-grams, then keep the most frequent as the vocabulary. + ngram_counter = collect_ngrams_from_dataframe( + df=dataset, + input_columns=input_cols, + event_size=tuple_size, + min_ngram_freq=self.getMinNgramFreq(), + ) + vocabulary = EventNgramVocabulary( + ngrams=build_vocabulary( + ngram_counter=ngram_counter, vocab_size=self.getVocabSize() + ) + ) + + # Pre-compute the top-k tokens for every event tuple observed in the corpus. + tuple_to_tokens = build_tuple_lookup_table( + df=dataset, + input_columns=input_cols, + vocabulary=vocabulary, + top_k=top_k, + event_size=tuple_size, + ) + + # Optionally derive the per-token ID-level bitmask ("type") lookup. + include_token_types = self.getIncludeTokenTypes() + token_type_lookup = ( + vocabulary.build_type_lookup() if include_token_types else None + ) + + # Flatten the table into two parallel int lists so the fitted pipeline can be + # saved (Spark writes params to JSON). + lookup_keys: List[int] = [] + lookup_values: List[int] = [] + for id_tuple, tokens in tuple_to_tokens.items(): + lookup_keys.extend(id_tuple) + lookup_values.extend(tokens) + + # Fail at fit time if the Keras layer could not pack the tuples into int64 keys. + compute_key_bits(lookup_keys, tuple_size) + + return EventNgramLookupTransformer( + inputCols=input_cols, + outputCols=output_cols, + inputDtype=self.getInputDtype(), + outputDtype=self.getOutputDtype(), + layerName=self.getLayerName(), + numEventsPerInput=num_events_per_input, + tupleSize=tuple_size, + topK=top_k, + vocabularySize=vocabulary.vocab_size, + lookupKeys=lookup_keys, + lookupValues=lookup_values, + includeTokenTypes=include_token_types, + tokenTypeLookup=token_type_lookup, + ) diff --git a/src/kamae/spark/params/__init__.py b/src/kamae/spark/params/__init__.py index cc301c99..7233d4a2 100644 --- a/src/kamae/spark/params/__init__.py +++ b/src/kamae/spark/params/__init__.py @@ -29,6 +29,7 @@ DateTimeParams, DefaultIntValueParams, DropUnseenParams, + EventNgramLookupParams, HashIndexParams, ImputeMethodParams, LabelsArrayParams, diff --git a/src/kamae/spark/params/shared.py b/src/kamae/spark/params/shared.py index 36a60dfa..6fea8e1d 100644 --- a/src/kamae/spark/params/shared.py +++ b/src/kamae/spark/params/shared.py @@ -1038,3 +1038,136 @@ def getMaskValue(self) -> str: :returns: Str value of the mask value. """ return self.getOrDefault(self.maskValue) + + +class EventNgramLookupParams(Params): + """ + Mixin class for the event n-gram lookup tokenizer parameters shared between the + ``EventNgramLookupEstimator`` and the ``EventNgramLookupTransformer``. + """ + + numEventsPerInput = Param( + Params._dummy(), + "numEventsPerInput", + "Number of events per input column, e.g. [10, 10, 1].", + typeConverter=TypeConverters.toListInt, + ) + + tupleSize = Param( + Params._dummy(), + "tupleSize", + "Number of discrete ID values (ID levels) per event tuple.", + typeConverter=TypeConverters.toInt, + ) + + topK = Param( + Params._dummy(), + "topK", + "Number of tokens emitted per event tuple.", + typeConverter=TypeConverters.toInt, + ) + + includeTokenTypes = Param( + Params._dummy(), + "includeTokenTypes", + "Whether to also emit, per input column, a parallel '_types' column " + "giving each token's ID-level bitmask.", + typeConverter=TypeConverters.toBoolean, + ) + + def setNumEventsPerInput(self, value: List[int]) -> "EventNgramLookupParams": + """ + Sets the numEventsPerInput parameter. + + :param value: Number of events per input column. + :raises ValueError: If the list is empty or any column has fewer than one event. + :returns: Instance of class mixed in. + """ + if not value or any(num_events < 1 for num_events in value): + raise ValueError( + f"numEventsPerInput must give at least one event for each input " + f"column. Got {value}" + ) + return self._set(numEventsPerInput=value) + + def getNumEventsPerInput(self) -> List[int]: + """ + Gets the numEventsPerInput parameter. + + :returns: Number of events per input column. + """ + return self.getOrDefault(self.numEventsPerInput) + + def setTupleSize(self, value: int) -> "EventNgramLookupParams": + """ + Sets the tupleSize parameter. + + :param value: Number of discrete ID values per event tuple. + :raises ValueError: If the tuple size is not a positive integer. + :returns: Instance of class mixed in. + """ + if value < 1: + raise ValueError(f"tupleSize must be a positive integer. Got {value}") + return self._set(tupleSize=value) + + def getTupleSize(self) -> int: + """ + Gets the tupleSize parameter. + + :returns: Number of discrete ID values per event tuple. + """ + return self.getOrDefault(self.tupleSize) + + def setTopK(self, value: int) -> "EventNgramLookupParams": + """ + Sets the topK parameter. + + :param value: Number of tokens emitted per event tuple. + :raises ValueError: If the number of tokens is not a positive integer. + :returns: Instance of class mixed in. + """ + if value < 1: + raise ValueError(f"topK must be a positive integer. Got {value}") + return self._set(topK=value) + + def getTopK(self) -> int: + """ + Gets the topK parameter. + + :returns: Number of tokens emitted per event tuple. + """ + return self.getOrDefault(self.topK) + + def setIncludeTokenTypes(self, value: bool) -> "EventNgramLookupParams": + """ + Sets the includeTokenTypes parameter. + + :param value: Whether to also emit per-input token-type columns. + :returns: Instance of class mixed in. + """ + return self._set(includeTokenTypes=value) + + def getIncludeTokenTypes(self) -> bool: + """ + Gets the includeTokenTypes parameter. + + :returns: Whether to also emit per-input token-type columns. + """ + return self.getOrDefault(self.includeTokenTypes) + + def getTokenTypeCols(self, token_cols: List[str]) -> List[str]: + """ + Gets the names of the derived token-type columns. + + With includeTokenTypes set, every token column ```` is accompanied by a + ``_types`` column. These are not listed in outputCols, so the estimator + and transformer both declare them alongside the token columns, in the order + the Keras layer returns its type tensors. + + :param token_cols: The token (output) column names. + :returns: One type column name per token column, or an empty list if + includeTokenTypes is not set. + """ + if not self.getIncludeTokenTypes(): + return [] + return [f"{col}_types" for col in token_cols] diff --git a/src/kamae/spark/transformers/__init__.py b/src/kamae/spark/transformers/__init__.py index a61ae8ea..4ffc2aa8 100644 --- a/src/kamae/spark/transformers/__init__.py +++ b/src/kamae/spark/transformers/__init__.py @@ -38,6 +38,7 @@ DateTimeToUnixTimestampTransformer, ) from .divide import DivideTransformer # noqa: F401 +from .event_ngram_lookup import EventNgramLookupTransformer # noqa: F401 from .exp import ExpTransformer # noqa: F401 from .exponent import ExponentTransformer # noqa: F401 from .hash_index import HashIndexTransformer # noqa: F401 diff --git a/src/kamae/spark/transformers/event_ngram_lookup.py b/src/kamae/spark/transformers/event_ngram_lookup.py new file mode 100644 index 00000000..f43a1059 --- /dev/null +++ b/src/kamae/spark/transformers/event_ngram_lookup.py @@ -0,0 +1,389 @@ +# Copyright [2024] Expedia, Inc. +# +# 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. + +""" +EventNgramLookupTransformer: tokenizes discrete IDs with a pre-computed lookup table. + +Applies the ``tuple -> top-k tokens`` table fitted by ``EventNgramLookupEstimator``. +The Spark path (``_transform``) and the Keras layer (``get_keras_layer``) apply the +same table, so both produce identical tokens. + +When ``includeTokenTypes`` is set, each token column ```` is accompanied by a +parallel ``_types`` column giving each token's ID-level bitmask. + +Token id conventions (shared with the vocabulary and the Keras layer): + - ``0`` = padding (````): emitted for all-zero / missing events. + - ``1`` = unknown (````): a non-zero tuple absent from the table yields a + single ```` followed by padding. +""" + +# pylint: disable=unused-argument +# pylint: disable=invalid-name +# pylint: disable=too-many-ancestors +# pylint: disable=no-member +from functools import partial +from typing import Dict, List, Optional, Tuple + +import pyspark.sql.functions as F +import tensorflow as tf +from pyspark import keyword_only +from pyspark.ml.param import Param, Params, TypeConverters +from pyspark.sql import DataFrame +from pyspark.sql.types import ( + ArrayType, + DataType, + IntegerType, + LongType, + StructField, + StructType, +) + +from kamae.keras.core.backend import TENSORFLOW_ONLY +from kamae.keras.tensorflow.layers import EventNgramLookupLayer +from kamae.spark.params import EventNgramLookupParams, MultiInputMultiOutputParams +from kamae.spark.utils import ( + tokenize_events, + tokenize_events_with_types, + validate_event_column_lengths, + validate_event_id_columns, +) + +from .base import BaseTransformer + + +class EventNgramLookupTransformerParams(Params): + """ + Mixin class containing the fitted state of the EventNgramLookupTransformer. + + These are produced by the ``EventNgramLookupEstimator``, so they live on the + transformer only. The lookup table is held as two parallel flat int lists because + Spark ML writes params to JSON, where a dict keyed by id tuples cannot be written. + """ + + vocabularySize = Param( + Params._dummy(), + "vocabularySize", + "Number of token ids in the fitted vocabulary, including the reserved pad and " + "unk tokens.", + typeConverter=TypeConverters.toInt, + ) + + lookupKeys = Param( + Params._dummy(), + "lookupKeys", + "Flattened event tuples of the fitted lookup table, tupleSize ids per tuple.", + typeConverter=TypeConverters.toListInt, + ) + + lookupValues = Param( + Params._dummy(), + "lookupValues", + "Flattened token lists of the fitted lookup table, topK tokens per tuple, " + "positionally aligned with lookupKeys.", + typeConverter=TypeConverters.toListInt, + ) + + tokenTypeLookup = Param( + Params._dummy(), + "tokenTypeLookup", + "Per-token-id list mapping each token to a bitmask of the ID levels its n-gram " + "spans (0 for pad/unk). Used only when includeTokenTypes is True.", + typeConverter=TypeConverters.toListInt, + ) + + def setVocabularySize(self, value: int) -> "EventNgramLookupTransformerParams": + """ + Sets the vocabularySize parameter. + + :param value: Number of token ids in the fitted vocabulary. + :returns: Instance of class mixed in. + """ + return self._set(vocabularySize=value) + + def getVocabularySize(self) -> int: + """ + Gets the vocabularySize parameter. + + :returns: Number of token ids in the fitted vocabulary. + """ + return self.getOrDefault(self.vocabularySize) + + def setLookupKeys(self, value: List[int]) -> "EventNgramLookupTransformerParams": + """ + Sets the lookupKeys parameter. + + :param value: Flattened event tuples, tupleSize ids per tuple. + :returns: Instance of class mixed in. + """ + return self._set(lookupKeys=value) + + def getLookupKeys(self) -> Optional[List[int]]: + """ + Gets the lookupKeys parameter. + + :returns: Flattened event tuples, or None if not set. + """ + return self.getOrDefault(self.lookupKeys) + + def setLookupValues(self, value: List[int]) -> "EventNgramLookupTransformerParams": + """ + Sets the lookupValues parameter. + + :param value: Flattened token lists, topK tokens per tuple. + :returns: Instance of class mixed in. + """ + return self._set(lookupValues=value) + + def getLookupValues(self) -> Optional[List[int]]: + """ + Gets the lookupValues parameter. + + :returns: Flattened token lists, or None if not set. + """ + return self.getOrDefault(self.lookupValues) + + def setTokenTypeLookup( + self, value: List[int] + ) -> "EventNgramLookupTransformerParams": + """ + Sets the tokenTypeLookup parameter. + + :param value: Per-token-id list of ID-level bitmasks. + :returns: Instance of class mixed in. + """ + return self._set(tokenTypeLookup=value) + + def getTokenTypeLookup(self) -> Optional[List[int]]: + """ + Gets the tokenTypeLookup parameter. + + :returns: Per-token-id list of ID-level bitmasks, or None. + """ + return self.getOrDefault(self.tokenTypeLookup) + + def getTupleToTokens(self) -> Dict[Tuple[int, ...], List[int]]: + """ + Rebuilds the ``event tuple -> top-k token list`` table from the flat params. + + :returns: Mapping from each event tuple to its list of token ids. + """ + keys = self.getLookupKeys() or [] + values = self.getLookupValues() or [] + tuple_size = self.getTupleSize() + top_k = self.getTopK() + return { + tuple(keys[i : i + tuple_size]): values[j : j + top_k] + for i, j in zip( + range(0, len(keys), tuple_size), range(0, len(values), top_k) + ) + } + + +class EventNgramLookupTransformer( + BaseTransformer, + MultiInputMultiOutputParams, + EventNgramLookupParams, + EventNgramLookupTransformerParams, +): + """ + Tokenizes discrete ID values using the pre-computed ``tuple -> top-k tokens`` table. + + Each input column holds a sequence of events; every event is a ``tupleSize``-long + group of IDs. For each event the transformer emits the tuple's ``topK`` token ids, + producing a flat ``numEvents * topK`` array per input column. When + ``includeTokenTypes`` is set, a parallel ``_types`` column of the same shape is + also produced, giving each token's ID-level bitmask. + + A row whose ids are not a whole number of events raises, as in the Keras layer. + """ + + supported_backends = TENSORFLOW_ONLY + jit_compatible = False + + @keyword_only + def __init__( + self, + inputCols: Optional[List[str]] = None, + outputCols: Optional[List[str]] = None, + inputDtype: Optional[str] = None, + outputDtype: Optional[str] = None, + layerName: Optional[str] = None, + numEventsPerInput: Optional[List[int]] = None, + tupleSize: int = 4, + topK: Optional[int] = None, + vocabularySize: Optional[int] = None, + lookupKeys: Optional[List[int]] = None, + lookupValues: Optional[List[int]] = None, + includeTokenTypes: bool = False, + tokenTypeLookup: Optional[List[int]] = None, + ) -> None: + """ + Initializes the EventNgramLookupTransformer transformer. + + :param inputCols: Input column names holding discrete ID values. + :param outputCols: Output column names for the token arrays. + :param inputDtype: Data type to cast the input columns to before transforming. + :param outputDtype: Data type to cast the output columns to after transforming. + :param layerName: Name of the layer. Used as the name of the Keras layer in the + keras model. If not set, we use the uid of the Spark transformer. + :param numEventsPerInput: Number of events per input column, e.g. [10, 10, 1]. + :param tupleSize: Number of discrete ID values (ID levels) per event tuple. + :param topK: Number of tokens emitted per event tuple. + :param vocabularySize: Number of token ids in the fitted vocabulary. + :param lookupKeys: Fitted lookup table keys, as the event tuples flattened to a + single int list (``tupleSize`` ids per tuple). + :param lookupValues: Fitted lookup table values, as the token lists flattened to + a single int list (``topK`` tokens per tuple), aligned with ``lookupKeys``. + :param includeTokenTypes: If `True`, also emit a ``_types`` column per + input giving each token's ID-level bitmask. Defaults to `False`. + :param tokenTypeLookup: Per-token-id list of ID-level bitmasks (from the fitted + vocabulary). Required when ``includeTokenTypes`` is `True`. + :returns: None - class instantiated. + """ + super().__init__() + self._setDefault( + numEventsPerInput=None, + tupleSize=4, + topK=None, + vocabularySize=None, + lookupKeys=None, + lookupValues=None, + includeTokenTypes=False, + tokenTypeLookup=None, + ) + kwargs = self._input_kwargs + self.setParams(**kwargs) + + @property + def compatible_dtypes(self) -> Optional[List[DataType]]: + """ + List of compatible data types for the layer. + If the computation can be performed on any data type, return None. + + :returns: List of compatible data types for the layer. + """ + return [IntegerType(), LongType()] + + def _transform(self, dataset: DataFrame) -> DataFrame: + """ + Tokenizes each input column with a per-column UDF. + + Each input value is a flat integer array, split into consecutive events of + ``tupleSize`` ids. The output is a flat integer array of length + ``numEvents * topK``. When ``includeTokenTypes`` is set, the same UDF also + returns each token's type bitmask, written to the ``_types`` column. + + :param dataset: Input DataFrame. + :raises ValueError: If the input columns, output columns and event counts + differ in length, or a column is missing or is not a single-level array. + :returns: DataFrame with the tokenized output columns, and the optional type + columns. + """ + input_cols = self.getInputCols() + output_cols = self.getOutputCols() + num_events_per_input = self.getNumEventsPerInput() + validate_event_column_lengths(input_cols, output_cols, num_events_per_input) + validate_event_id_columns(dataset, input_cols) + + table_kwargs = { + "lookup_table": self.getTupleToTokens(), + "tuple_size": self.getTupleSize(), + "top_k": self.getTopK(), + } + token_array_type = ArrayType(IntegerType()) + if not self.getIncludeTokenTypes(): + for input_col, output_col, num_events in zip( + input_cols, output_cols, num_events_per_input + ): + tokenize_udf = F.udf( + partial(tokenize_events, num_events=num_events, **table_kwargs), + token_array_type, + ) + dataset = dataset.withColumn(output_col, tokenize_udf(F.col(input_col))) + return dataset + + # Tokens and types come from one UDF returning a struct, so each row crosses + # into Python once. + struct_type = StructType( + [ + StructField("tokens", token_array_type), + StructField("types", token_array_type), + ] + ) + for input_col, output_col, type_col, num_events in zip( + input_cols, + output_cols, + self.getTokenTypeCols(output_cols), + num_events_per_input, + ): + tokenize_udf = F.udf( + partial( + tokenize_events_with_types, + token_type_lookup=self.getTokenTypeLookup(), + num_events=num_events, + **table_kwargs, + ), + struct_type, + ) + struct_col = f"{output_col}__tokens_and_types" + # The type columns are not in outputCols, so the base class's output cast + # does not reach them; cast them here to match the Keras layer. + casted_types, _ = self._cast_output_columns( + [F.col(f"{struct_col}.types")], [token_array_type] + )[0] + dataset = ( + dataset.withColumn(struct_col, tokenize_udf(F.col(input_col))) + .withColumn(output_col, F.col(f"{struct_col}.tokens")) + .withColumn(type_col, casted_types) + .drop(struct_col) + ) + return dataset + + def get_keras_layer(self) -> tf.keras.layers.Layer: + """ + Gets the Keras layer for the EventNgramLookup transformer. + + Returns a single ``EventNgramLookupLayer`` that tokenizes all input columns, so + the fitted lookup table is embedded once. The layer outputs one token tensor + per input, plus one type tensor per input when ``includeTokenTypes`` is set. + + :returns: The consolidated ``EventNgramLookupLayer``. + """ + lookup_table = self.getTupleToTokens() + return EventNgramLookupLayer( + num_events_per_input=self.getNumEventsPerInput(), + top_k=self.getTopK(), + tuple_size=self.getTupleSize(), + lookup_keys=[list(key) for key in lookup_table], + lookup_values=list(lookup_table.values()), + token_type_lookup=( + self.getTokenTypeLookup() if self.getIncludeTokenTypes() else None + ), + input_dtype=self.getInputKerasDtype(), + output_dtype=self.getOutputKerasDtype(), + name=self.getLayerName(), + ) + + def get_layer_inputs_outputs(self) -> Tuple[List[str], List[str]]: + """ + Gets the input and output column names, including the token-type columns. + + Overrides the base method because, with ``includeTokenTypes`` set, the layer + returns the per-input type tensors after the token tensors, so the + ``_types`` columns are appended in that same order. + + :returns: Tuple of the input column names and the output column names. + """ + inputs, token_cols = super().get_layer_inputs_outputs() + return inputs, token_cols + self.getTokenTypeCols(token_cols) diff --git a/src/kamae/spark/utils/__init__.py b/src/kamae/spark/utils/__init__.py index 39e0c3e9..478d253c 100644 --- a/src/kamae/spark/utils/__init__.py +++ b/src/kamae/spark/utils/__init__.py @@ -33,6 +33,16 @@ check_listwise_columns, get_listwise_condition_and_window, ) +from .ngram_utils import ( # noqa: F401 + EventNgramVocabulary, + build_tuple_lookup_table, + build_vocabulary, + collect_ngrams_from_dataframe, + tokenize_events, + tokenize_events_with_types, + validate_event_column_lengths, + validate_event_id_columns, +) from .transform_utils import ( # noqa: F401 multi_input_single_output_array_transform, multi_input_single_output_scalar_transform, diff --git a/src/kamae/spark/utils/ngram_utils.py b/src/kamae/spark/utils/ngram_utils.py new file mode 100644 index 00000000..055c5fef --- /dev/null +++ b/src/kamae/spark/utils/ngram_utils.py @@ -0,0 +1,391 @@ +# Copyright [2024] Expedia, Inc. +# +# 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. + +""" +N-gram vocabulary utilities for event-level discrete-ID tokenization. + +Each event is a fixed-size tuple of IDs (e.g. 4 levels ``L0, L1, L2, L3``) +and n-grams are learned only within event boundaries, never across events. The public +helpers are ``collect_ngrams_from_dataframe`` (distributed n-gram counting), +``build_vocabulary`` (frequency selection) and ``build_tuple_lookup_table`` (pre-compute +the tuple -> top-k tokens table for O(1) inference). +""" + +import logging +from collections import Counter +from typing import Dict, List, Optional, Tuple + +from pyspark.sql import DataFrame + +from kamae.keras.tensorflow.layers.event_ngram_lookup import PAD_TOKEN_ID, UNK_TOKEN_ID +from kamae.spark.utils.array_utils import get_array_nesting_level +from kamae.spark.utils.ngram_worker_functions import ( + NUM_RESERVED_TOKENS, + encode_tuple, + extract_ngrams_from_column_worker, + extract_tuples_from_column_worker, + iter_events, +) + +logger = logging.getLogger(__name__) + + +class EventNgramVocabulary: + """ + Vocabulary of n-grams learned from within-event patterns only. + + Each event is a tuple of IDs and n-grams are learned only within event + boundaries. Token id 0 is reserved for ```` and 1 for ````; learned + n-grams are assigned ids from 2 upwards, contiguously and in descending frequency + order, so a smaller token id always means a more frequent n-gram. + """ + + def __init__(self, ngrams: Dict[Tuple[str, ...], int]) -> None: + """ + Initializes the n-gram vocabulary. + + :param ngrams: Mapping from n-gram tuple to token id, as returned by + ``build_vocabulary``. + """ + self.ngrams = ngrams + + @property + def vocab_size(self) -> int: + """ + Returns the total number of token ids, including the two reserved tokens. + + Token ids are contiguous from ``NUM_RESERVED_TOKENS``, so this is also one past + the largest token id. + + :returns: Number of token ids in the vocabulary. + """ + return len(self.ngrams) + NUM_RESERVED_TOKENS + + def build_type_lookup(self) -> List[int]: + """ + Builds a per-token "type" lookup: token id -> bitmask of contributing ID levels. + + Each n-gram is a tuple of level-prefixed ids (``L_``). The type + encodes *which* ID levels the n-gram spans as a bitmask with bit ``level`` set + for each present level (e.g. an ``(L0, L2)`` skip-gram -> ``0b0101 = 5``). The + reserved ```` and ```` tokens map to ``0``, and the type cardinality + is ``2 ** tupleSize``. + + :returns: List indexed by token id giving each token's level bitmask. + """ + lookup = [0] * self.vocab_size + for ngram, token_id in self.ngrams.items(): + bitmask = 0 + for part in ngram: + # part is "L_"; take the integer level after the "L". + bitmask |= 1 << int(part[1:].split("_")[0]) + lookup[token_id] = bitmask + return lookup + + +def validate_event_column_lengths( + input_columns: List[str], + output_columns: List[str], + num_events_per_input: Optional[List[int]], +) -> None: + """ + Checks that there is one output column and one event count per input column. + + The columns are tokenized pairwise, so a length mismatch would otherwise be + silently truncated to the shortest list. + + :param input_columns: Column names holding discrete ID values. + :param output_columns: Column names for the token arrays. + :param num_events_per_input: Number of events per input column. + :raises ValueError: If the three lists do not all have the same length. + :returns: None - the lengths are validated. + """ + if len(input_columns) != len(output_columns): + raise ValueError( + f"inputCols and outputCols must have the same length. Got " + f"{len(input_columns)} inputs and {len(output_columns)} outputs." + ) + if num_events_per_input is None or len(input_columns) != len(num_events_per_input): + n_events = 0 if num_events_per_input is None else len(num_events_per_input) + raise ValueError( + f"numEventsPerInput must have one entry per input column. Got " + f"{len(input_columns)} inputs and {n_events} event counts." + ) + + +def validate_event_id_columns(df: DataFrame, input_columns: List[str]) -> None: + """ + Checks that every input column is present and is a flat array of ids. + + Each column value is a flat integer array chunked into consecutive events, so a + scalar column has no ids to chunk and a nested array (e.g. ``array>``) + would be chunked over its sub-arrays rather than its ids. Nested (e.g. listwise) + columns are therefore not yet supported on the Spark side, although the Keras + layer accepts rank-3 inputs. + + :param df: DataFrame holding the discrete-ID columns. + :param input_columns: Column names holding discrete ID values. + :raises ValueError: If a column is missing, or is not a single-level array. + :returns: None - the columns are validated. + """ + missing_cols = [c for c in input_columns if c not in df.columns] + if missing_cols: + # Training the vocabulary on a subset of the requested columns would silently + # produce a vocabulary that does not cover every column being tokenized. + raise ValueError( + f"Input columns not found on the DataFrame: {missing_cols}. " + f"Available columns: {df.columns}" + ) + + for column_name in input_columns: + nesting_level = get_array_nesting_level( + column_dtype=df.schema[column_name].dataType + ) + if nesting_level != 1: + raise ValueError( + f"Input column {column_name} must be a single-level array of discrete " + f"ID values, but it has an array nesting level of {nesting_level}." + ) + + +def collect_ngrams_from_dataframe( + df: DataFrame, + input_columns: List[str], + event_size: int, + min_ngram_freq: int = 1, +) -> Counter: + """ + Extracts and counts within-event n-grams across the given columns, distributed. + + All columns are counted in a single pass: every n-gram of every column of a row is + emitted by one ``flatMap``, so the dataset is read once no matter how many columns + are tokenized. Counts are aggregated with ``reduceByKey`` (map-side combine) and + filtered by ``min_ngram_freq`` on the executors, so only the distinct n-grams that + can enter the vocabulary reach the driver. + + :param df: Input DataFrame with the discrete-ID columns. + :param input_columns: Column names holding discrete ID values. + :param event_size: Number of discrete ID values per event. + :param min_ngram_freq: Minimum corpus frequency for an n-gram to be collected. + Defaults to 1 (every n-gram). + :returns: Counter mapping each collected n-gram tuple to its corpus frequency. + :raises ValueError: If a column is missing, or is not a single-level array. + """ + validate_event_id_columns(df, input_columns) + + logger.info("Collecting n-grams from %d columns in one pass...", len(input_columns)) + + ngram_counts = ( + df.select(*input_columns) + .rdd.flatMap( + lambda row: [ + ngram + for col_value in row + for ngram in extract_ngrams_from_column_worker(col_value, event_size) + ] + ) + .map(lambda ngram: (ngram, 1)) + .reduceByKey(lambda a, b: a + b) + .filter(lambda ngram_count: ngram_count[1] >= min_ngram_freq) + .collect() + ) + + all_ngrams = Counter(dict(ngram_counts)) + logger.info( + "Found %d unique n-grams with frequency >= %d across all columns", + len(all_ngrams), + min_ngram_freq, + ) + + return all_ngrams + + +def log_vocabulary_by_length(ngram_to_id: Dict[Tuple[str, ...], int]) -> None: + """ + Logs how the vocabulary splits across n-gram lengths. + + A vocabulary dominated by 1-grams means the n-gram combinations are earning + little, which is the signal for tuning ``vocab_size`` and ``min_ngram_freq``. + + :param ngram_to_id: The fitted vocabulary, n-gram tuple to token id. + :returns: None - the distribution is logged. + """ + if not ngram_to_id: + logger.info("Vocabulary by n-gram length: empty vocabulary") + return + by_length = Counter(len(ngram) for ngram in ngram_to_id) + total = len(ngram_to_id) + parts = [ + f"{length}-gram {by_length[length]:,} " + f"({100.0 * by_length[length] / total:.1f}%)" + for length in sorted(by_length) + ] + logger.info("Vocabulary by n-gram length: %s", " | ".join(parts)) + + +def build_vocabulary( + ngram_counter: Counter, + vocab_size: int, +) -> Dict[Tuple[str, ...], int]: + """ + Builds the vocabulary by keeping the most frequent n-grams. + + The counter is expected to be already filtered by frequency, as returned by + ``collect_ngrams_from_dataframe``. N-grams are sorted by descending frequency + (ties broken by the n-gram itself, for determinism) and the top + ``vocab_size - NUM_RESERVED_TOKENS`` are assigned contiguous ids from + ``NUM_RESERVED_TOKENS`` upwards. Because ids ascend as frequency descends, + "the top-k most frequent matching n-grams" is later just an ascending sort of ids. + + :param ngram_counter: Counter mapping n-gram tuples to frequencies. + :param vocab_size: Target vocabulary size including the two reserved tokens. + :returns: Mapping from each kept n-gram tuple to its token id. + """ + # Sort by frequency descending, then by the n-gram itself so that equal-frequency + # n-grams are ordered deterministically across runs. + top_ngrams = sorted(ngram_counter.items(), key=lambda x: (-x[1], x[0]))[ + : max(vocab_size - NUM_RESERVED_TOKENS, 0) + ] + ngram_to_id = { + ngram: token_id + for token_id, (ngram, _) in enumerate(top_ngrams, start=NUM_RESERVED_TOKENS) + } + + logger.info( + "Final vocabulary size: %d (%d n-grams + %d reserved tokens)", + len(ngram_to_id) + NUM_RESERVED_TOKENS, + len(ngram_to_id), + NUM_RESERVED_TOKENS, + ) + log_vocabulary_by_length(ngram_to_id) + + return ngram_to_id + + +def build_tuple_lookup_table( + df: DataFrame, + input_columns: List[str], + vocabulary: EventNgramVocabulary, + top_k: int, + event_size: int, +) -> Dict[Tuple[int, ...], List[int]]: + """ + Pre-computes the ``id_tuple -> top-k token ids`` lookup table, distributed. + + All input columns are read in one pass and deduped, so each distinct event tuple is + encoded exactly once, on the workers, via ``encode_tuple``. That encoding is a pure + function of the tuple and the fitted vocabulary, so the assembled table does not + depend on how the work was partitioned. + + Tuples that matched no n-gram are left out of the table, so they resolve to + ```` in the same way as a tuple never seen during fitting. + + :param df: DataFrame with the discrete-ID columns. + :param input_columns: Column names holding discrete ID values. + :param vocabulary: Fitted ``EventNgramVocabulary``. + :param top_k: Number of tokens per tuple. + :param event_size: Number of discrete ID values per tuple. + :returns: Mapping from each token-bearing event tuple to its ``top_k`` token ids. + """ + logger.info("Building tuple->tokens lookup table (distributed)...") + ngrams = vocabulary.ngrams + + tuple_to_tokens = ( + df.select(*input_columns) + .rdd.flatMap( + lambda row: [ + id_tuple + for col_value in row + for id_tuple in extract_tuples_from_column_worker(col_value, event_size) + ] + ) + .distinct() + .map(lambda id_tuple: (id_tuple, encode_tuple(id_tuple, ngrams, top_k))) + .filter(lambda kv: kv[1] is not None) + .collectAsMap() + ) + + logger.info("Encoded %d unique tuples", len(tuple_to_tokens)) + + return tuple_to_tokens + + +def tokenize_events( + ids: Optional[List[int]], + lookup_table: Dict[Tuple[int, ...], List[int]], + num_events: int, + tuple_size: int, + top_k: int, +) -> List[int]: + """ + Tokenizes one row's discrete ID values into a flat token array via the lookup table. + + Splits the row into per-event tuples and maps each to its tokens: an all-zero event + yields padding, a tuple in the table yields its stored tokens, and any other + non-zero tuple yields a single ```` followed by padding. Mirrored by the + Keras ``EventNgramLookupLayer``. + + :param ids: One row's discrete ID values as a flat int array, split into + consecutive events of ``tuple_size`` ids. May be `None` or empty, and individual + ids may be null (read as ``0``). + :param lookup_table: Fitted mapping from event tuple to its top-k token list. + :param num_events: Number of events the output is padded/truncated to. + :param tuple_size: Number of discrete ID values per event tuple. + :param top_k: Number of tokens per event tuple. + :raises ValueError: If the number of ids is not a multiple of ``tuple_size``. + :returns: Flat token array of length ``num_events * top_k``. + """ + pad_tokens = [PAD_TOKEN_ID] * top_k + unk_tokens = [UNK_TOKEN_ID] + [PAD_TOKEN_ID] * (top_k - 1) + + if not ids: + return pad_tokens * num_events + + all_tokens: List[int] = [] + for event_ids in iter_events(ids, tuple_size): + event_tuple = tuple(event_ids) + if all(id_value == 0 for id_value in event_tuple): + all_tokens.extend(pad_tokens) + else: + all_tokens.extend(lookup_table.get(event_tuple, unk_tokens)) + + # Pad/truncate to the fixed output length. + expected_length = num_events * top_k + if len(all_tokens) < expected_length: + all_tokens.extend([PAD_TOKEN_ID] * (expected_length - len(all_tokens))) + return all_tokens[:expected_length] + + +def tokenize_events_with_types( + ids: Optional[List[int]], + lookup_table: Dict[Tuple[int, ...], List[int]], + token_type_lookup: List[int], + num_events: int, + tuple_size: int, + top_k: int, +) -> Tuple[List[int], List[int]]: + """ + Tokenizes one row as ``tokenize_events`` does, and gathers each token's type. + + :param ids: One row's discrete ID values as a flat int array. + :param lookup_table: Fitted mapping from event tuple to its top-k token list. + :param token_type_lookup: Per-token-id list of ID-level bitmasks. + :param num_events: Number of events the output is padded/truncated to. + :param tuple_size: Number of discrete ID values per event tuple. + :param top_k: Number of tokens per event tuple. + :returns: Tuple of (token ids, type bitmasks), both of length + ``num_events * top_k``. + """ + tokens = tokenize_events(ids, lookup_table, num_events, tuple_size, top_k) + return tokens, [token_type_lookup[token] for token in tokens] diff --git a/src/kamae/spark/utils/ngram_worker_functions.py b/src/kamae/spark/utils/ngram_worker_functions.py new file mode 100644 index 00000000..a434f58b --- /dev/null +++ b/src/kamae/spark/utils/ngram_worker_functions.py @@ -0,0 +1,144 @@ +# Copyright [2024] Expedia, Inc. +# +# 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. + +""" +Worker functions for n-gram extraction in Spark RDD operations. + +These run inside ``flatMap`` / ``map`` on the Spark workers and are pure Python over +plain ints and tuples. +""" + +from itertools import combinations +from typing import Any, Dict, Iterator, List, Optional, Sequence, Tuple + +from kamae.keras.tensorflow.layers.event_ngram_lookup import PAD_TOKEN_ID + +NUM_RESERVED_TOKENS = 2 # and ; learned n-grams get ids from here upwards + + +def iter_events(col_value: Any, event_size: int) -> Iterator[List[int]]: + """ + Yields each event of a column value as a list of ``event_size`` int ids. + + The column value is a flat integer array chunked into consecutive events, so its + length must be a whole number of events. A null id is read as ``0``, i.e. as an + absent ID level. + + :param col_value: Column value containing discrete ID values. + :param event_size: Number of discrete ID values per event. + :raises ValueError: If the number of ids is not a multiple of ``event_size``. + :returns: Iterator over per-event lists of int ids. + """ + if col_value is None: + return + if len(col_value) % event_size != 0: + raise ValueError( + f"Number of discrete ID values must be a whole number of events: got " + f"{len(col_value)} ids, which is not a multiple of the event size " + f"{event_size}." + ) + for i in range(0, len(col_value), event_size): + yield [ + 0 if id_value is None else id_value + for id_value in col_value[i : i + event_size] + ] + + +def event_ngrams(event_ids: Sequence[int]) -> Iterator[Tuple[str, ...]]: + """ + Yields every within-event n-gram of a single event. + + The event's non-zero ids are level-prefixed (``L_``) and every + combination of length 1..n is emitted, so n-grams that skip ID levels are included. + Zero ids denote absent ID levels and take no part in any n-gram. Used both when + counting n-grams and when encoding a tuple against a fitted vocabulary. + + :param event_ids: One event's ids, one per ID level. + :returns: Iterator over the event's n-gram tuples. + """ + present = [ + f"L{level}_{id_value}" + for level, id_value in enumerate(event_ids) + if id_value != 0 + ] + for length in range(1, len(present) + 1): + yield from combinations(present, length) + + +def extract_ngrams_from_column_worker( + col_value: Any, event_size: int +) -> List[Tuple[str, ...]]: + """ + Extracts all within-event n-grams from a single row's column value. + + N-grams are never formed across event boundaries. + + :param col_value: Column value containing discrete ID values. + :param event_size: Number of discrete ID values per event. + :returns: List of n-gram tuples. + """ + return [ + ngram + for event_ids in iter_events(col_value, event_size) + for ngram in event_ngrams(event_ids) + ] + + +def extract_tuples_from_column_worker( + col_value: Any, event_size: int +) -> List[Tuple[int, ...]]: + """ + Extracts all non-zero event tuples from a single row's column value. + + :param col_value: Column value containing discrete ID values. + :param event_size: Number of discrete ID values per tuple. + :returns: List of ID tuples (all-zero events are skipped). + """ + return [ + tuple(event_ids) + for event_ids in iter_events(col_value, event_size) + if any(id_value != 0 for id_value in event_ids) + ] + + +def encode_tuple( + id_tuple: Tuple[int, ...], + ngrams: Dict[Tuple[str, ...], int], + top_k: int, +) -> Optional[List[int]]: + """ + Encodes one event tuple into its top-k token ids against a fitted vocabulary. + + Keeps the tuple's n-grams that are present in ``ngrams`` and returns the ``top_k`` + smallest ids; ids are assigned by descending frequency, so a smaller id is a more + frequent n-gram. + + Returns `None` when the tuple is all-zero or none of its n-grams is in the + vocabulary. Such a tuple is left out of the lookup table, so at inference it + resolves to ```` like a tuple never seen during fitting. + + :param id_tuple: Event tuple of IDs. + :param ngrams: Fitted mapping from n-gram tuple to token id. + :param top_k: Number of tokens to return. + :returns: List of token ids of length ``top_k`` (padded if needed), or `None` if the + tuple matched no n-gram. + """ + token_ids = sorted( + ngrams[ngram] for ngram in event_ngrams(id_tuple) if ngram in ngrams + ) + if not token_ids: + return None + + tokens = token_ids[:top_k] + return tokens + [PAD_TOKEN_ID] * (top_k - len(tokens)) diff --git a/tests/kamae/keras/tensorflow/layers/test_event_ngram_lookup.py b/tests/kamae/keras/tensorflow/layers/test_event_ngram_lookup.py new file mode 100644 index 00000000..fc981827 --- /dev/null +++ b/tests/kamae/keras/tensorflow/layers/test_event_ngram_lookup.py @@ -0,0 +1,282 @@ +# Copyright [2024] Expedia, Inc. +# +# 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. + +import keras +import pytest +import tensorflow as tf + +from kamae.keras.tensorflow.layers import EventNgramLookupLayer + +# Fitted table shared across the tests: two known 4-tuples, top_k = 3, held as the +# parallel key/value lists the layer takes. +LOOKUP_KEYS = [[1, 2, 3, 4], [5, 6, 7, 8]] +LOOKUP_VALUES = [[2, 3, 0], [4, 5, 0]] +TOP_K = 3 +TUPLE_SIZE = 4 + + +def _layer(**overrides) -> EventNgramLookupLayer: + params = dict( + name="event_ngram_lookup", + num_events_per_input=[2], + top_k=TOP_K, + tuple_size=TUPLE_SIZE, + lookup_keys=LOOKUP_KEYS, + lookup_values=LOOKUP_VALUES, + ) + params.update(overrides) + return EventNgramLookupLayer(**params) + + +class TestEventNgramLookupLayer: + def test_layer_name_and_dtypes(self): + layer = _layer() + assert layer.name == "event_ngram_lookup", "Layer name is not set properly" + assert layer.compatible_dtypes == ["int32", "int64"] + + @pytest.mark.parametrize("input_dtype", [tf.int32, tf.int64]) + def test_known_unknown_and_padding_tuples(self, input_dtype): + # given + inputs = tf.constant( + [ + [1, 2, 3, 4, 5, 6, 7, 8], # known, known + [1, 2, 3, 4, 9, 9, 9, 9], # known, unknown -> unk + [0, 0, 0, 0, 0, 0, 0, 0], # padding, padding -> pad + ], + dtype=input_dtype, + ) + + # when + output = _layer()(inputs) + + # then + assert output.shape == (3, 2 * TOP_K) + tf.debugging.assert_equal( + output, + tf.constant( + [[2, 3, 0, 4, 5, 0], [2, 3, 0, 1, 0, 0], [0, 0, 0, 0, 0, 0]], + dtype=output.dtype, + ), + ) + + def test_preserves_list_dimension_for_rank_three_input(self): + # given: (batch=2, list_size=2, num_events * tuple_size=8) + inputs = tf.constant( + [ + [[1, 2, 3, 4, 5, 6, 7, 8], [0, 0, 0, 0, 0, 0, 0, 0]], + [[5, 6, 7, 8, 1, 2, 3, 4], [9, 9, 9, 9, 0, 0, 0, 0]], + ], + dtype=tf.int32, + ) + + # when + output = _layer()(inputs) + + # then + assert output.shape == (2, 2, 2 * TOP_K) + tf.debugging.assert_equal( + output, + tf.constant( + [ + [[2, 3, 0, 4, 5, 0], [0, 0, 0, 0, 0, 0]], + [[4, 5, 0, 2, 3, 0], [1, 0, 0, 0, 0, 0]], + ], + dtype=output.dtype, + ), + ) + + @pytest.mark.parametrize( + "inputs, expected", + [ + # Short id axis is zero-padded up to num_events * tuple_size. + ([[1, 2, 3, 4]], [[2, 3, 0, 0, 0, 0]]), + # Long id axis is truncated to num_events * tuple_size. + ([[1, 2, 3, 4, 5, 6, 7, 8, 1, 2, 3, 4]], [[2, 3, 0, 4, 5, 0]]), + ], + ) + def test_pads_and_truncates_the_id_axis(self, inputs, expected): + output = _layer()(tf.constant(inputs, dtype=tf.int32)) + tf.debugging.assert_equal(output, tf.constant(expected, dtype=output.dtype)) + + @pytest.mark.parametrize("width", [3, 6, 9]) + def test_raises_when_a_known_width_is_not_a_whole_number_of_events(self, width): + # Matches the Spark transformer, which raises on the same inputs. + with pytest.raises(ValueError, match="whole number of events"): + _layer()(tf.ones((1, width), dtype=tf.int32)) + with pytest.raises(ValueError, match="whole number of events"): + _layer()(keras.Input(shape=(None, width), dtype="int32")) + + def test_unknown_width_is_checked_per_batch(self): + # With no static width, the whole-events check runs in the graph instead. + tokenize = tf.function( + _layer(), input_signature=[tf.TensorSpec([None, None], tf.int32)] + ) + output = tokenize(tf.constant([[1, 2, 3, 4]], dtype=tf.int32)) + tf.debugging.assert_equal( + output, tf.constant([[2, 3, 0, 0, 0, 0]], dtype=output.dtype) + ) + with pytest.raises(tf.errors.InvalidArgumentError, match="whole number"): + tokenize(tf.constant([[1, 2, 3, 4, 5, 6]], dtype=tf.int32)) + + def test_tokenizes_every_input_with_its_own_event_count(self): + # given: two inputs with different numbers of events, one shared table. + layer = _layer(num_events_per_input=[2, 1]) + clicks = tf.constant([[1, 2, 3, 4, 5, 6, 7, 8]], dtype=tf.int32) + prop = tf.constant([[5, 6, 7, 8]], dtype=tf.int32) + + # when + tokens = layer([clicks, prop]) + + # then + assert len(tokens) == 2 + assert tokens[0].shape == (1, 2 * TOP_K) + assert tokens[1].shape == (1, 1 * TOP_K) + + def test_emits_a_type_tensor_per_input_after_the_token_tensors(self): + # given: token id -> ID-level bitmask. + type_lookup = [0, 0, 5, 1, 8, 3] + layer = _layer(token_type_lookup=type_lookup) + inputs = tf.constant([[1, 2, 3, 4, 5, 6, 7, 8]], dtype=tf.int32) + + # when + tokens, types = layer(inputs) + + # then + assert types.shape == tokens.shape + tf.debugging.assert_equal( + types, + tf.constant( + [[type_lookup[t] for t in tokens.numpy()[0].tolist()]], + dtype=types.dtype, + ), + ) + + @pytest.mark.parametrize("tuple_size", [1, 2, 3, 5]) + def test_supports_tuple_sizes_other_than_four(self, tuple_size): + # given + layer = _layer( + num_events_per_input=[1], + tuple_size=tuple_size, + lookup_keys=[list(range(1, tuple_size + 1))], + lookup_values=[[2] * TOP_K], + ) + + # when + output = layer(tf.constant([list(range(1, tuple_size + 1))], dtype=tf.int32)) + + # then + tf.debugging.assert_equal( + output, tf.constant([[2] * TOP_K], dtype=output.dtype) + ) + + def test_empty_table_maps_every_non_zero_tuple_to_unk(self): + layer = _layer(lookup_keys=[], lookup_values=[]) + output = layer(tf.constant([[1, 2, 3, 4, 0, 0, 0, 0]], dtype=tf.int32)) + tf.debugging.assert_equal( + output, tf.constant([[1, 0, 0, 0, 0, 0]], dtype=output.dtype) + ) + + def test_out_of_range_id_misses_instead_of_aliasing_a_valid_key(self): + # given: key_bits is sized from the largest id in the table (8 -> 4 bits), so + # id 17 does not fit. Clamping alone would turn (17, 2, 3, 4) into (15, 2, 3, 4) + # and could alias onto a real key, so out-of-range tuples must miss. + layer = _layer(num_events_per_input=[1]) + assert layer.key_bits == 4 + + # when + output = layer(tf.constant([[17, 2, 3, 4]], dtype=tf.int32)) + + # then + tf.debugging.assert_equal(output, tf.constant([[1, 0, 0]], dtype=output.dtype)) + + def test_negative_id_misses_instead_of_aliasing_a_valid_key(self): + # given: a table key whose first ID level is 0, which is what clamping a + # negative id would produce, so a negative id must miss rather than resolve + # to this key's tokens. + layer = _layer( + num_events_per_input=[1], + lookup_keys=[[0, 6, 7, 8]], + lookup_values=[[2, 3, 0]], + ) + + # when + output = layer(tf.constant([[-3, 6, 7, 8]], dtype=tf.int32)) + + # then + tf.debugging.assert_equal(output, tf.constant([[1, 0, 0]], dtype=output.dtype)) + + def test_raises_when_the_table_contains_a_negative_id(self): + # Packing allots a fixed number of bits per ID level, so a negative id would + # collide with the non-negative tuple that shares its remaining levels. + with pytest.raises(ValueError, match="non-negative"): + _layer( + num_events_per_input=[1], + lookup_keys=[[-1, 6, 7, 8]], + lookup_values=[[2, 3, 0]], + ) + + def test_raises_when_keys_do_not_fit_in_an_int64(self): + # 16 levels of a 2**32-sized id space needs far more than 63 bits. + with pytest.raises(ValueError, match="int64"): + _layer( + num_events_per_input=[1], + tuple_size=16, + lookup_keys=[[2**32] * 16], + lookup_values=[[2] * TOP_K], + ) + + def test_get_config_round_trips_the_table(self): + # given + layer = _layer(token_type_lookup=[0, 0, 5, 1, 8, 3]) + inputs = tf.constant([[1, 2, 3, 4, 9, 9, 9, 9]], dtype=tf.int32) + expected_tokens, expected_types = layer(inputs) + + # when: the table is persisted as parallel key/value lists, so the rebuilt + # layer must tokenize identically without a custom from_config. + rebuilt = EventNgramLookupLayer.from_config(layer.get_config()) + tokens, types = rebuilt(inputs) + + # then + tf.debugging.assert_equal(tokens, expected_tokens) + tf.debugging.assert_equal(types, expected_types) + + def test_symbolic_outputs_are_int32_with_the_tokenized_shape(self): + # Keras infers the symbolic outputs by tracing _call, so they carry the real + # int32 dtype and a downstream integer-only layer can be chained onto them. + layer = _layer( + num_events_per_input=[2, 1], token_type_lookup=[0, 0, 5, 1, 8, 3] + ) + clicks = keras.Input(shape=(8,), dtype="int32") + prop = keras.Input(shape=(None, 4), dtype="int64") + + outputs = layer([clicks, prop]) + + assert [(output.dtype, output.shape) for output in outputs] == [ + ("int32", (None, 2 * TOP_K)), + ("int32", (None, None, TOP_K)), + ("int32", (None, 2 * TOP_K)), + ("int32", (None, None, TOP_K)), + ] + + def test_raises_when_inputs_do_not_match_num_events_per_input(self): + # Inputs are paired with their event counts, so a mismatch would otherwise + # silently drop the unpaired inputs' outputs. + layer = _layer(num_events_per_input=[2, 1]) + with pytest.raises(ValueError, match="num_events_per_input"): + layer(tf.constant([[1, 2, 3, 4, 5, 6, 7, 8]], dtype=tf.int32)) + + @pytest.mark.parametrize("shape", [(8,), (1, 1, 1, 8)]) + def test_raises_for_inputs_that_are_not_rank_two_or_three(self, shape): + with pytest.raises(ValueError, match="rank"): + _layer()(tf.zeros(shape, dtype=tf.int32)) diff --git a/tests/kamae/keras/test_jit_compatibility.py b/tests/kamae/keras/test_jit_compatibility.py index 2ec2d600..2e7839f9 100644 --- a/tests/kamae/keras/test_jit_compatibility.py +++ b/tests/kamae/keras/test_jit_compatibility.py @@ -70,6 +70,7 @@ DateDiffLayer, DateParseLayer, DateTimeToUnixTimestampLayer, + EventNgramLookupLayer, HashIndexLayer, IfStatementLayer, LambdaFunctionLayer, @@ -285,6 +286,17 @@ [tf.strings.as_string(tf.random.normal((100, 23, 32, 1)))], {"num_hash_fns": 3, "num_bins": 100}, ), + ( + EventNgramLookupLayer, + [tf.constant([[1, 2, 3, 4, 5, 0, 0, 0]], dtype=tf.int32)], + { + "num_events_per_input": [2], + "top_k": 3, + "tuple_size": 4, + "lookup_keys": [[1, 2, 3, 4], [5, 0, 0, 0]], + "lookup_values": [[2, 3, 0], [4, 0, 0]], + }, + ), (CurrentDateLayer, [tf.constant(100, shape=(100, 10, 1))], None), (CurrentDateTimeLayer, [tf.constant(100, shape=(100, 10, 1))], None), ( diff --git a/tests/kamae/keras/test_layer_serialisation.py b/tests/kamae/keras/test_layer_serialisation.py index a4679777..e3537bf9 100644 --- a/tests/kamae/keras/test_layer_serialisation.py +++ b/tests/kamae/keras/test_layer_serialisation.py @@ -77,6 +77,7 @@ DateDiffLayer, DateParseLayer, DateTimeToUnixTimestampLayer, + EventNgramLookupLayer, HashIndexLayer, IfStatementLayer, LambdaFunctionLayer, @@ -242,6 +243,45 @@ {"unit": "s"}, False, ), + ( + EventNgramLookupLayer, + [ + tf.constant( + [ + [1, 2, 3, 4, 5, 0, 0, 0], + [9, 9, 9, 9, 0, 0, 0, 0], + [0, 0, 0, 0, 1, 2, 3, 4], + ], + dtype=tf.int32, + ) + ], + { + "num_events_per_input": [2], + "top_k": 3, + "tuple_size": 4, + "lookup_keys": [[1, 2, 3, 4], [5, 0, 0, 0]], + "lookup_values": [[2, 3, 0], [4, 0, 0]], + }, + False, + ), + ( + EventNgramLookupLayer, + [ + tf.constant( + [[1, 2, 3, 4, 5, 0, 0, 0], [9, 9, 9, 9, 0, 0, 0, 0]], + dtype=tf.int32, + ) + ], + { + "num_events_per_input": [2], + "top_k": 3, + "tuple_size": 4, + "lookup_keys": [[1, 2, 3, 4], [5, 0, 0, 0]], + "lookup_values": [[2, 3, 0], [4, 0, 0]], + "token_type_lookup": [0, 0, 5, 1, 8], + }, + False, + ), (DivideLayer, [tf.random.normal((100, 10, 5))], {"divisor": 2}, False), (ExpLayer, [tf.random.normal((100, 10, 5))], None, False), (ExponentLayer, [tf.random.normal((100, 10, 5))], {"exponent": 2}, False), diff --git a/tests/kamae/spark/estimators/test_event_ngram_lookup.py b/tests/kamae/spark/estimators/test_event_ngram_lookup.py new file mode 100644 index 00000000..365acdc0 --- /dev/null +++ b/tests/kamae/spark/estimators/test_event_ngram_lookup.py @@ -0,0 +1,280 @@ +# Copyright [2024] Expedia, Inc. +# +# 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. + +import numpy as np +import pytest +import tensorflow as tf +from pyspark.sql.types import ArrayType, IntegerType, LongType, StructField, StructType + +from kamae.spark.estimators import EventNgramLookupEstimator +from kamae.spark.pipeline import KamaeSparkPipeline +from kamae.spark.transformers import EventNgramLookupTransformer + + +class TestEventNgramLookupEstimator: + @pytest.fixture + def id_df(self, spark_session): + # Two events x 4 ids per row. The tuple (1, 2, 3, 4) recurs many times so its + # n-grams clear minNgramFreq, while (9, 9, 9, 9) appears once (stays unknown). + rows = ( + [([1, 2, 3, 4, 5, 6, 7, 8],)] * 5 + + [([1, 2, 3, 4, 0, 0, 0, 0],)] * 5 + + [([9, 9, 9, 9, 9, 9, 9, 9],)] + ) + schema = StructType([StructField("clicks_ids", ArrayType(IntegerType()), True)]) + return spark_session.createDataFrame(rows, schema) + + def _estimator(self, **overrides) -> EventNgramLookupEstimator: + params = dict( + inputCols=["clicks_ids"], + outputCols=["clicks_tokens"], + numEventsPerInput=[2], + tupleSize=4, + topK=3, + vocabSize=100, + minNgramFreq=1, + ) + params.update(overrides) + return EventNgramLookupEstimator(**params) + + def test_fit_returns_transformer_with_propagated_params(self, id_df): + estimator = self._estimator(layerName="clicks_tokens") + transformer = estimator.fit(id_df) + + assert isinstance(transformer, EventNgramLookupTransformer) + assert transformer.getInputCols() == ["clicks_ids"] + assert transformer.getOutputCols() == ["clicks_tokens"] + assert transformer.getTupleSize() == 4 + assert transformer.getTopK() == 3 + assert transformer.getNumEventsPerInput() == [2] + # Vocabulary holds the two reserved tokens plus the learned n-grams. + assert transformer.getVocabularySize() > 2 + # The frequent tuple was observed and pre-encoded into the lookup table. + assert (1, 2, 3, 4) in transformer.getTupleToTokens() + + @pytest.mark.parametrize( + "param", + ["vocabularySize", "lookupKeys", "lookupValues", "tokenTypeLookup"], + ) + def test_fitted_state_is_not_an_estimator_param(self, param): + # The lookup table is produced by fitting, so only the transformer carries it. + assert not self._estimator().hasParam(param) + + def test_layer_name_defaults_to_uid(self, id_df): + estimator = self._estimator() + transformer = estimator.fit(id_df) + assert transformer.getLayerName() == estimator.uid + + def test_fitted_transformer_tokenizes_frequent_tuple(self, id_df): + transformer = self._estimator(layerName="clicks_tokens").fit(id_df) + tokens = ( + transformer.transform(id_df) + .select("clicks_tokens") + .collect()[0]["clicks_tokens"] + ) + assert len(tokens) == 2 * 3 + # The frequent tuple maps to at least one real (non pad/unk) token id (>= 2). + assert any(token >= 2 for token in tokens) + + def test_validates_input_output_lengths(self, id_df): + estimator = self._estimator(outputCols=["a", "b"]) + with pytest.raises(ValueError): + estimator.fit(id_df) + + def test_validates_events_per_input_length(self, id_df): + estimator = self._estimator(numEventsPerInput=[2, 2]) + with pytest.raises(ValueError): + estimator.fit(id_df) + + def test_sample_fraction_still_fits(self, id_df): + estimator = self._estimator(sampleFraction=0.9) + transformer = estimator.fit(id_df) + assert isinstance(transformer, EventNgramLookupTransformer) + + def test_multiple_inputs_share_one_vocabulary(self, spark_session): + rows = [([1, 2, 3, 4], [5, 6, 7, 8])] * 5 + schema = StructType( + [ + StructField("clicks_ids", ArrayType(IntegerType()), True), + StructField("prop_ids", ArrayType(IntegerType()), True), + ] + ) + df = spark_session.createDataFrame(rows, schema) + estimator = self._estimator( + inputCols=["clicks_ids", "prop_ids"], + outputCols=["clicks_tokens", "prop_tokens"], + numEventsPerInput=[1, 1], + ) + transformer = estimator.fit(df) + lookup = transformer.getTupleToTokens() + # Tuples from both input columns end up in the single shared lookup table. + assert (1, 2, 3, 4) in lookup + assert (5, 6, 7, 8) in lookup + + def test_tuple_matching_no_ngram_is_left_out_of_the_lookup_table(self, id_df): + # minNgramFreq=3 keeps the frequent tuple's n-grams but drops those of + # (9, 9, 9, 9), which occurs twice. That tuple then carries no token, so it is + # omitted from the table and takes the same table-miss path at inference as a + # tuple that was never seen while fitting. + transformer = self._estimator(layerName="clicks_tokens", minNgramFreq=3).fit( + id_df + ) + + lookup = transformer.getTupleToTokens() + assert (1, 2, 3, 4) in lookup + assert (9, 9, 9, 9) not in lookup + + def test_compatible_dtypes_are_integer_only(self): + # The Keras layer consumes integer ids only, so the Spark side must not + # advertise dtypes it cannot achieve parity on. The dtypes are compared + # against the element type of the input column, so they name the element + # types rather than the array that holds them. + estimator = self._estimator() + assert estimator.compatible_dtypes == [IntegerType(), LongType()] + + def test_use_fit_sample_defaults_to_false(self, id_df): + # KamaeSparkPipeline calls getUseFitSample() on every estimator declaring the + # param, so an unset default would raise. False is also the right value here: + # the vocabulary is thresholded on exact corpus frequencies. + estimator = self._estimator() + assert estimator.getUseFitSample() is False + + @pytest.mark.parametrize( + "param, value", + [ + ("vocabSize", 2), + ("minNgramFreq", 0), + ("tupleSize", 0), + ("topK", 0), + ("numEventsPerInput", []), + ("numEventsPerInput", [1, 0]), + ], + ) + def test_invalid_fit_params_raise(self, param, value): + with pytest.raises(ValueError): + self._estimator(**{param: value}) + + def test_include_token_types_propagates_type_lookup(self, id_df): + transformer = self._estimator( + layerName="clicks_tokens", includeTokenTypes=True + ).fit(id_df) + + assert transformer.getIncludeTokenTypes() is True + type_lookup = transformer.getTokenTypeLookup() + assert type_lookup is not None + # pad (0) and unk (1) are always type 0. + assert type_lookup[0] == 0 + assert type_lookup[1] == 0 + # Every learned token's type is a non-zero bitmask within [1, 2**tupleSize). + vocab_size = transformer.getVocabularySize() + assert len(type_lookup) == vocab_size + for token_id in range(2, vocab_size): + assert 1 <= type_lookup[token_id] < 2**4 + + def test_construct_layer_info_declares_type_outputs(self): + # The pipeline builds its fit-time graph from the estimator, so the type + # columns its fitted transformer produces must be declared here too. + info = self._estimator(includeTokenTypes=True).construct_layer_info() + assert info["outputs"] == ["clicks_tokens", "clicks_tokens_types"] + + @pytest.mark.parametrize( + "unpackable_id", + [ + # 2**15 needs 16 bits, and 4 levels x 16 bits exceeds a signed int64. + 2**15, + -3, + ], + ) + def test_fit_raises_when_the_keras_layer_could_not_pack_the_ids( + self, spark_session, unpackable_id + ): + schema = StructType([StructField("clicks_ids", ArrayType(IntegerType()), True)]) + df = spark_session.createDataFrame([([unpackable_id, 2, 3, 4],)], schema) + estimator = self._estimator(numEventsPerInput=[1]) + + with pytest.raises(ValueError) as excinfo: + estimator.fit(df) + + # BaseEstimator.fit re-raises with the params, so the cause is the context. + assert "int64" in str(excinfo.value.__context__) + + def test_pipeline_keras_model_matches_spark_with_an_estimator_on_the_types( + self, spark_session + ): + # An estimator reading a derived type column must be fitted after the + # tokenizer has been applied; a second tokenizer over the types stands in for + # any such downstream stage. + rows = ( + [([1, 2, 3, 4, 5, 6, 7, 8], [1, 2, 3, 4])] * 5 + + [([1, 2, 3, 4, 0, 0, 0, 0], [5, 6, 7, 8])] * 5 + + [([9, 9, 9, 9, 0, 0, 0, 0], [0, 0, 0, 0])] + ) + schema = StructType( + [ + StructField("clicks_ids", ArrayType(IntegerType()), True), + StructField("prop_ids", ArrayType(LongType()), True), + ] + ) + df = spark_session.createDataFrame(rows, schema) + pipeline = KamaeSparkPipeline( + stages=[ + self._estimator( + inputCols=["clicks_ids", "prop_ids"], + outputCols=["clicks_tokens", "prop_tokens"], + numEventsPerInput=[2, 1], + minNgramFreq=2, + includeTokenTypes=True, + layerName="tokenizer", + ), + self._estimator( + inputCols=["clicks_tokens_types"], + outputCols=["clicks_types_tokens"], + numEventsPerInput=[2], + tupleSize=3, + topK=2, + layerName="types_tokenizer", + ), + ] + ) + + pipeline_model = pipeline.fit(df) + keras_model = pipeline_model.build_keras_model( + input_schema=[ + {"name": "clicks_ids", "shape": (8,), "dtype": "int32"}, + {"name": "prop_ids", "shape": (4,), "dtype": "int64"}, + ] + ) + keras_outputs = keras_model( + { + "clicks_ids": tf.constant([row[0] for row in rows], dtype=tf.int32), + "prop_ids": tf.constant([row[1] for row in rows], dtype=tf.int64), + } + ) + spark_rows = pipeline_model.transform(df).collect() + + # The symbolic outputs carry the real dtype, so that integer-only layers can be + # chained onto them when the model is built. + assert {output.dtype for output in keras_model.output.values()} == {"int32"} + # clicks_tokens_types feeds the second tokenizer, so it is not a model output. + assert sorted(keras_outputs) == [ + "clicks_tokens", + "clicks_types_tokens", + "prop_tokens", + "prop_tokens_types", + ] + for output_name, keras_output in keras_outputs.items(): + np.testing.assert_array_equal( + keras_output.numpy(), + np.array([row[output_name] for row in spark_rows]), + ) diff --git a/tests/kamae/spark/transformers/test_event_ngram_lookup.py b/tests/kamae/spark/transformers/test_event_ngram_lookup.py new file mode 100644 index 00000000..355d39fa --- /dev/null +++ b/tests/kamae/spark/transformers/test_event_ngram_lookup.py @@ -0,0 +1,333 @@ +# Copyright [2024] Expedia, Inc. +# +# 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. + +import pytest +import tensorflow as tf +from pyspark.sql.types import ( + ArrayType, + IntegerType, + StringType, + StructField, + StructType, +) + +from kamae.keras.tensorflow.layers import EventNgramLookupLayer +from kamae.spark.transformers import EventNgramLookupTransformer + +# Fitted lookup table shared across the tests: two known event tuples, top_k = 3. +LOOKUP_TABLE = {(1, 2, 3, 4): [2, 3, 0], (5, 6, 7, 8): [4, 5, 0]} +# The transformer persists the table as two flat, JSON-serializable int lists. +LOOKUP_KEYS = [i for key in LOOKUP_TABLE for i in key] +LOOKUP_VALUES = [v for value in LOOKUP_TABLE.values() for v in value] +TOP_K = 3 +TUPLE_SIZE = 4 +VOCAB_SIZE = 6 + + +class TestEventNgramLookupTransformer: + def _transformer(self, **overrides) -> EventNgramLookupTransformer: + params = dict( + inputCols=["clicks_ids"], + outputCols=["clicks_tokens"], + numEventsPerInput=[2], + tupleSize=TUPLE_SIZE, + topK=TOP_K, + vocabularySize=VOCAB_SIZE, + lookupKeys=LOOKUP_KEYS, + lookupValues=LOOKUP_VALUES, + layerName="clicks_tokens", + ) + params.update(overrides) + return EventNgramLookupTransformer(**params) + + @pytest.fixture + def id_df(self, spark_session): + # Each row is a flat int array of 2 events x 4 ids: a known tuple, an unknown + # tuple and an all-zero (padding) tuple. + data = [ + ([1, 2, 3, 4, 5, 6, 7, 8],), # known, known + ([1, 2, 3, 4, 9, 9, 9, 9],), # known, unknown + ([0, 0, 0, 0, 0, 0, 0, 0],), # padding, padding + ] + schema = StructType([StructField("clicks_ids", ArrayType(IntegerType()), True)]) + return spark_session.createDataFrame(data, schema) + + def test_transform_known_unknown_padding(self, id_df): + transformer = self._transformer() + actual = [ + row["clicks_tokens"] + for row in transformer.transform(id_df).select("clicks_tokens").collect() + ] + expected = [ + [2, 3, 0, 4, 5, 0], # both tuples known + [2, 3, 0, 1, 0, 0], # known then a single (id 1), then padding + [0, 0, 0, 0, 0, 0], # both padding (id 0) + ] + assert actual == expected + + def test_spark_tf_parity_search_level(self, id_df): + transformer = self._transformer() + spark_out = [ + row["clicks_tokens"] + for row in transformer.transform(id_df).select("clicks_tokens").collect() + ] + layer = transformer.get_keras_layer() + tf_in = tf.constant( + [ + [1, 2, 3, 4, 5, 6, 7, 8], + [1, 2, 3, 4, 9, 9, 9, 9], + [0, 0, 0, 0, 0, 0, 0, 0], + ], + dtype=tf.int32, + ) + tf_out = layer(tf_in).numpy().tolist() + assert spark_out == tf_out + + @pytest.mark.parametrize( + "ids, expected", + [ + # A short row is right-padded up to numEvents * topK. + ([1, 2, 3, 4], [2, 3, 0, 0, 0, 0]), + # A row longer than numEvents is truncated. + ([1, 2, 3, 4, 5, 6, 7, 8, 1, 2, 3, 4], [2, 3, 0, 4, 5, 0]), + # Null and empty rows are all padding. + (None, [0, 0, 0, 0, 0, 0]), + ([], [0, 0, 0, 0, 0, 0]), + ], + ) + def test_transform_pads_and_truncates_to_fixed_length( + self, spark_session, ids, expected + ): + # given + schema = StructType([StructField("clicks_ids", ArrayType(IntegerType()), True)]) + df = spark_session.createDataFrame([(ids,)], schema) + + # when + actual = ( + self._transformer() + .transform(df) + .select("clicks_tokens") + .collect()[0]["clicks_tokens"] + ) + + # then + assert actual == expected + + def test_transform_rejects_string_input_column(self, spark_session): + # The Keras layer can only consume integer ids, so the Spark side declares the + # same dtypes rather than silently accepting a column it cannot serve. + schema = StructType([StructField("clicks_ids", ArrayType(StringType()), True)]) + df = spark_session.createDataFrame([(["1,2,3,4"],)], schema) + with pytest.raises(TypeError): + self._transformer().transform(df).collect() + + def test_transform_raises_when_column_lengths_differ(self, id_df): + # Columns are tokenized pairwise, so a mismatch must not silently drop one. + transformer = self._transformer(outputCols=["clicks_tokens", "extra_tokens"]) + with pytest.raises(ValueError): + transformer.transform(id_df) + + def test_spark_tf_parity_with_list_dimension(self, id_df): + # Rank-3 inputs (batch, list_size, ids) arise in listwise models; the list + # dimension must be preserved and each item tokenized independently. + transformer = self._transformer() + spark_out = [ + row["clicks_tokens"] + for row in transformer.transform(id_df).select("clicks_tokens").collect() + ] + layer = transformer.get_keras_layer() + # One batch of 3 list items, mirroring the 3 Spark rows. + tf_in = tf.constant( + [ + [ + [1, 2, 3, 4, 5, 6, 7, 8], + [1, 2, 3, 4, 9, 9, 9, 9], + [0, 0, 0, 0, 0, 0, 0, 0], + ] + ], + dtype=tf.int32, + ) + tf_out = layer(tf_in) + assert tf_out.shape == (1, 3, 2 * TOP_K) + assert tf_out.numpy()[0].tolist() == spark_out + + def test_get_keras_layer_returns_single_consolidated_layer(self): + layer = self._transformer().get_keras_layer() + assert isinstance(layer, EventNgramLookupLayer) + assert layer.name == "clicks_tokens" + assert layer.num_events_per_input == [2] + assert layer.top_k == TOP_K + assert layer.tuple_size == TUPLE_SIZE + + def test_spark_tf_parity_multiple_inputs_with_types(self, spark_session): + # A single layer tokenizes every input against the shared table, returning all + # the token tensors and then all the type tensors, in the declared order. + transformer = self._transformer( + inputCols=["clicks_ids", "prop_ids"], + outputCols=["clicks_tokens", "prop_tokens"], + numEventsPerInput=[2, 1], + includeTokenTypes=True, + tokenTypeLookup=[0, 0, 5, 1, 8, 3], + layerName="tokenizer", + ) + clicks = [[1, 2, 3, 4, 5, 6, 7, 8], [1, 2, 3, 4, 9, 9, 9, 9], [0] * 8] + prop = [[5, 6, 7, 8], [9, 9, 9, 9], [1, 2, 3, 4]] + schema = StructType( + [ + StructField("clicks_ids", ArrayType(IntegerType()), True), + StructField("prop_ids", ArrayType(IntegerType()), True), + ] + ) + df = spark_session.createDataFrame(list(zip(clicks, prop)), schema) + _, output_cols = transformer.get_layer_inputs_outputs() + + spark_rows = transformer.transform(df).select(*output_cols).collect() + tf_outputs = transformer.get_keras_layer()( + [tf.constant(clicks, dtype=tf.int32), tf.constant(prop, dtype=tf.int32)] + ) + + assert output_cols == [ + "clicks_tokens", + "prop_tokens", + "clicks_tokens_types", + "prop_tokens_types", + ] + assert len(tf_outputs) == len(output_cols) + for output_col, tf_output in zip(output_cols, tf_outputs): + assert tf_output.numpy().tolist() == [row[output_col] for row in spark_rows] + + def test_construct_layer_info(self): + info = self._transformer().construct_layer_info() + assert info["name"] == "clicks_tokens" + assert info["inputs"] == ["clicks_ids"] + assert info["outputs"] == ["clicks_tokens"] + + def test_transform_defaults(self): + transformer = EventNgramLookupTransformer( + inputCols=["clicks_ids"], + outputCols=["clicks_tokens"], + numEventsPerInput=[2], + topK=TOP_K, + lookupKeys=LOOKUP_KEYS, + lookupValues=LOOKUP_VALUES, + ) + assert transformer.getLayerName() == transformer.uid + assert transformer.getTupleSize() == 4 + assert transformer.getIncludeTokenTypes() is False + assert transformer.getTokenTypeLookup() is None + + def test_getters(self): + transformer = self._transformer() + assert transformer.getTupleSize() == TUPLE_SIZE + assert transformer.getTopK() == TOP_K + assert transformer.getNumEventsPerInput() == [2] + assert transformer.getVocabularySize() == VOCAB_SIZE + assert transformer.getTupleToTokens() == LOOKUP_TABLE + assert transformer.getIncludeTokenTypes() is False + + def test_params_survive_save_and_load(self, spark_session, tmp_path): + # Spark writes params to JSON metadata, so the fitted table has to stay in a + # JSON-serializable shape for the transformer, and any pipeline holding it, to + # be saveable at all. + transformer = self._transformer( + includeTokenTypes=True, tokenTypeLookup=[0, 0, 5, 1, 8, 3] + ) + path = str(tmp_path / "transformer") + transformer.write().overwrite().save(path) + + reloaded = EventNgramLookupTransformer.load(path) + assert reloaded.getLookupKeys() == LOOKUP_KEYS + assert reloaded.getLookupValues() == LOOKUP_VALUES + assert reloaded.getTupleToTokens() == LOOKUP_TABLE + assert reloaded.getTokenTypeLookup() == [0, 0, 5, 1, 8, 3] + + def test_reloaded_transformer_tokenizes_identically( + self, spark_session, id_df, tmp_path + ): + transformer = self._transformer() + expected = transformer.transform(id_df).collect() + + path = str(tmp_path / "transformer") + transformer.write().overwrite().save(path) + reloaded = EventNgramLookupTransformer.load(path) + + assert reloaded.transform(id_df).collect() == expected + + def test_transform_emits_type_columns(self, id_df): + # token id -> ID-level bitmask; tokens 2/3 are types 5/1, unk(1)->0, pad(0)->0. + type_lookup = [0, 0, 5, 1, 8, 3] + transformer = self._transformer( + includeTokenTypes=True, tokenTypeLookup=type_lookup + ) + rows = ( + transformer.transform(id_df) + .select("clicks_tokens", "clicks_tokens_types") + .collect() + ) + for row in rows: + # Each type is the bitmask gathered at the corresponding token id. + assert row["clicks_tokens_types"] == [ + type_lookup[t] for t in row["clicks_tokens"] + ] + + def test_output_dtype_applies_to_the_type_columns(self, id_df): + # The type columns are derived rather than listed in outputCols, so they must + # be cast alongside them to stay the same dtype as the layer's type tensors. + transformer = self._transformer( + includeTokenTypes=True, + tokenTypeLookup=[0, 0, 5, 1, 8, 3], + outputDtype="float", + ) + + schema = transformer.transform(id_df).schema + + assert ( + schema["clicks_tokens"].dataType == schema["clicks_tokens_types"].dataType + ) + + def test_construct_layer_info_appends_type_outputs(self): + # The consolidated layer returns all token tensors then all type tensors, so + # the declared output columns must follow the same grouping. + info = self._transformer( + includeTokenTypes=True, tokenTypeLookup=[0, 0, 5, 1, 8, 3] + ).construct_layer_info() + assert info["inputs"] == ["clicks_ids"] + assert info["outputs"] == ["clicks_tokens", "clicks_tokens_types"] + + def test_spark_tf_parity_with_types(self, id_df): + type_lookup = [0, 0, 5, 1, 8, 3] + transformer = self._transformer( + includeTokenTypes=True, tokenTypeLookup=type_lookup + ) + spark_tokens, spark_types = [], [] + for row in ( + transformer.transform(id_df) + .select("clicks_tokens", "clicks_tokens_types") + .collect() + ): + spark_tokens.append(row["clicks_tokens"]) + spark_types.append(row["clicks_tokens_types"]) + + layer = transformer.get_keras_layer() + tf_in = tf.constant( + [ + [1, 2, 3, 4, 5, 6, 7, 8], + [1, 2, 3, 4, 9, 9, 9, 9], + [0, 0, 0, 0, 0, 0, 0, 0], + ], + dtype=tf.int32, + ) + tf_tokens, tf_types = layer(tf_in) + assert tf_tokens.numpy().tolist() == spark_tokens + assert tf_types.numpy().tolist() == spark_types diff --git a/tests/kamae/spark/utils/test_ngram_utils.py b/tests/kamae/spark/utils/test_ngram_utils.py new file mode 100644 index 00000000..f0fc572e --- /dev/null +++ b/tests/kamae/spark/utils/test_ngram_utils.py @@ -0,0 +1,318 @@ +# Copyright [2024] Expedia, Inc. +# +# 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. + +import logging +from collections import Counter + +import pytest +from pyspark.sql.types import ArrayType, IntegerType, StructField, StructType + +from kamae.keras.tensorflow.layers.event_ngram_lookup import PAD_TOKEN_ID, UNK_TOKEN_ID +from kamae.spark.utils.ngram_utils import ( + EventNgramVocabulary, + build_tuple_lookup_table, + build_vocabulary, + collect_ngrams_from_dataframe, + log_vocabulary_by_length, + tokenize_events, + tokenize_events_with_types, +) +from kamae.spark.utils.ngram_worker_functions import ( + encode_tuple, + extract_ngrams_from_column_worker, + extract_tuples_from_column_worker, +) + + +class TestCollectNgramsFromDataFrame: + @pytest.fixture(scope="class") + def id_df(self, spark_session): + schema = StructType([StructField("clicks_ids", ArrayType(IntegerType()), True)]) + return spark_session.createDataFrame([([1, 2, 3, 4],)], schema) + + def test_raises_when_an_input_column_is_missing(self, id_df): + # Silently training on whichever columns happen to exist would produce a + # vocabulary that does not cover every column being tokenized. + with pytest.raises(ValueError, match="not found on the DataFrame"): + collect_ngrams_from_dataframe(id_df, ["clicks_ids", "absent_col"], 4) + + def test_counts_across_all_given_columns(self, id_df): + counter = collect_ngrams_from_dataframe(id_df, ["clicks_ids"], 4) + # One event of 4 non-zero levels -> 2**4 - 1 = 15 combinations, each seen once. + assert len(counter) == 15 + assert set(counter.values()) == {1} + + def test_drops_ngrams_below_min_frequency_before_collecting(self, spark_session): + # The two events share L0_1, L1_2 and L2_3, so only the 2**3 - 1 = 7 + # combinations of those levels are seen twice. + schema = StructType([StructField("clicks_ids", ArrayType(IntegerType()), True)]) + df = spark_session.createDataFrame([([1, 2, 3, 4],), ([1, 2, 3, 5],)], schema) + + counter = collect_ngrams_from_dataframe(df, ["clicks_ids"], 4, min_ngram_freq=2) + + assert len(counter) == 7 + assert set(counter.values()) == {2} + assert all("L3_4" not in ngram and "L3_5" not in ngram for ngram in counter) + + def test_raises_when_an_input_column_is_a_nested_array(self, spark_session): + # A nested array would be chunked over its sub-arrays rather than over its ids. + schema = StructType( + [StructField("nested_ids", ArrayType(ArrayType(IntegerType())), True)] + ) + nested_df = spark_session.createDataFrame([([[1, 2, 3, 4]],)], schema) + with pytest.raises(ValueError, match="array nesting level"): + collect_ngrams_from_dataframe(nested_df, ["nested_ids"], 4) + + +class TestExtractNgramsFromColumnWorker: + @pytest.mark.parametrize( + "col_value, event_size, expected", + [ + # Every combination of the present levels, so level-skipping n-grams + # such as (L0, L2) are included. + ( + [1, 2, 0], + 3, + [("L0_1",), ("L1_2",), ("L0_1", "L1_2")], + ), + # Zero ids are dropped before combining, so L1 never appears. + ( + [1, 0, 3], + 3, + [("L0_1",), ("L2_3",), ("L0_1", "L2_3")], + ), + # An all-zero event contributes nothing. + ([0, 0, 0], 3, []), + (None, 3, []), + ([], 3, []), + # A null id is an absent ID level, exactly like a 0, rather than becoming + # an "L1_None" token that could never match the vocabulary. + ( + [1, None, 3], + 3, + [("L0_1",), ("L2_3",), ("L0_1", "L2_3")], + ), + ([None, None, None], 3, []), + ], + ) + def test_extracts_within_event_combinations(self, col_value, event_size, expected): + actual = extract_ngrams_from_column_worker(col_value, event_size) + assert sorted(actual) == sorted(expected) + + def test_never_crosses_event_boundaries(self): + # Two single-level events: no n-gram may contain ids from both. + ngrams = extract_ngrams_from_column_worker([1, 0, 0, 5, 0, 0], 3) + assert all( + not ({"L0_1"} <= set(ngram) and {"L0_5"} <= set(ngram)) for ngram in ngrams + ) + + def test_raises_when_the_ids_are_not_a_whole_number_of_events(self): + with pytest.raises(ValueError, match="whole number of events"): + extract_ngrams_from_column_worker([1, 2, 3, 4, 5, 6], 4) + + +class TestExtractTuplesFromColumnWorker: + @pytest.mark.parametrize( + "col_value, expected", + [ + ([1, 2, 3, 4, 5, 6, 7, 8], [(1, 2, 3, 4), (5, 6, 7, 8)]), + # All-zero events are padding and are not vocabulary entries. + ([1, 2, 3, 4, 0, 0, 0, 0], [(1, 2, 3, 4)]), + # A partially-zero event is still a real event. + ([1, 0, 0, 0], [(1, 0, 0, 0)]), + ([0, 0, 0, 0], []), + (None, []), + ], + ) + def test_extracts_non_zero_tuples(self, col_value, expected): + assert extract_tuples_from_column_worker(col_value, 4) == expected + + +class TestBuildVocabulary: + def test_assigns_ids_by_descending_frequency_from_two(self): + counter = Counter({("L0_1",): 100, ("L1_2",): 50, ("L0_1", "L1_2"): 10}) + ngram_to_id = build_vocabulary(counter, vocab_size=100) + # Ids start at 2 (0 and 1 are reserved) and ascend as frequency descends, so + # a smaller id always means a more frequent n-gram. + assert ngram_to_id == { + ("L0_1",): 2, + ("L1_2",): 3, + ("L0_1", "L1_2"): 4, + } + + def test_vocab_size_reserves_two_special_tokens(self): + counter = Counter({("L0_1",): 100, ("L1_2",): 90, ("L2_3",): 80}) + ngram_to_id = build_vocabulary(counter, vocab_size=4) + # vocab_size=4 leaves room for 2 learned n-grams alongside pad and unk. + assert len(ngram_to_id) == 2 + + def test_ties_broken_deterministically_by_ngram(self): + counter = Counter({("L1_2",): 10, ("L0_1",): 10}) + first = build_vocabulary(counter, vocab_size=100) + second = build_vocabulary( + Counter({("L0_1",): 10, ("L1_2",): 10}), vocab_size=100 + ) + assert first == second + + +class TestEventNgramVocabularyBuildTypeLookup: + def _vocabulary(self, ngrams): + return EventNgramVocabulary(ngrams=ngrams) + + @pytest.mark.parametrize( + "ngram, expected_bitmask", + [ + (("L0_1",), 0b0001), + (("L1_2",), 0b0010), + (("L3_4",), 0b1000), + # A level-skipping n-gram sets exactly the bits of the levels it spans. + (("L0_1", "L2_3"), 0b0101), + (("L0_1", "L1_2", "L2_3", "L3_4"), 0b1111), + ], + ) + def test_bitmask_names_the_spanned_levels(self, ngram, expected_bitmask): + vocabulary = self._vocabulary({ngram: 2}) + assert vocabulary.build_type_lookup()[2] == expected_bitmask + + def test_special_tokens_have_type_zero(self): + lookup = self._vocabulary({("L0_1",): 2}).build_type_lookup() + assert lookup[PAD_TOKEN_ID] == 0 + assert lookup[UNK_TOKEN_ID] == 0 + + def test_supports_more_than_four_levels(self): + # The bitmask is parsed from the level integer, so it is not capped at L3. + vocabulary = self._vocabulary({("L4_9", "L6_1"): 2}) + assert vocabulary.build_type_lookup()[2] == (1 << 4) | (1 << 6) + + +class TestEncodeTuple: + # ("L0_1",) is the most frequent, then the pair, then ("L1_2",). + NGRAMS = {("L0_1",): 2, ("L0_1", "L1_2"): 3, ("L1_2",): 4} + + def test_returns_matching_ngrams_most_frequent_first(self): + assert encode_tuple((1, 2), self.NGRAMS, top_k=3) == [2, 3, 4] + + def test_truncates_to_top_k(self): + assert encode_tuple((1, 2), self.NGRAMS, top_k=2) == [2, 3] + + def test_pads_when_fewer_matches_than_top_k(self): + assert encode_tuple((1, 0), self.NGRAMS, top_k=3) == [ + 2, + PAD_TOKEN_ID, + PAD_TOKEN_ID, + ] + + def test_all_zero_tuple_has_no_encoding(self): + assert encode_tuple((0, 0), self.NGRAMS, top_k=3) is None + + def test_unmatched_tuple_has_no_encoding_so_it_is_left_out_of_the_table(self): + # A tuple with no n-gram in the vocabulary carries no token, so it is omitted + # from the lookup table and resolves to at inference by the same rule as + # a tuple that was never seen while fitting. + assert encode_tuple((7, 8), self.NGRAMS, top_k=4) is None + assert tokenize_events([7, 8], {}, num_events=1, tuple_size=2, top_k=4) == [ + UNK_TOKEN_ID, + PAD_TOKEN_ID, + PAD_TOKEN_ID, + PAD_TOKEN_ID, + ] + + +class TestTokenizeEvents: + LOOKUP = {(1, 2, 3, 4): [2, 3], (5, 6, 7, 8): [4, 5]} + + @pytest.mark.parametrize( + "ids, expected", + [ + ([1, 2, 3, 4, 5, 6, 7, 8], [2, 3, 4, 5]), + # A non-zero tuple missing from the table is unknown, not padding, and is + # marked by a single rather than filling every token slot. + ([1, 2, 3, 4, 9, 9, 9, 9], [2, 3, UNK_TOKEN_ID, PAD_TOKEN_ID]), + # An all-zero event is padding. + ([1, 2, 3, 4, 0, 0, 0, 0], [2, 3, 0, 0]), + # Short rows are right-padded to num_events * top_k. + ([1, 2, 3, 4], [2, 3, 0, 0]), + # Long rows are truncated. + ([1, 2, 3, 4, 5, 6, 7, 8, 1, 2, 3, 4], [2, 3, 4, 5]), + (None, [0, 0, 0, 0]), + ([], [0, 0, 0, 0]), + # Null ids read as 0, so an all-null event is padding and a partially-null + # event is looked up as if the null levels were absent. + ([None, None, None, None, 1, 2, 3, 4], [0, 0, 2, 3]), + ([1, 2, 3, None, 0, 0, 0, 0], [UNK_TOKEN_ID, PAD_TOKEN_ID, 0, 0]), + ], + ) + def test_tokenizes_to_fixed_length(self, ids, expected): + actual = tokenize_events(ids, self.LOOKUP, num_events=2, tuple_size=4, top_k=2) + assert actual == expected + assert len(actual) == 4 + + def test_raises_when_the_ids_are_not_a_whole_number_of_events(self): + with pytest.raises(ValueError, match="whole number of events"): + tokenize_events( + [1, 2, 3, 4, 5, 6], self.LOOKUP, num_events=2, tuple_size=4, top_k=2 + ) + + def test_empty_lookup_table_maps_every_event_to_unk(self): + assert tokenize_events( + [1, 2, 3, 4], {}, num_events=1, tuple_size=4, top_k=2 + ) == [UNK_TOKEN_ID, PAD_TOKEN_ID] + + def test_with_types_gathers_each_tokens_type(self): + # Type lookup indexed by token id: pad and unk are 0, tokens 2..5 are learned. + type_lookup = [0, 0, 1, 3, 2, 15] + tokens, types = tokenize_events_with_types( + [1, 2, 3, 4, 9, 9, 9, 9], + self.LOOKUP, + type_lookup, + num_events=2, + tuple_size=4, + top_k=2, + ) + assert tokens == [2, 3, UNK_TOKEN_ID, PAD_TOKEN_ID] + assert types == [1, 3, 0, 0] + + +class TestLogVocabularyByLength: + def test_reports_the_share_of_each_n_gram_length(self, caplog): + # A vocabulary dominated by 1-grams means the n-gram combinations are + # earning little, which is what this reports. + vocabulary = { + ("L0_1",): 2, + ("L0_2",): 3, + ("L0_1", "L1_5"): 4, + ("L0_1", "L1_5", "L2_9"): 5, + ("L0_1", "L1_5", "L2_9", "L3_2"): 6, + } + with caplog.at_level(logging.INFO, logger="kamae.spark.utils.ngram_utils"): + log_vocabulary_by_length(vocabulary) + + message = caplog.text + assert "1-gram 2 (40.0%)" in message + assert "2-gram 1 (20.0%)" in message + assert "3-gram 1 (20.0%)" in message + assert "4-gram 1 (20.0%)" in message + + def test_reports_an_empty_vocabulary_without_dividing_by_zero(self, caplog): + with caplog.at_level(logging.INFO, logger="kamae.spark.utils.ngram_utils"): + log_vocabulary_by_length({}) + assert "empty vocabulary" in caplog.text + + def test_build_vocabulary_reports_the_distribution(self, caplog): + counter = Counter( + {("L0_1",): 10, ("L0_1", "L1_2"): 8, ("L0_1", "L1_2", "L2_3"): 6} + ) + with caplog.at_level(logging.INFO, logger="kamae.spark.utils.ngram_utils"): + build_vocabulary(ngram_counter=counter, vocab_size=100) + assert "Vocabulary by n-gram length" in caplog.text