From 64bf5c553d0de54edc125a3bfed88ebe48c72237 Mon Sep 17 00:00:00 2001 From: Sbranikas Date: Sat, 19 Sep 2026 00:32:58 +0200 Subject: [PATCH 1/6] feat: add event n-gram tokenizer with a frequency-based shared vocabulary Co-Authored-By: Claude Opus 5 (1M context) --- README.md | 1 + src/kamae/keras/tensorflow/layers/__init__.py | 2 + .../tensorflow/layers/event_ngram_lookup.py | 332 ++++++++++++++++++ src/kamae/spark/estimators/__init__.py | 1 + .../spark/estimators/event_ngram_lookup.py | 290 +++++++++++++++ src/kamae/spark/params/__init__.py | 1 + src/kamae/spark/params/shared.py | 234 +++++++++++- src/kamae/spark/transformers/__init__.py | 1 + .../spark/transformers/event_ngram_lookup.py | 265 ++++++++++++++ src/kamae/spark/utils/ngram_utils.py | 299 ++++++++++++++++ .../spark/utils/ngram_worker_functions.py | 132 +++++++ .../layers/test_event_ngram_lookup.py | 229 ++++++++++++ tests/kamae/keras/test_jit_compatibility.py | 12 + tests/kamae/keras/test_layer_serialisation.py | 40 +++ .../estimators/test_event_ngram_lookup.py | 153 ++++++++ .../transformers/test_event_ngram_lookup.py | 300 ++++++++++++++++ tests/kamae/spark/utils/test_ngram_utils.py | 251 +++++++++++++ 17 files changed, 2542 insertions(+), 1 deletion(-) create mode 100644 src/kamae/keras/tensorflow/layers/event_ngram_lookup.py create mode 100644 src/kamae/spark/estimators/event_ngram_lookup.py create mode 100644 src/kamae/spark/transformers/event_ngram_lookup.py create mode 100644 src/kamae/spark/utils/ngram_utils.py create mode 100644 src/kamae/spark/utils/ngram_worker_functions.py create mode 100644 tests/kamae/keras/tensorflow/layers/test_event_ngram_lookup.py create mode 100644 tests/kamae/spark/estimators/test_event_ngram_lookup.py create mode 100644 tests/kamae/spark/transformers/test_event_ngram_lookup.py create mode 100644 tests/kamae/spark/utils/test_ngram_utils.py diff --git a/README.md b/README.md index 88ad70ba..1b636726 100644 --- a/README.md +++ b/README.md @@ -115,6 +115,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 | Tokenises 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..70b58ffc --- /dev/null +++ b/src/kamae/keras/tensorflow/layers/event_ngram_lookup.py @@ -0,0 +1,332 @@ +# 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, applied fully +vectorised (no per-element ``map_fn`` / ``cond`` and 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 from a ``.keras`` file and +exports cleanly to SavedModel. + +Token id conventions (shared with the vocabulary): + - ``0`` = padding (````): emitted for all-zero / missing events. + - ``1`` = unknown (````): emitted for a non-zero tuple absent from the table. +""" + +from typing import Any, Dict, 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 + + +@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 (e.g. + per listwise item) is preserved. The id axis is padded/truncated to + ``num_events * tuple_size`` before being split into tuples. + + 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 bijectively packed into a single ``int64`` by allotting + ``key_bits`` bits per ID level, where ``key_bits`` is the bit length of the + largest ID in the table. 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 the packed key would not fit in a signed ``int64``. + """ + max_id = max((int(x) for k in keys for x in k), default=1) + self.key_bits = max(1, max_id.bit_length()) + if self.tuple_size * self.key_bits > 63: + raise ValueError( + f"Cannot pack event tuples into a signed int64: tuple_size " + f"({self.tuple_size}) * key_bits ({self.key_bits}, from a largest ID " + f"of {max_id}) = {self.tuple_size * self.key_bits} bits, which exceeds " + f"63. Reduce tuple_size or the ID cardinality." + ) + 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] * self.top_k, 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. + :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 == 3: + batch_size = tf.shape(inputs)[0] + list_or_1 = tf.shape(inputs)[1] + num_ids = tf.shape(inputs)[2] + inputs_flat = tf.reshape(inputs, [-1, num_ids]) + eff_batch = batch_size * list_or_1 + restore_list = True + else: + inputs_flat = inputs + eff_batch = tf.shape(inputs_flat)[0] + restore_list = False + + # Pad/truncate the id axis to num_events * tuple_size, then split into tuples. + # Padding by the full expected length before slicing covers both the short and + # the long case without branching on the input width. + 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 wider than key_bits 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 <= 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, so the (top_k,) unk/pad constants apply as they are, + # with no per-tuple copy materialised. + 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_or_1, output_length]) + return tf.reshape(tokens, [eff_batch, 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. + :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] + + 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 compute_output_shape( + self, input_shape: Union[tuple, List[tuple]] + ) -> Union[tuple, List[tuple]]: + """ + Declares the output shape(s) for functional model building. + + :param input_shape: A single input shape, or a list of input shapes (one per + input column). + :returns: The matching output shape(s), with each id axis replaced by + ``num_events * top_k``. When types are emitted, the per-input type shapes are + appended after the token shapes. + """ + single = not isinstance(input_shape[0], (list, tuple)) + shapes = [input_shape] if single else list(input_shape) + + token_shapes = [] + for shape, num_events in zip(shapes, self.num_events_per_input): + output_length = num_events * self.top_k + if len(shape) == 3: + token_shapes.append((shape[0], shape[1], output_length)) + elif len(shape) == 2: + token_shapes.append((shape[0], output_length)) + else: + token_shapes.append(shape) + + if self._token_type_lookup is not None: + outputs = token_shapes + list(token_shapes) + elif single: + return token_shapes[0] + else: + outputs = token_shapes + return outputs + + 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..bf51281a --- /dev/null +++ b/src/kamae/spark/estimators/event_ngram_lookup.py @@ -0,0 +1,290 @@ +# 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. +""" + +# pylint: disable=unused-argument +# pylint: disable=invalid-name +# pylint: disable=too-many-ancestors +# pylint: disable=no-member +from typing import List, Optional + +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 + +from kamae.keras.core.backend import TENSORFLOW_ONLY +from kamae.spark.params import ( + EventNgramLookupParams, + MultiInputMultiOutputParams, + SampleFractionParams, +) +from kamae.spark.transformers import EventNgramLookupTransformer +from kamae.spark.utils.ngram_utils import ( + EventNgramVocabulary, + build_tuple_lookup_table, + build_vocabulary, + collect_ngrams_from_dataframe, +) + +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(), ArrayType(IntegerType())] + + 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. + """ + input_cols = self.getInputCols() + output_cols = self.getOutputCols() + num_events_per_input = self.getNumEventsPerInput() + + if len(input_cols) != len(output_cols): + raise ValueError( + f"inputCols and outputCols must have the same length. Got " + f"{len(input_cols)} inputs and {len(output_cols)} outputs." + ) + if num_events_per_input is None or len(input_cols) != 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_cols)} inputs and {n_events} event counts." + ) + + 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, + ) + vocabulary = EventNgramVocabulary( + ngrams=build_vocabulary( + ngram_counter=ngram_counter, + vocab_size=self.getVocabSize(), + min_ngram_freq=self.getMinNgramFreq(), + ) + ) + + # 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 that it is JSON-serialisable + # and the fitted pipeline can be saved. One pass over the items keeps the keys + # and values positionally aligned. + 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) + + 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..c3418c2f 100644 --- a/src/kamae/spark/params/shared.py +++ b/src/kamae/spark/params/shared.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import List, Union +from typing import Dict, List, Optional, Tuple, Union from pyspark.ml.param import Param, Params, TypeConverters @@ -1038,3 +1038,235 @@ 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, + ) + + vocabularySize = Param( + Params._dummy(), + "vocabularySize", + "Number of token ids the fitted vocabulary actually produced, including the " + "reserved pad and unk tokens. This is the fitted result, not the requested " + "target (which is the estimator's vocabSize and is only an upper bound): it is " + "published so a downstream embedding table can be sized to the tokens that " + "were really learned.", + typeConverter=TypeConverters.toInt, + ) + + # The fitted lookup table is held as two parallel flat int arrays rather than as a + # dict, because Spark ML writes params to JSON metadata and a dict keyed by id + # tuples is not JSON-serialisable, which would make the fitted pipeline unsaveable. + # Same parallel-arrays shape as StringMapParams. + 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, + ) + + includeTokenTypes = Param( + Params._dummy(), + "includeTokenTypes", + "Whether to also emit, per input column, a parallel '_types' column giving " + "each token's ID-level bitmask (a compact categorical feature for the model).", + typeConverter=TypeConverters.toBoolean, + ) + + 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 setNumEventsPerInput(self, value: List[int]) -> "EventNgramLookupParams": + """ + Sets the numEventsPerInput parameter. + + :param value: Number of events per input column. + :returns: Instance of class mixed in. + """ + 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. + :returns: Instance of class mixed in. + """ + 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. + :returns: Instance of class mixed in. + """ + 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 setVocabularySize(self, value: int) -> "EventNgramLookupParams": + """ + Sets the vocabularySize parameter. + + :param value: Total number of unique tokens in the fitted vocabulary. + :returns: Instance of class mixed in. + """ + return self._set(vocabularySize=value) + + def getVocabularySize(self) -> int: + """ + Gets the vocabularySize parameter. + + :returns: Total number of unique tokens in the fitted vocabulary. + """ + return self.getOrDefault(self.vocabularySize) + + def setLookupKeys(self, value: List[int]) -> "EventNgramLookupParams": + """ + 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]) -> "EventNgramLookupParams": + """ + 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 getTupleToTokens(self) -> Dict[Tuple[int, ...], List[int]]: + """ + Rebuilds the ``event tuple -> top-k token list`` lookup table. + + Derived from the flat ``lookupKeys`` / ``lookupValues`` params, which are what + the fitted pipeline actually persists. + + :returns: Mapping from each event tuple to its list of token ids. + """ + keys = self.getLookupKeys() + values = self.getLookupValues() + if not keys or not values: + return {} + 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) + ) + } + + 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 setTokenTypeLookup(self, value: List[int]) -> "EventNgramLookupParams": + """ + 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) diff --git a/src/kamae/spark/transformers/__init__.py b/src/kamae/spark/transformers/__init__.py index de563832..1943127a 100644 --- a/src/kamae/spark/transformers/__init__.py +++ b/src/kamae/spark/transformers/__init__.py @@ -37,6 +37,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..bda25ed9 --- /dev/null +++ b/src/kamae/spark/transformers/event_ngram_lookup.py @@ -0,0 +1,265 @@ +# 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``. +Both the Spark path (``_transform``) and the TensorFlow path (``get_keras_layer``) are +plain ``O(1)`` lookups, so the same event tuple always maps to the same tokens and the +two paths produce identical output. + +When ``includeTokenTypes`` is set, each token column ```` is accompanied by a +parallel ``_types`` column giving each token's ID-level bitmask (a compact +categorical feature for the model). The type is a pure function of the token id, so it +is derived by gathering the fitted ``tokenTypeLookup`` at the token ids. + +Token id conventions (shared with the vocabulary and the TF layer): + - ``0`` = padding (````): emitted for all-zero / missing events. + - ``1`` = unknown (````): emitted for a non-zero tuple absent from the table. +""" + +# pylint: disable=unused-argument +# pylint: disable=invalid-name +# pylint: disable=too-many-ancestors +# pylint: disable=no-member +from typing import Any, Dict, List, Optional, Tuple + +import pyspark.sql.functions as F +import tensorflow as tf +from pyspark import keyword_only +from pyspark.sql import DataFrame +from pyspark.sql.types import ArrayType, DataType, IntegerType, 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.ngram_utils import tokenize_events + +from .base import BaseTransformer + + +class EventNgramLookupTransformer( + BaseTransformer, + MultiInputMultiOutputParams, + EventNgramLookupParams, +): + """ + 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 (padding for all-zero events, ```` for non-zero tuples absent from + the table), 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. + """ + + 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: Total number of unique tokens 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(), ArrayType(IntegerType())] + + def _transform(self, dataset: DataFrame) -> DataFrame: + """ + Tokenizes each input column with a per-column UDF over ``tokenize_events``. + + 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, a ``_types`` + column is added by gathering the type bitmask at each token id. + + :param dataset: Input DataFrame. + :returns: DataFrame with the tokenized output columns, and the optional type + columns. + """ + tuple_size = self.getTupleSize() + top_k = self.getTopK() + lookup_table = self.getTupleToTokens() + include_types = self.getIncludeTokenTypes() + type_lookup = self.getTokenTypeLookup() if include_types else None + + token_array_type = ArrayType(IntegerType()) + for input_col, output_col, num_events in zip( + self.getInputCols(), + self.getOutputCols(), + self.getNumEventsPerInput(), + ): + # num_events is bound per column via the default argument so each UDF + # captures its own value rather than the last loop iteration's. + if not include_types: + tokenize_udf = F.udf( + lambda ids, num_events=num_events: tokenize_events( + ids, lookup_table, num_events, tuple_size, top_k + ), + token_array_type, + ) + dataset = dataset.withColumn(output_col, tokenize_udf(F.col(input_col))) + continue + + # A token's type is a pure function of its id, so both are produced by one + # UDF returning a struct, keeping the column to a single Python round-trip. + def _tokenize_with_types( + ids: Optional[List[int]], num_events: int = num_events + ) -> Tuple[List[int], List[int]]: + """Tokenizes one row and gathers each token's type bitmask. + + :param ids: One row's discrete ID values. + :param num_events: Number of events this column is sized to. + :returns: Tuple of (token ids, type bitmasks), same length. + """ + tokens = tokenize_events( + ids, lookup_table, num_events, tuple_size, top_k + ) + return tokens, [type_lookup[token] for token in tokens] + + tokenize_udf = F.udf( + _tokenize_with_types, + StructType( + [ + StructField("tokens", token_array_type), + StructField("types", token_array_type), + ] + ), + ) + struct_col = f"{output_col}__tokens_and_types" + dataset = ( + dataset.withColumn(struct_col, tokenize_udf(F.col(input_col))) + .withColumn(output_col, F.col(f"{struct_col}.tokens")) + .withColumn(f"{output_col}_types", F.col(f"{struct_col}.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), for in-graph tokenization on + TensorFlow (needed when reading raw files via TFParquet). The layer outputs one + token tensor per input, plus one type tensor per input when + ``includeTokenTypes`` is set. + + :returns: The consolidated ``EventNgramLookupLayer``. + """ + # The layer takes the table as nested key/value lists, so the flat params are + # reshaped back into one list per tuple. + tuple_size = self.getTupleSize() + top_k = self.getTopK() + keys = self.getLookupKeys() or [] + values = self.getLookupValues() or [] + return EventNgramLookupLayer( + num_events_per_input=self.getNumEventsPerInput(), + top_k=top_k, + tuple_size=tuple_size, + lookup_keys=[ + keys[i : i + tuple_size] for i in range(0, len(keys), tuple_size) + ], + lookup_values=[values[i : i + top_k] for i in range(0, len(values), top_k)], + token_type_lookup=( + self.getTokenTypeLookup() if self.getIncludeTokenTypes() else None + ), + input_dtype=self.getInputKerasDtype(), + output_dtype=self.getOutputKerasDtype(), + name=f"{self.getLayerName()}_tokenizer", + ) + + def construct_layer_info(self) -> Dict[str, Any]: + """ + Constructs the layer info dictionary, appending the token-type output columns. + + Overrides the base method because the consolidated layer emits more outputs + than there are output columns when ``includeTokenTypes`` is set: it returns the + per-input type tensors after the token tensors, so the ``_types`` columns + must be appended in that same order for the pipeline graph to zip them up + correctly. + + :returns: Dictionary with the layer name, Keras layer, inputs and outputs. + """ + inputs, token_cols = self.get_layer_inputs_outputs() + + outputs = list(token_cols) + if self.getIncludeTokenTypes(): + outputs += [f"{col}_types" for col in token_cols] + + return { + "name": self.getOrDefault("layerName"), + "layer": self.get_keras_layer(), + "inputs": inputs, + "outputs": outputs, + } diff --git a/src/kamae/spark/utils/ngram_utils.py b/src/kamae/spark/utils/ngram_utils.py new file mode 100644 index 00000000..7907df74 --- /dev/null +++ b/src/kamae/spark/utils/ngram_utils.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. + +""" +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.spark.utils.ngram_worker_functions import ( + NUM_RESERVED_TOKENS, + PAD_TOKEN_ID, + UNK_TOKEN_ID, + encode_tuple, + extract_ngrams_from_column_worker, + extract_tuples_from_column_worker, +) + +logger = logging.getLogger(__name__) + +DEFAULT_EVENT_SIZE = 4 # Default discrete ID values per event (e.g. L0, L1, L2, L3) + + +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``). It is + compact categorical feature the model can embed alongside the token itself. The + reserved ```` and ```` tokens map to ``0``. + + The bitmask is parametric in the number of ID levels: the level integer is + parsed from each n-gram part, so it works for any ``tupleSize`` (and + multi-digit levels). + 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 collect_ngrams_from_dataframe( + df: DataFrame, + input_columns: List[str], + event_size: int = DEFAULT_EVENT_SIZE, +) -> 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), so + only the distinct n-grams 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. + :returns: Counter mapping each n-gram tuple to its corpus frequency. + :raises ValueError: If any of the input columns is missing from the DataFrame. + """ + 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}" + ) + + logger.info(f"Collecting n-grams from {len(input_columns)} columns in one pass...") + + 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) + .collect() + ) + + all_ngrams = Counter(dict(ngram_counts)) + logger.info(f"Found {len(all_ngrams):,} unique n-grams across all columns") + + return all_ngrams + + +def build_vocabulary( + ngram_counter: Counter, + vocab_size: int = 50000, + min_ngram_freq: int = 10, +) -> Dict[Tuple[str, ...], int]: + """ + Builds the vocabulary by filtering on frequency and keeping the top n-grams. + + N-grams below ``min_ngram_freq`` are dropped, the rest 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. + :param min_ngram_freq: Minimum frequency for an n-gram to be included. + :returns: Mapping from each kept n-gram tuple to its token id. + """ + filtered_ngrams = { + ngram: freq for ngram, freq in ngram_counter.items() if freq >= min_ngram_freq + } + logger.info( + f"After min_freq={min_ngram_freq}: {len(filtered_ngrams):,} n-grams " + f"(from {len(ngram_counter):,} total)" + ) + + # Sort by frequency descending, then by the n-gram itself so that equal-frequency + # n-grams are ordered deterministically across runs. + top_ngrams = sorted(filtered_ngrams.items(), key=lambda x: (-x[1], x[0]))[ + : vocab_size - NUM_RESERVED_TOKENS + ] + ngram_to_id = { + ngram: token_id + for token_id, (ngram, _) in enumerate(top_ngrams, start=NUM_RESERVED_TOKENS) + } + + logger.info( + f"Final vocabulary size: {len(ngram_to_id) + NUM_RESERVED_TOKENS:,} " + f"({len(ngram_to_id):,} n-grams + {NUM_RESERVED_TOKENS} reserved tokens)" + ) + + return ngram_to_id + + +def build_tuple_lookup_table( + df: DataFrame, + input_columns: List[str], + vocabulary: EventNgramVocabulary, + top_k: int, + event_size: int = DEFAULT_EVENT_SIZE, +) -> 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. + + :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 event tuple to its list of ``top_k`` token ids. + """ + if not input_columns: + return {} + + logger.info("Building tuple->tokens lookup table (distributed)...") + + # The fitted n-gram vocabulary is bounded by vocab_size (not by the number of + # tuples), so it is captured by the encoding closure and shipped to the workers. + 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))) + .collectAsMap() + ) + + logger.info(f"Encoded {len(tuple_to_tokens):,} unique tuples") + + 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 or + incomplete event yields padding, a tuple in the table yields its stored tokens, and + any other non-zero tuple yields ````. This is the single source of truth for + the row-level tokenization used by the Spark transform UDF, and is mirrored by the + TensorFlow ``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. + :returns: Flat token array of length ``num_events * top_k``. + """ + pad_tokens = [PAD_TOKEN_ID] * top_k + unk_tokens = [UNK_TOKEN_ID] * top_k + + if not ids: + return pad_tokens * num_events + + # A null id is read as 0 (an absent ID level), matching the fitting path and the + # dense TensorFlow input, rather than producing a tuple that can never be found. + event_tuples = [ + tuple( + 0 if id_value is None else id_value for id_value in ids[i : i + tuple_size] + ) + for i in range(0, len(ids), tuple_size) + ] + + all_tokens: List[int] = [] + for event_tuple in event_tuples: + if len(event_tuple) != tuple_size or all(v == 0 for v 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] 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..79c37039 --- /dev/null +++ b/src/kamae/spark/utils/ngram_worker_functions.py @@ -0,0 +1,132 @@ +# 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 functions run on Spark workers, so they depend only on the Python standard +library (no TensorFlow / kamae imports). +""" + +from itertools import combinations +from typing import Any, Dict, Iterator, List, Tuple + +PAD_TOKEN_ID = 0 +UNK_TOKEN_ID = 1 +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. A + trailing chunk with fewer than ``event_size`` ids is skipped, as it cannot be a + complete event. A null id is read as ``0``, i.e. as an absent ID level, which is + what the dense TensorFlow input carries for a missing id. + + :param col_value: Column value containing discrete ID values. + :param event_size: Number of discrete ID values per event. + :returns: Iterator over per-event lists of int ids. + """ + if col_value is None: + return + for i in range(0, len(col_value), event_size): + chunk = col_value[i : i + event_size] + if len(chunk) == event_size: + yield [0 if id_value is None else id_value for id_value in chunk] + + +def extract_ngrams_from_column_worker( + col_value: Any, event_size: int = 4 +) -> List[Tuple[str, ...]]: + """ + Extracts all within-event n-grams from a single row's column value. + + Each 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). 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. + """ + ngrams_list = [] + for event_ids in _iter_events(col_value, event_size): + event_strs = [f"L{j}_{v}" for j, v in enumerate(event_ids) if v != 0] + for length in range(1, len(event_strs) + 1): + for combo in combinations(range(len(event_strs)), length): + ngrams_list.append(tuple(event_strs[i] for i in combo)) + return ngrams_list + + +def extract_tuples_from_column_worker( + col_value: Any, event_size: int = 4 +) -> 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). + """ + tuples_list = [] + for event_ids in _iter_events(col_value, event_size): + if any(v != 0 for v in event_ids): + tuples_list.append(tuple(event_ids)) + return tuples_list + + +def encode_tuple( + id_tuple: Tuple[int, ...], + ngrams: Dict[Tuple[str, ...], int], + top_k: int, +) -> List[int]: + """ + Encodes one event tuple into its top-k token ids against a fitted vocabulary. + + Level-prefixes the non-zero ids (``L_``), enumerates every combination + of length 1..n (so n-grams skipping ID levels are considered), keeps those present + in ``ngrams``, and returns the ``top_k`` smallest ids (ids are assigned by + frequency, so smaller id = more frequent). An all-zero tuple yields padding; a + non-zero tuple whose every combination missed the vocabulary yields ````, one + per present ID level. + + Runs inside a Spark ``map`` on the workers, and is the single source of truth for + per-tuple encoding used by ``build_tuple_lookup_table``. + + :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). + """ + event = [f"L{j}_{v}" for j, v in enumerate(id_tuple) if v != 0] + if not event: + return [PAD_TOKEN_ID] * top_k + + candidate_ids = [] + for length in range(len(event), 0, -1): + for combo in combinations(range(len(event)), length): + ngram = tuple(event[i] for i in combo) + if ngram in ngrams: + candidate_ids.append(ngrams[ngram]) + + if not candidate_ids: + num_unk = min(len(event), top_k) + return [UNK_TOKEN_ID] * num_unk + [PAD_TOKEN_ID] * (top_k - num_unk) + + # Ids are unique per n-gram, so sorting ascending gives the most frequent first. + candidate_ids.sort() + tokens = candidate_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..5b247c15 --- /dev/null +++ b/tests/kamae/keras/tensorflow/layers/test_event_ngram_lookup.py @@ -0,0 +1,229 @@ +# 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 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, 1, 1], [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, 1, 1, 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)) + + 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, 1, 1, 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, 1, 1]], dtype=output.dtype)) + + 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_compute_output_shape_matches_the_call_output(self): + layer = _layer(num_events_per_input=[2, 1]) + shapes = layer.compute_output_shape([(None, 8), (None, 4)]) + assert shapes == [(None, 2 * TOP_K), (None, 1 * TOP_K)] + + def test_compute_output_shape_single_input_is_unwrapped(self): + # A lone input with no types returns one shape, not a list of one, matching + # what _call returns. + assert _layer().compute_output_shape((None, 8)) == (None, 2 * TOP_K) + + def test_compute_output_shape_preserves_the_list_dimension(self): + # Rank-3 in, rank-3 out: only the id axis is replaced. + assert _layer().compute_output_shape((None, 30, 8)) == (None, 30, 2 * TOP_K) + + def test_compute_output_shape_passes_through_unexpected_ranks(self): + # Neither rank-2 nor rank-3, so there is no id axis to resize; the shape is + # returned untouched rather than guessed at. + assert _layer().compute_output_shape((8,)) == (8,) + + def test_compute_output_shape_appends_type_shapes(self): + layer = _layer(token_type_lookup=[0, 0, 5, 1, 8, 3]) + shapes = layer.compute_output_shape((None, 8)) + assert shapes == [(None, 2 * TOP_K), (None, 2 * TOP_K)] diff --git a/tests/kamae/keras/test_jit_compatibility.py b/tests/kamae/keras/test_jit_compatibility.py index bd9c2354..65fbec50 100644 --- a/tests/kamae/keras/test_jit_compatibility.py +++ b/tests/kamae/keras/test_jit_compatibility.py @@ -69,6 +69,7 @@ DateDiffLayer, DateParseLayer, DateTimeToUnixTimestampLayer, + EventNgramLookupLayer, HashIndexLayer, IfStatementLayer, LambdaFunctionLayer, @@ -279,6 +280,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 e634e7b3..08ad9b39 100644 --- a/tests/kamae/keras/test_layer_serialisation.py +++ b/tests/kamae/keras/test_layer_serialisation.py @@ -76,6 +76,7 @@ DateDiffLayer, DateParseLayer, DateTimeToUnixTimestampLayer, + EventNgramLookupLayer, HashIndexLayer, IfStatementLayer, LambdaFunctionLayer, @@ -229,6 +230,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..3fc90bc5 --- /dev/null +++ b/tests/kamae/spark/estimators/test_event_ngram_lookup.py @@ -0,0 +1,153 @@ +# 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 +from pyspark.sql.types import ArrayType, IntegerType, StructField, StructType + +from kamae.spark.estimators import EventNgramLookupEstimator +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() + + 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_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. + estimator = self._estimator() + assert estimator.compatible_dtypes == [ + IntegerType(), + ArrayType(IntegerType()), + ] + + 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)], + ) + 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 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..5f446847 --- /dev/null +++ b/tests/kamae/spark/transformers/test_event_ngram_lookup.py @@ -0,0 +1,300 @@ +# 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-serialisable 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 search_level_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, search_level_df): + transformer = self._transformer() + actual = [ + row["clicks_tokens"] + for row in transformer.transform(search_level_df) + .select("clicks_tokens") + .collect() + ] + expected = [ + [2, 3, 0, 4, 5, 0], # both tuples known + [2, 3, 0, 1, 1, 1], # known then (id 1) + [0, 0, 0, 0, 0, 0], # both padding (id 0) + ] + assert actual == expected + + def test_spark_tf_parity_search_level(self, search_level_df): + transformer = self._transformer() + spark_out = [ + row["clicks_tokens"] + for row in transformer.transform(search_level_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 trailing partial event cannot be a tuple, so it is dropped. + ([1, 2, 3, 4, 5, 6], [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(Exception): + self._transformer().transform(df).collect() + + def test_spark_tf_parity_with_list_dimension(self, search_level_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(search_level_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_tokenizer" + assert layer.num_events_per_input == [2] + assert layer.top_k == TOP_K + assert layer.tuple_size == TUPLE_SIZE + + def test_multiple_inputs_one_layer_all_events(self): + transformer = EventNgramLookupTransformer( + inputCols=["clicks_ids", "prop_ids"], + outputCols=["clicks_tokens", "prop_tokens"], + numEventsPerInput=[2, 1], + tupleSize=TUPLE_SIZE, + topK=TOP_K, + vocabularySize=VOCAB_SIZE, + lookupKeys=LOOKUP_KEYS, + lookupValues=LOOKUP_VALUES, + layerName="tokenizer", + ) + # A single layer tokenizes all inputs (the table is embedded once). + layer = transformer.get_keras_layer() + assert isinstance(layer, EventNgramLookupLayer) + assert layer.num_events_per_input == [2, 1] + + 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-serialisable 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, search_level_df, tmp_path + ): + transformer = self._transformer() + expected = transformer.transform(search_level_df).collect() + + path = str(tmp_path / "transformer") + transformer.write().overwrite().save(path) + reloaded = EventNgramLookupTransformer.load(path) + + assert reloaded.transform(search_level_df).collect() == expected + + def test_transform_emits_type_columns(self, search_level_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(search_level_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_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, search_level_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(search_level_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..28e3ccbb --- /dev/null +++ b/tests/kamae/spark/utils/test_ngram_utils.py @@ -0,0 +1,251 @@ +# 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. + +from collections import Counter + +import pytest +from pyspark.sql.types import ArrayType, IntegerType, StructField, StructType + +from kamae.spark.utils.ngram_utils import ( + EventNgramVocabulary, + build_tuple_lookup_table, + build_vocabulary, + collect_ngrams_from_dataframe, + tokenize_events, +) +from kamae.spark.utils.ngram_worker_functions import ( + PAD_TOKEN_ID, + UNK_TOKEN_ID, + 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} + + +class TestBuildTupleLookupTable: + def test_no_input_columns_yields_an_empty_table(self, spark_session): + vocabulary = EventNgramVocabulary(ngrams={("L0_1",): 2}) + assert build_tuple_lookup_table(None, [], vocabulary, top_k=3) == {} + + +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, []), + # A trailing partial event is not a complete event and is skipped. + ([1, 0, 0, 5], 3, [("L0_1",)]), + (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], 3) + assert all( + not ({"L0_1"} <= set(ngram) and {"L0_5"} <= set(ngram)) for ngram in ngrams + ) + + +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, min_ngram_freq=1) + # 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_drops_ngrams_below_min_frequency(self): + counter = Counter({("L0_1",): 100, ("L1_2",): 5}) + ngram_to_id = build_vocabulary(counter, vocab_size=100, min_ngram_freq=10) + assert ngram_to_id == {("L0_1",): 2} + + 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, min_ngram_freq=1) + # 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, min_ngram_freq=1) + second = build_vocabulary( + Counter({("L0_1",): 10, ("L1_2",): 10}), vocab_size=100, min_ngram_freq=1 + ) + 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_is_padding(self): + assert encode_tuple((0, 0), self.NGRAMS, top_k=3) == [PAD_TOKEN_ID] * 3 + + def test_unmatched_tuple_yields_one_unk_per_present_level(self): + # Two non-zero levels, none in the vocabulary -> two unks, then padding. + assert encode_tuple((7, 8), self.NGRAMS, top_k=4) == [ + UNK_TOKEN_ID, + UNK_TOKEN_ID, + PAD_TOKEN_ID, + PAD_TOKEN_ID, + ] + + def test_unk_count_never_exceeds_top_k(self): + assert encode_tuple((7, 8, 9), self.NGRAMS, top_k=2) == [UNK_TOKEN_ID] * 2 + + +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. + ([1, 2, 3, 4, 9, 9, 9, 9], [2, 3, UNK_TOKEN_ID, UNK_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]), + # A trailing partial event is dropped. + ([1, 2, 3, 4, 5, 6], [2, 3, 0, 0]), + (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, UNK_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_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, UNK_TOKEN_ID] From 45480eadef20f48ea01827502606db3c9bf7f089 Mon Sep 17 00:00:00 2001 From: Sbranikas Date: Mon, 21 Sep 2026 16:02:34 +0200 Subject: [PATCH 2/6] feat: log the vocabulary's n-gram length distribution --- src/kamae/spark/utils/ngram_utils.py | 26 +++++++++++++++ tests/kamae/spark/utils/test_ngram_utils.py | 36 +++++++++++++++++++++ 2 files changed, 62 insertions(+) diff --git a/src/kamae/spark/utils/ngram_utils.py b/src/kamae/spark/utils/ngram_utils.py index 7907df74..cd9619e3 100644 --- a/src/kamae/spark/utils/ngram_utils.py +++ b/src/kamae/spark/utils/ngram_utils.py @@ -150,6 +150,31 @@ def collect_ngrams_from_dataframe( 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 tokenizer is mostly learning + single ID levels and the n-gram combinations are earning little, which is the + signal for tuning ``vocab_size`` and ``min_ngram_freq``. Counting is one pass + over the kept n-grams, so this is cheap enough to always report. + + :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(f"Vocabulary by n-gram length: {' | '.join(parts)}") + + def build_vocabulary( ngram_counter: Counter, vocab_size: int = 50000, @@ -191,6 +216,7 @@ def build_vocabulary( f"Final vocabulary size: {len(ngram_to_id) + NUM_RESERVED_TOKENS:,} " f"({len(ngram_to_id):,} n-grams + {NUM_RESERVED_TOKENS} reserved tokens)" ) + log_vocabulary_by_length(ngram_to_id) return ngram_to_id diff --git a/tests/kamae/spark/utils/test_ngram_utils.py b/tests/kamae/spark/utils/test_ngram_utils.py index 28e3ccbb..f3551bcc 100644 --- a/tests/kamae/spark/utils/test_ngram_utils.py +++ b/tests/kamae/spark/utils/test_ngram_utils.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import logging from collections import Counter import pytest @@ -22,6 +23,7 @@ build_tuple_lookup_table, build_vocabulary, collect_ngrams_from_dataframe, + log_vocabulary_by_length, tokenize_events, ) from kamae.spark.utils.ngram_worker_functions import ( @@ -249,3 +251,37 @@ 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, UNK_TOKEN_ID] + + +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, min_ngram_freq=1) + assert "Vocabulary by n-gram length" in caplog.text From fe12b3a4a1c29f116e02a2188d784c3f64c7624a Mon Sep 17 00:00:00 2001 From: Sbranikas Date: Wed, 23 Sep 2026 11:46:41 +0200 Subject: [PATCH 3/6] fix: Improve tokeniser input contract, and add sanitization.: --- .../tensorflow/layers/event_ngram_lookup.py | 34 +++--- .../spark/estimators/event_ngram_lookup.py | 6 +- src/kamae/spark/params/shared.py | 12 ++ .../spark/transformers/event_ngram_lookup.py | 26 ++++- src/kamae/spark/utils/__init__.py | 8 ++ src/kamae/spark/utils/ngram_utils.py | 84 ++++++++------ .../spark/utils/ngram_worker_functions.py | 109 ++++++++++-------- .../layers/test_event_ngram_lookup.py | 26 +++++ .../estimators/test_event_ngram_lookup.py | 20 ++-- .../transformers/test_event_ngram_lookup.py | 19 ++- tests/kamae/spark/utils/test_ngram_utils.py | 29 +++-- 11 files changed, 250 insertions(+), 123 deletions(-) diff --git a/src/kamae/keras/tensorflow/layers/event_ngram_lookup.py b/src/kamae/keras/tensorflow/layers/event_ngram_lookup.py index 70b58ffc..69236aea 100644 --- a/src/kamae/keras/tensorflow/layers/event_ngram_lookup.py +++ b/src/kamae/keras/tensorflow/layers/event_ngram_lookup.py @@ -143,9 +143,18 @@ def _build_lookup_table( :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 the packed key would not fit in a signed ``int64``. + :raises ValueError: If an ID is negative, or the packed key would not fit in a + signed ``int64``. """ - max_id = max((int(x) for k in keys for x in k), default=1) + key_ids = [int(x) for k in keys for x in k] + min_id = min(key_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(key_ids, default=1) self.key_bits = max(1, max_id.bit_length()) if self.tuple_size * self.key_bits > 63: raise ValueError( @@ -203,13 +212,10 @@ def _tokenize(self, inputs: KerasTensor, num_events: int) -> KerasTensor: if input_rank == 3: batch_size = tf.shape(inputs)[0] list_or_1 = tf.shape(inputs)[1] - num_ids = tf.shape(inputs)[2] - inputs_flat = tf.reshape(inputs, [-1, num_ids]) - eff_batch = batch_size * list_or_1 + inputs_flat = tf.reshape(inputs, [-1, tf.shape(inputs)[2]]) restore_list = True else: inputs_flat = inputs - eff_batch = tf.shape(inputs_flat)[0] restore_list = False # Pad/truncate the id axis to num_events * tuple_size, then split into tuples. @@ -224,12 +230,12 @@ def _tokenize(self, inputs: KerasTensor, num_events: int) -> KerasTensor: # 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 wider than key_bits cannot be + # 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 <= max_id, axis=1) + 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) @@ -245,7 +251,7 @@ def _tokenize(self, inputs: KerasTensor, num_events: int) -> KerasTensor: if restore_list: return tf.reshape(tokens, [batch_size, list_or_1, output_length]) - return tf.reshape(tokens, [eff_batch, output_length]) + return tf.reshape(tokens, [-1, output_length]) def _call( self, inputs: Union[KerasTensor, List[KerasTensor]], **kwargs: Any @@ -300,12 +306,10 @@ def compute_output_shape( token_shapes.append(shape) if self._token_type_lookup is not None: - outputs = token_shapes + list(token_shapes) - elif single: - return token_shapes[0] - else: - outputs = token_shapes - return outputs + # One type tensor per input, each the shape of its tokens, appended after + # all the token tensors. + return token_shapes + token_shapes + return token_shapes[0] if single else token_shapes def get_config(self) -> Dict[str, Any]: """ diff --git a/src/kamae/spark/estimators/event_ngram_lookup.py b/src/kamae/spark/estimators/event_ngram_lookup.py index bf51281a..150a1cb5 100644 --- a/src/kamae/spark/estimators/event_ngram_lookup.py +++ b/src/kamae/spark/estimators/event_ngram_lookup.py @@ -32,7 +32,7 @@ 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 +from pyspark.sql.types import DataType, IntegerType, LongType from kamae.keras.core.backend import TENSORFLOW_ONLY from kamae.spark.params import ( @@ -41,7 +41,7 @@ SampleFractionParams, ) from kamae.spark.transformers import EventNgramLookupTransformer -from kamae.spark.utils.ngram_utils import ( +from kamae.spark.utils import ( EventNgramVocabulary, build_tuple_lookup_table, build_vocabulary, @@ -205,7 +205,7 @@ def compatible_dtypes(self) -> Optional[List[DataType]]: :returns: List of compatible data types for the estimator. """ - return [IntegerType(), ArrayType(IntegerType())] + return [IntegerType(), LongType()] def _fit(self, dataset: DataFrame) -> EventNgramLookupTransformer: """ diff --git a/src/kamae/spark/params/shared.py b/src/kamae/spark/params/shared.py index c3418c2f..f037669e 100644 --- a/src/kamae/spark/params/shared.py +++ b/src/kamae/spark/params/shared.py @@ -1118,8 +1118,14 @@ 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]: @@ -1135,8 +1141,11 @@ 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: @@ -1152,8 +1161,11 @@ 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: diff --git a/src/kamae/spark/transformers/event_ngram_lookup.py b/src/kamae/spark/transformers/event_ngram_lookup.py index bda25ed9..e8818430 100644 --- a/src/kamae/spark/transformers/event_ngram_lookup.py +++ b/src/kamae/spark/transformers/event_ngram_lookup.py @@ -40,12 +40,19 @@ import tensorflow as tf from pyspark import keyword_only from pyspark.sql import DataFrame -from pyspark.sql.types import ArrayType, DataType, IntegerType, StructField, StructType +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.ngram_utils import tokenize_events +from kamae.spark.utils import tokenize_events, validate_input_columns from .base import BaseTransformer @@ -131,7 +138,7 @@ def compatible_dtypes(self) -> Optional[List[DataType]]: :returns: List of compatible data types for the layer. """ - return [IntegerType(), ArrayType(IntegerType())] + return [IntegerType(), LongType()] def _transform(self, dataset: DataFrame) -> DataFrame: """ @@ -143,9 +150,12 @@ def _transform(self, dataset: DataFrame) -> DataFrame: column is added by gathering the type bitmask at each token id. :param dataset: Input DataFrame. + :raises ValueError: If a column is missing, or is not a single-level array. :returns: DataFrame with the tokenized output columns, and the optional type columns. """ + validate_input_columns(dataset, self.getInputCols()) + tuple_size = self.getTupleSize() top_k = self.getTopK() lookup_table = self.getTupleToTokens() @@ -196,10 +206,16 @@ def _tokenize_with_types( ), ) struct_col = f"{output_col}__tokens_and_types" + # The type columns are not in outputCols, so the egress cast applied to the + # output columns does not reach them. Cast them here so that they carry the + # same dtype as the type tensors the Keras layer returns. + 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(f"{output_col}_types", F.col(f"{struct_col}.types")) + .withColumn(f"{output_col}_types", casted_types) .drop(struct_col) ) @@ -236,7 +252,7 @@ def get_keras_layer(self) -> tf.keras.layers.Layer: ), input_dtype=self.getInputKerasDtype(), output_dtype=self.getOutputKerasDtype(), - name=f"{self.getLayerName()}_tokenizer", + name=self.getLayerName(), ) def construct_layer_info(self) -> Dict[str, Any]: diff --git a/src/kamae/spark/utils/__init__.py b/src/kamae/spark/utils/__init__.py index 39e0c3e9..9f43f621 100644 --- a/src/kamae/spark/utils/__init__.py +++ b/src/kamae/spark/utils/__init__.py @@ -33,6 +33,14 @@ 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, + validate_input_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 index cd9619e3..feed91bc 100644 --- a/src/kamae/spark/utils/ngram_utils.py +++ b/src/kamae/spark/utils/ngram_utils.py @@ -28,6 +28,7 @@ from pyspark.sql import DataFrame +from kamae.spark.utils.array_utils import get_array_nesting_level from kamae.spark.utils.ngram_worker_functions import ( NUM_RESERVED_TOKENS, PAD_TOKEN_ID, @@ -35,12 +36,11 @@ encode_tuple, extract_ngrams_from_column_worker, extract_tuples_from_column_worker, + iter_events, ) logger = logging.getLogger(__name__) -DEFAULT_EVENT_SIZE = 4 # Default discrete ID values per event (e.g. L0, L1, L2, L3) - class EventNgramVocabulary: """ @@ -100,10 +100,43 @@ def build_type_lookup(self) -> List[int]: return lookup +def validate_input_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. + + :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 = DEFAULT_EVENT_SIZE, + event_size: int, ) -> Counter: """ Extracts and counts within-event n-grams across the given columns, distributed. @@ -117,16 +150,9 @@ def collect_ngrams_from_dataframe( :param input_columns: Column names holding discrete ID values. :param event_size: Number of discrete ID values per event. :returns: Counter mapping each n-gram tuple to its corpus frequency. - :raises ValueError: If any of the input columns is missing from the DataFrame. + :raises ValueError: If a column is missing, or is not a single-level array. """ - 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}" - ) + validate_input_columns(df, input_columns) logger.info(f"Collecting n-grams from {len(input_columns)} columns in one pass...") @@ -177,8 +203,8 @@ def log_vocabulary_by_length(ngram_to_id: Dict[Tuple[str, ...], int]) -> None: def build_vocabulary( ngram_counter: Counter, - vocab_size: int = 50000, - min_ngram_freq: int = 10, + vocab_size: int, + min_ngram_freq: int, ) -> Dict[Tuple[str, ...], int]: """ Builds the vocabulary by filtering on frequency and keeping the top n-grams. @@ -205,7 +231,7 @@ def build_vocabulary( # Sort by frequency descending, then by the n-gram itself so that equal-frequency # n-grams are ordered deterministically across runs. top_ngrams = sorted(filtered_ngrams.items(), key=lambda x: (-x[1], x[0]))[ - : vocab_size - NUM_RESERVED_TOKENS + : max(vocab_size - NUM_RESERVED_TOKENS, 0) ] ngram_to_id = { ngram: token_id @@ -226,7 +252,7 @@ def build_tuple_lookup_table( input_columns: List[str], vocabulary: EventNgramVocabulary, top_k: int, - event_size: int = DEFAULT_EVENT_SIZE, + event_size: int, ) -> Dict[Tuple[int, ...], List[int]]: """ Pre-computes the ``id_tuple -> top-k token ids`` lookup table, distributed. @@ -243,9 +269,6 @@ def build_tuple_lookup_table( :param event_size: Number of discrete ID values per tuple. :returns: Mapping from each event tuple to its list of ``top_k`` token ids. """ - if not input_columns: - return {} - logger.info("Building tuple->tokens lookup table (distributed)...") # The fitted n-gram vocabulary is bounded by vocab_size (not by the number of @@ -281,10 +304,10 @@ def tokenize_events( """ 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 or - incomplete event yields padding, a tuple in the table yields its stored tokens, and - any other non-zero tuple yields ````. This is the single source of truth for - the row-level tokenization used by the Spark transform UDF, and is mirrored by the + 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 ````. This is the single source of truth for the + row-level tokenization used by the Spark transform UDF, and is mirrored by the TensorFlow ``EventNgramLookupLayer``. :param ids: One row's discrete ID values as a flat int array, split into @@ -294,6 +317,7 @@ def tokenize_events( :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 @@ -302,18 +326,10 @@ def tokenize_events( if not ids: return pad_tokens * num_events - # A null id is read as 0 (an absent ID level), matching the fitting path and the - # dense TensorFlow input, rather than producing a tuple that can never be found. - event_tuples = [ - tuple( - 0 if id_value is None else id_value for id_value in ids[i : i + tuple_size] - ) - for i in range(0, len(ids), tuple_size) - ] - all_tokens: List[int] = [] - for event_tuple in event_tuples: - if len(event_tuple) != tuple_size or all(v == 0 for v in event_tuple): + 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)) diff --git a/src/kamae/spark/utils/ngram_worker_functions.py b/src/kamae/spark/utils/ngram_worker_functions.py index 79c37039..d766452e 100644 --- a/src/kamae/spark/utils/ngram_worker_functions.py +++ b/src/kamae/spark/utils/ngram_worker_functions.py @@ -20,59 +20,84 @@ """ from itertools import combinations -from typing import Any, Dict, Iterator, List, Tuple +from typing import Any, Dict, Iterator, List, Sequence, Tuple PAD_TOKEN_ID = 0 UNK_TOKEN_ID = 1 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]]: +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. A - trailing chunk with fewer than ``event_size`` ids is skipped, as it cannot be a - complete event. A null id is read as ``0``, i.e. as an absent ID level, which is - what the dense TensorFlow input carries for a missing id. + 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): - chunk = col_value[i : i + event_size] - if len(chunk) == event_size: - yield [0 if id_value is None else id_value for id_value in chunk] + 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. This is the single + definition of an n-gram, used both when counting them 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 = 4 + col_value: Any, event_size: int ) -> List[Tuple[str, ...]]: """ Extracts all within-event n-grams from a single row's column value. - Each 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). N-grams are never formed across event boundaries. + 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. """ - ngrams_list = [] - for event_ids in _iter_events(col_value, event_size): - event_strs = [f"L{j}_{v}" for j, v in enumerate(event_ids) if v != 0] - for length in range(1, len(event_strs) + 1): - for combo in combinations(range(len(event_strs)), length): - ngrams_list.append(tuple(event_strs[i] for i in combo)) - return ngrams_list + 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 = 4 + col_value: Any, event_size: int ) -> List[Tuple[int, ...]]: """ Extracts all non-zero event tuples from a single row's column value. @@ -81,11 +106,11 @@ def extract_tuples_from_column_worker( :param event_size: Number of discrete ID values per tuple. :returns: List of ID tuples (all-zero events are skipped). """ - tuples_list = [] - for event_ids in _iter_events(col_value, event_size): - if any(v != 0 for v in event_ids): - tuples_list.append(tuple(event_ids)) - return tuples_list + 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( @@ -96,12 +121,10 @@ def encode_tuple( """ Encodes one event tuple into its top-k token ids against a fitted vocabulary. - Level-prefixes the non-zero ids (``L_``), enumerates every combination - of length 1..n (so n-grams skipping ID levels are considered), keeps those present - in ``ngrams``, and returns the ``top_k`` smallest ids (ids are assigned by - frequency, so smaller id = more frequent). An all-zero tuple yields padding; a - non-zero tuple whose every combination missed the vocabulary yields ````, one - per present ID level. + 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. An all-zero tuple yields padding, and a non-zero tuple whose every + n-gram missed the vocabulary yields ````, one per present ID level. Runs inside a Spark ``map`` on the workers, and is the single source of truth for per-tuple encoding used by ``build_tuple_lookup_table``. @@ -111,22 +134,16 @@ def encode_tuple( :param top_k: Number of tokens to return. :returns: List of token ids of length ``top_k`` (padded if needed). """ - event = [f"L{j}_{v}" for j, v in enumerate(id_tuple) if v != 0] - if not event: + present_levels = sum(1 for id_value in id_tuple if id_value != 0) + if present_levels == 0: return [PAD_TOKEN_ID] * top_k - candidate_ids = [] - for length in range(len(event), 0, -1): - for combo in combinations(range(len(event)), length): - ngram = tuple(event[i] for i in combo) - if ngram in ngrams: - candidate_ids.append(ngrams[ngram]) - - if not candidate_ids: - num_unk = min(len(event), top_k) + token_ids = sorted( + ngrams[ngram] for ngram in event_ngrams(id_tuple) if ngram in ngrams + ) + if not token_ids: + num_unk = min(present_levels, top_k) return [UNK_TOKEN_ID] * num_unk + [PAD_TOKEN_ID] * (top_k - num_unk) - # Ids are unique per n-gram, so sorting ascending gives the most frequent first. - candidate_ids.sort() - tokens = candidate_ids[:top_k] + 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 index 5b247c15..34d438c4 100644 --- a/tests/kamae/keras/tensorflow/layers/test_event_ngram_lookup.py +++ b/tests/kamae/keras/tensorflow/layers/test_event_ngram_lookup.py @@ -179,6 +179,32 @@ def test_out_of_range_id_misses_instead_of_aliasing_a_valid_key(self): # then tf.debugging.assert_equal(output, tf.constant([[1, 1, 1]], 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, 1, 1]], 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"): diff --git a/tests/kamae/spark/estimators/test_event_ngram_lookup.py b/tests/kamae/spark/estimators/test_event_ngram_lookup.py index 3fc90bc5..38da7376 100644 --- a/tests/kamae/spark/estimators/test_event_ngram_lookup.py +++ b/tests/kamae/spark/estimators/test_event_ngram_lookup.py @@ -13,7 +13,7 @@ # limitations under the License. import pytest -from pyspark.sql.types import ArrayType, IntegerType, StructField, StructType +from pyspark.sql.types import ArrayType, IntegerType, LongType, StructField, StructType from kamae.spark.estimators import EventNgramLookupEstimator from kamae.spark.transformers import EventNgramLookupTransformer @@ -113,12 +113,11 @@ def test_multiple_inputs_share_one_vocabulary(self, spark_session): 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. + # 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(), - ArrayType(IntegerType()), - ] + 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 @@ -129,7 +128,14 @@ def test_use_fit_sample_defaults_to_false(self, id_df): @pytest.mark.parametrize( "param, value", - [("vocabSize", 2), ("minNgramFreq", 0)], + [ + ("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): diff --git a/tests/kamae/spark/transformers/test_event_ngram_lookup.py b/tests/kamae/spark/transformers/test_event_ngram_lookup.py index 5f446847..2edcf877 100644 --- a/tests/kamae/spark/transformers/test_event_ngram_lookup.py +++ b/tests/kamae/spark/transformers/test_event_ngram_lookup.py @@ -103,8 +103,6 @@ def test_spark_tf_parity_search_level(self, search_level_df): [ # A short row is right-padded up to numEvents * topK. ([1, 2, 3, 4], [2, 3, 0, 0, 0, 0]), - # A trailing partial event cannot be a tuple, so it is dropped. - ([1, 2, 3, 4, 5, 6], [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. @@ -167,7 +165,7 @@ def test_spark_tf_parity_with_list_dimension(self, search_level_df): 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_tokenizer" + assert layer.name == "clicks_tokens" assert layer.num_events_per_input == [2] assert layer.top_k == TOP_K assert layer.tuple_size == TUPLE_SIZE @@ -263,6 +261,21 @@ def test_transform_emits_type_columns(self, search_level_df): type_lookup[t] for t in row["clicks_tokens"] ] + def test_output_dtype_applies_to_the_type_columns(self, search_level_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(search_level_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. diff --git a/tests/kamae/spark/utils/test_ngram_utils.py b/tests/kamae/spark/utils/test_ngram_utils.py index f3551bcc..73aba6be 100644 --- a/tests/kamae/spark/utils/test_ngram_utils.py +++ b/tests/kamae/spark/utils/test_ngram_utils.py @@ -53,11 +53,14 @@ def test_counts_across_all_given_columns(self, id_df): assert len(counter) == 15 assert set(counter.values()) == {1} - -class TestBuildTupleLookupTable: - def test_no_input_columns_yields_an_empty_table(self, spark_session): - vocabulary = EventNgramVocabulary(ngrams={("L0_1",): 2}) - assert build_tuple_lookup_table(None, [], vocabulary, top_k=3) == {} + 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: @@ -79,8 +82,6 @@ class TestExtractNgramsFromColumnWorker: ), # An all-zero event contributes nothing. ([0, 0, 0], 3, []), - # A trailing partial event is not a complete event and is skipped. - ([1, 0, 0, 5], 3, [("L0_1",)]), (None, 3, []), ([], 3, []), # A null id is an absent ID level, exactly like a 0, rather than becoming @@ -99,11 +100,15 @@ def test_extracts_within_event_combinations(self, col_value, event_size, expecte 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], 3) + 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( @@ -232,8 +237,6 @@ class TestTokenizeEvents: ([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]), - # A trailing partial event is dropped. - ([1, 2, 3, 4, 5, 6], [2, 3, 0, 0]), (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 @@ -247,6 +250,12 @@ def test_tokenizes_to_fixed_length(self, ids, expected): 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 From cc45bee7d8fd0c7a177412d7922b587aff61d977 Mon Sep 17 00:00:00 2001 From: Sbranikas Date: Wed, 23 Sep 2026 19:14:35 +0200 Subject: [PATCH 4/6] refactor: harden EventNgramLookup pipeline integration and unknown handling --- .../tensorflow/layers/event_ngram_lookup.py | 117 +++++++++--------- .../spark/estimators/event_ngram_lookup.py | 43 ++++--- src/kamae/spark/params/shared.py | 17 +++ .../spark/transformers/event_ngram_lookup.py | 57 ++++----- src/kamae/spark/utils/__init__.py | 3 +- src/kamae/spark/utils/ngram_utils.py | 76 +++++++++--- .../spark/utils/ngram_worker_functions.py | 29 +++-- .../layers/test_event_ngram_lookup.py | 52 ++++---- .../estimators/test_event_ngram_lookup.py | 113 +++++++++++++++++ .../transformers/test_event_ngram_lookup.py | 52 ++++++-- tests/kamae/spark/utils/test_ngram_utils.py | 39 ++++-- 11 files changed, 416 insertions(+), 182 deletions(-) diff --git a/src/kamae/keras/tensorflow/layers/event_ngram_lookup.py b/src/kamae/keras/tensorflow/layers/event_ngram_lookup.py index 69236aea..e305344b 100644 --- a/src/kamae/keras/tensorflow/layers/event_ngram_lookup.py +++ b/src/kamae/keras/tensorflow/layers/event_ngram_lookup.py @@ -36,11 +36,15 @@ exports cleanly to SavedModel. Token id conventions (shared with the vocabulary): - - ``0`` = padding (````): emitted for all-zero / missing events. - - ``1`` = unknown (````): emitted for a non-zero tuple absent from the table. + - ``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. Every unrecognised tuple takes this one + path, whether it was unseen while fitting or matched no n-gram, so the encoding + depends on the tuple rather than on whether it appeared in the fitting corpus. """ -from typing import Any, Dict, List, Optional, Union +from typing import Any, Dict, Iterable, List, Optional, Union import keras import tensorflow as tf @@ -54,6 +58,42 @@ 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): """ @@ -134,35 +174,18 @@ def _build_lookup_table( """ Builds the ``IntegerLookup`` sublayer and gathered values tensor. - Each tuple key is bijectively packed into a single ``int64`` by allotting - ``key_bits`` bits per ID level, where ``key_bits`` is the bit length of the - largest ID in the table. 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. + 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``. """ - key_ids = [int(x) for k in keys for x in k] - min_id = min(key_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(key_ids, default=1) - self.key_bits = max(1, max_id.bit_length()) - if self.tuple_size * self.key_bits > 63: - raise ValueError( - f"Cannot pack event tuples into a signed int64: tuple_size " - f"({self.tuple_size}) * key_bits ({self.key_bits}, from a largest ID " - f"of {max_id}) = {self.tuple_size * self.key_bits} bits, which exceeds " - f"63. Reduce tuple_size or the ID cardinality." - ) + 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 @@ -185,7 +208,9 @@ def _build_lookup_table( mask_token=None, name=f"{self.name}_key_lookup", ) - self.unk_pattern = tf.constant([UNK_TOKEN_ID] * self.top_k, dtype=tf.int32) + 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 @@ -261,12 +286,19 @@ def _call( :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) @@ -280,37 +312,6 @@ def _call( types = [tf.gather(self._type_tensor, tok) for tok in tokens] return tokens + types - def compute_output_shape( - self, input_shape: Union[tuple, List[tuple]] - ) -> Union[tuple, List[tuple]]: - """ - Declares the output shape(s) for functional model building. - - :param input_shape: A single input shape, or a list of input shapes (one per - input column). - :returns: The matching output shape(s), with each id axis replaced by - ``num_events * top_k``. When types are emitted, the per-input type shapes are - appended after the token shapes. - """ - single = not isinstance(input_shape[0], (list, tuple)) - shapes = [input_shape] if single else list(input_shape) - - token_shapes = [] - for shape, num_events in zip(shapes, self.num_events_per_input): - output_length = num_events * self.top_k - if len(shape) == 3: - token_shapes.append((shape[0], shape[1], output_length)) - elif len(shape) == 2: - token_shapes.append((shape[0], output_length)) - else: - token_shapes.append(shape) - - if self._token_type_lookup is not None: - # One type tensor per input, each the shape of its tokens, appended after - # all the token tensors. - return token_shapes + token_shapes - return token_shapes[0] if single else token_shapes - def get_config(self) -> Dict[str, Any]: """ Gets the configuration of the EventNgramLookupLayer layer. diff --git a/src/kamae/spark/estimators/event_ngram_lookup.py b/src/kamae/spark/estimators/event_ngram_lookup.py index 150a1cb5..6b35613f 100644 --- a/src/kamae/spark/estimators/event_ngram_lookup.py +++ b/src/kamae/spark/estimators/event_ngram_lookup.py @@ -21,13 +21,14 @@ 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 +from typing import List, Optional, Tuple from pyspark import keyword_only from pyspark.ml.param import Param, Params, TypeConverters @@ -35,6 +36,7 @@ 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, @@ -46,6 +48,7 @@ build_tuple_lookup_table, build_vocabulary, collect_ngrams_from_dataframe, + validate_event_column_lengths, ) from .base import BaseEstimator @@ -207,6 +210,20 @@ def compatible_dtypes(self) -> Optional[List[DataType]]: """ 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. @@ -214,38 +231,30 @@ def _fit(self, dataset: DataFrame) -> EventNgramLookupTransformer: :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. + do not all have the same length, or the observed IDs cannot be packed by the + Keras layer (negative, or too large for ``tupleSize``). """ input_cols = self.getInputCols() output_cols = self.getOutputCols() num_events_per_input = self.getNumEventsPerInput() - - if len(input_cols) != len(output_cols): - raise ValueError( - f"inputCols and outputCols must have the same length. Got " - f"{len(input_cols)} inputs and {len(output_cols)} outputs." - ) - if num_events_per_input is None or len(input_cols) != 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_cols)} inputs and {n_events} event counts." - ) + validate_event_column_lengths(input_cols, output_cols, num_events_per_input) tuple_size = self.getTupleSize() top_k = self.getTopK() + min_ngram_freq = self.getMinNgramFreq() # 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=min_ngram_freq, ) vocabulary = EventNgramVocabulary( ngrams=build_vocabulary( ngram_counter=ngram_counter, vocab_size=self.getVocabSize(), - min_ngram_freq=self.getMinNgramFreq(), + min_ngram_freq=min_ngram_freq, ) ) @@ -273,6 +282,10 @@ def _fit(self, dataset: DataFrame) -> EventNgramLookupTransformer: lookup_keys.extend(id_tuple) lookup_values.extend(tokens) + # Fail now, rather than when the Keras layer is built, if the Keras layer could + # not pack the observed tuples into its int64 keys. + compute_key_bits(lookup_keys, tuple_size) + return EventNgramLookupTransformer( inputCols=input_cols, outputCols=output_cols, diff --git a/src/kamae/spark/params/shared.py b/src/kamae/spark/params/shared.py index f037669e..f854796f 100644 --- a/src/kamae/spark/params/shared.py +++ b/src/kamae/spark/params/shared.py @@ -1282,3 +1282,20 @@ def getTokenTypeLookup(self) -> Optional[List[int]]: :returns: Per-token-id list of ID-level bitmasks, or None. """ return self.getOrDefault(self.tokenTypeLookup) + + 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/event_ngram_lookup.py b/src/kamae/spark/transformers/event_ngram_lookup.py index e8818430..d694dfb5 100644 --- a/src/kamae/spark/transformers/event_ngram_lookup.py +++ b/src/kamae/spark/transformers/event_ngram_lookup.py @@ -28,13 +28,17 @@ Token id conventions (shared with the vocabulary and the TF layer): - ``0`` = padding (````): emitted for all-zero / missing events. - ``1`` = unknown (````): emitted for a non-zero tuple absent from the table. + +The table only holds the event tuples that matched at least one n-gram during fitting, +so any other non-zero tuple, including a new combination of individually known IDs, +maps to a single ```` followed by padding. """ # pylint: disable=unused-argument # pylint: disable=invalid-name # pylint: disable=too-many-ancestors # pylint: disable=no-member -from typing import Any, Dict, List, Optional, Tuple +from typing import List, Optional, Tuple import pyspark.sql.functions as F import tensorflow as tf @@ -52,7 +56,11 @@ 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, validate_input_columns +from kamae.spark.utils import ( + tokenize_events, + validate_event_column_lengths, + validate_event_id_columns, +) from .base import BaseTransformer @@ -150,23 +158,27 @@ def _transform(self, dataset: DataFrame) -> DataFrame: column is added by gathering the type bitmask at each token id. :param dataset: Input DataFrame. - :raises ValueError: If a column is missing, or is not a single-level array. + :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. """ - validate_input_columns(dataset, self.getInputCols()) + 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) tuple_size = self.getTupleSize() top_k = self.getTopK() lookup_table = self.getTupleToTokens() include_types = self.getIncludeTokenTypes() type_lookup = self.getTokenTypeLookup() if include_types else None + type_cols = dict(zip(output_cols, self.getTokenTypeCols(output_cols))) token_array_type = ArrayType(IntegerType()) for input_col, output_col, num_events in zip( - self.getInputCols(), - self.getOutputCols(), - self.getNumEventsPerInput(), + input_cols, output_cols, num_events_per_input ): # num_events is bound per column via the default argument so each UDF # captures its own value rather than the last loop iteration's. @@ -215,7 +227,7 @@ def _tokenize_with_types( dataset = ( dataset.withColumn(struct_col, tokenize_udf(F.col(input_col))) .withColumn(output_col, F.col(f"{struct_col}.tokens")) - .withColumn(f"{output_col}_types", casted_types) + .withColumn(type_cols[output_col], casted_types) .drop(struct_col) ) @@ -255,27 +267,16 @@ def get_keras_layer(self) -> tf.keras.layers.Layer: name=self.getLayerName(), ) - def construct_layer_info(self) -> Dict[str, Any]: + def get_layer_inputs_outputs(self) -> Tuple[List[str], List[str]]: """ - Constructs the layer info dictionary, appending the token-type output columns. + Gets the input and output column names, including the token-type columns. - Overrides the base method because the consolidated layer emits more outputs - than there are output columns when ``includeTokenTypes`` is set: it returns the - per-input type tensors after the token tensors, so the ``_types`` columns - must be appended in that same order for the pipeline graph to zip them up - correctly. + Overrides the base method because, with ``includeTokenTypes`` set, the layer + emits more outputs than there are output columns: it returns the per-input + type tensors after the token tensors, so the ``_types`` columns are + appended in that same order for the pipeline graph to zip them up correctly. - :returns: Dictionary with the layer name, Keras layer, inputs and outputs. + :returns: Tuple of the input column names and the output column names. """ - inputs, token_cols = self.get_layer_inputs_outputs() - - outputs = list(token_cols) - if self.getIncludeTokenTypes(): - outputs += [f"{col}_types" for col in token_cols] - - return { - "name": self.getOrDefault("layerName"), - "layer": self.get_keras_layer(), - "inputs": inputs, - "outputs": outputs, - } + 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 9f43f621..82fa5abe 100644 --- a/src/kamae/spark/utils/__init__.py +++ b/src/kamae/spark/utils/__init__.py @@ -39,7 +39,8 @@ build_vocabulary, collect_ngrams_from_dataframe, tokenize_events, - validate_input_columns, + validate_event_column_lengths, + validate_event_id_columns, ) from .transform_utils import ( # noqa: F401 multi_input_single_output_array_transform, diff --git a/src/kamae/spark/utils/ngram_utils.py b/src/kamae/spark/utils/ngram_utils.py index feed91bc..0a2c206e 100644 --- a/src/kamae/spark/utils/ngram_utils.py +++ b/src/kamae/spark/utils/ngram_utils.py @@ -100,13 +100,45 @@ def build_type_lookup(self) -> List[int]: return lookup -def validate_input_columns(df: DataFrame, input_columns: List[str]) -> None: +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. + 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. @@ -137,22 +169,26 @@ 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), so - only the distinct n-grams reach the driver. + 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. - :returns: Counter mapping each n-gram tuple to its corpus frequency. + :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_input_columns(df, input_columns) + validate_event_id_columns(df, input_columns) logger.info(f"Collecting n-grams from {len(input_columns)} columns in one pass...") @@ -167,11 +203,15 @@ def collect_ngrams_from_dataframe( ) .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(f"Found {len(all_ngrams):,} unique n-grams across all columns") + logger.info( + f"Found {len(all_ngrams):,} unique n-grams with frequency >= " + f"{min_ngram_freq} across all columns" + ) return all_ngrams @@ -223,10 +263,7 @@ def build_vocabulary( filtered_ngrams = { ngram: freq for ngram, freq in ngram_counter.items() if freq >= min_ngram_freq } - logger.info( - f"After min_freq={min_ngram_freq}: {len(filtered_ngrams):,} n-grams " - f"(from {len(ngram_counter):,} total)" - ) + logger.info(f"After min_freq={min_ngram_freq}: {len(filtered_ngrams):,} n-grams") # Sort by frequency descending, then by the n-gram itself so that equal-frequency # n-grams are ordered deterministically across runs. @@ -262,12 +299,16 @@ def build_tuple_lookup_table( 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 + ```` at inference by the same rule as a tuple never seen during fitting. This + also keeps the table to the tuples that carry a token. + :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 event tuple to its list of ``top_k`` token ids. + :returns: Mapping from each token-bearing event tuple to its ``top_k`` token ids. """ logger.info("Building tuple->tokens lookup table (distributed)...") @@ -286,6 +327,7 @@ def build_tuple_lookup_table( ) .distinct() .map(lambda id_tuple: (id_tuple, encode_tuple(id_tuple, ngrams, top_k))) + .filter(lambda kv: kv[1] is not None) .collectAsMap() ) @@ -306,9 +348,11 @@ def tokenize_events( 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 ````. This is the single source of truth for the - row-level tokenization used by the Spark transform UDF, and is mirrored by the - TensorFlow ``EventNgramLookupLayer``. + non-zero tuple is unrecognised and yields a single ```` followed by padding. + One ```` marks that an event occurred but matched nothing, without giving an + unrecognised event more embedded tokens than a recognised one. This is the single + source of truth for the row-level tokenization used by the Spark transform UDF, and + is mirrored by the TensorFlow ``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 @@ -321,7 +365,7 @@ def tokenize_events( :returns: Flat token array of length ``num_events * top_k``. """ pad_tokens = [PAD_TOKEN_ID] * top_k - unk_tokens = [UNK_TOKEN_ID] * top_k + unk_tokens = [UNK_TOKEN_ID] + [PAD_TOKEN_ID] * (top_k - 1) if not ids: return pad_tokens * num_events diff --git a/src/kamae/spark/utils/ngram_worker_functions.py b/src/kamae/spark/utils/ngram_worker_functions.py index d766452e..f2e8439f 100644 --- a/src/kamae/spark/utils/ngram_worker_functions.py +++ b/src/kamae/spark/utils/ngram_worker_functions.py @@ -15,12 +15,14 @@ """ Worker functions for n-gram extraction in Spark RDD operations. -These functions run on Spark workers, so they depend only on the Python standard -library (no TensorFlow / kamae imports). +These run inside ``flatMap`` / ``map`` on the Spark workers and are pure Python over +plain ints and tuples (standard library only). Importing this module still imports the +``kamae`` package, and with it TensorFlow, so kamae must be installed on the workers, +as for kamae's other UDF-based transformers. """ from itertools import combinations -from typing import Any, Dict, Iterator, List, Sequence, Tuple +from typing import Any, Dict, Iterator, List, Optional, Sequence, Tuple PAD_TOKEN_ID = 0 UNK_TOKEN_ID = 1 @@ -117,14 +119,19 @@ def encode_tuple( id_tuple: Tuple[int, ...], ngrams: Dict[Tuple[str, ...], int], top_k: int, -) -> List[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. An all-zero tuple yields padding, and a non-zero tuple whose every - n-gram missed the vocabulary yields ````, one per present ID level. + frequent n-gram. + + Returns `None` when the tuple has no token to contribute, i.e. when it is all-zero + or when every one of its n-grams missed the vocabulary. Such a tuple is left out of + the lookup table, so at inference it takes the same table-miss path as a tuple that + was never seen during fitting. Unknown is therefore decided in one place, by one + rule, rather than once here and again at inference. Runs inside a Spark ``map`` on the workers, and is the single source of truth for per-tuple encoding used by ``build_tuple_lookup_table``. @@ -132,18 +139,14 @@ def encode_tuple( :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). + :returns: List of token ids of length ``top_k`` (padded if needed), or `None` if the + tuple matched no n-gram. """ - present_levels = sum(1 for id_value in id_tuple if id_value != 0) - if present_levels == 0: - return [PAD_TOKEN_ID] * top_k - token_ids = sorted( ngrams[ngram] for ngram in event_ngrams(id_tuple) if ngram in ngrams ) if not token_ids: - num_unk = min(present_levels, top_k) - return [UNK_TOKEN_ID] * num_unk + [PAD_TOKEN_ID] * (top_k - num_unk) + 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 index 34d438c4..8dfb6340 100644 --- a/tests/kamae/keras/tensorflow/layers/test_event_ngram_lookup.py +++ b/tests/kamae/keras/tensorflow/layers/test_event_ngram_lookup.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import keras import pytest import tensorflow as tf @@ -64,7 +65,7 @@ def test_known_unknown_and_padding_tuples(self, input_dtype): tf.debugging.assert_equal( output, tf.constant( - [[2, 3, 0, 4, 5, 0], [2, 3, 0, 1, 1, 1], [0, 0, 0, 0, 0, 0]], + [[2, 3, 0, 4, 5, 0], [2, 3, 0, 1, 0, 0], [0, 0, 0, 0, 0, 0]], dtype=output.dtype, ), ) @@ -89,7 +90,7 @@ def test_preserves_list_dimension_for_rank_three_input(self): tf.constant( [ [[2, 3, 0, 4, 5, 0], [0, 0, 0, 0, 0, 0]], - [[4, 5, 0, 2, 3, 0], [1, 1, 1, 0, 0, 0]], + [[4, 5, 0, 2, 3, 0], [1, 0, 0, 0, 0, 0]], ], dtype=output.dtype, ), @@ -163,7 +164,7 @@ 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, 1, 1, 0, 0, 0]], dtype=output.dtype) + 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): @@ -177,7 +178,7 @@ def test_out_of_range_id_misses_instead_of_aliasing_a_valid_key(self): output = layer(tf.constant([[17, 2, 3, 4]], dtype=tf.int32)) # then - tf.debugging.assert_equal(output, tf.constant([[1, 1, 1]], dtype=output.dtype)) + 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 @@ -193,7 +194,7 @@ def test_negative_id_misses_instead_of_aliasing_a_valid_key(self): output = layer(tf.constant([[-3, 6, 7, 8]], dtype=tf.int32)) # then - tf.debugging.assert_equal(output, tf.constant([[1, 1, 1]], dtype=output.dtype)) + 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 @@ -230,26 +231,27 @@ def test_get_config_round_trips_the_table(self): tf.debugging.assert_equal(tokens, expected_tokens) tf.debugging.assert_equal(types, expected_types) - def test_compute_output_shape_matches_the_call_output(self): - layer = _layer(num_events_per_input=[2, 1]) - shapes = layer.compute_output_shape([(None, 8), (None, 4)]) - assert shapes == [(None, 2 * TOP_K), (None, 1 * TOP_K)] - - def test_compute_output_shape_single_input_is_unwrapped(self): - # A lone input with no types returns one shape, not a list of one, matching - # what _call returns. - assert _layer().compute_output_shape((None, 8)) == (None, 2 * TOP_K) + 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") - def test_compute_output_shape_preserves_the_list_dimension(self): - # Rank-3 in, rank-3 out: only the id axis is replaced. - assert _layer().compute_output_shape((None, 30, 8)) == (None, 30, 2 * TOP_K) + outputs = layer([clicks, prop]) - def test_compute_output_shape_passes_through_unexpected_ranks(self): - # Neither rank-2 nor rank-3, so there is no id axis to resize; the shape is - # returned untouched rather than guessed at. - assert _layer().compute_output_shape((8,)) == (8,) + 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_compute_output_shape_appends_type_shapes(self): - layer = _layer(token_type_lookup=[0, 0, 5, 1, 8, 3]) - shapes = layer.compute_output_shape((None, 8)) - assert shapes == [(None, 2 * TOP_K), (None, 2 * 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)) diff --git a/tests/kamae/spark/estimators/test_event_ngram_lookup.py b/tests/kamae/spark/estimators/test_event_ngram_lookup.py index 38da7376..c49efb59 100644 --- a/tests/kamae/spark/estimators/test_event_ngram_lookup.py +++ b/tests/kamae/spark/estimators/test_event_ngram_lookup.py @@ -12,10 +12,13 @@ # 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 @@ -111,6 +114,19 @@ def test_multiple_inputs_share_one_vocabulary(self, spark_session): 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 @@ -157,3 +173,100 @@ def test_include_token_types_propagates_type_lookup(self, id_df): 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 index 2edcf877..597a1b2c 100644 --- a/tests/kamae/spark/transformers/test_event_ngram_lookup.py +++ b/tests/kamae/spark/transformers/test_event_ngram_lookup.py @@ -73,7 +73,7 @@ def test_transform_known_unknown_padding(self, search_level_df): ] expected = [ [2, 3, 0, 4, 5, 0], # both tuples known - [2, 3, 0, 1, 1, 1], # known then (id 1) + [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 @@ -133,9 +133,15 @@ def test_transform_rejects_string_input_column(self, spark_session): # 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(Exception): + with pytest.raises(TypeError): self._transformer().transform(df).collect() + def test_transform_raises_when_column_lengths_differ(self, search_level_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(search_level_df) + def test_spark_tf_parity_with_list_dimension(self, search_level_df): # Rank-3 inputs (batch, list_size, ids) arise in listwise models; the list # dimension must be preserved and each item tokenized independently. @@ -170,22 +176,42 @@ def test_get_keras_layer_returns_single_consolidated_layer(self): assert layer.top_k == TOP_K assert layer.tuple_size == TUPLE_SIZE - def test_multiple_inputs_one_layer_all_events(self): - transformer = EventNgramLookupTransformer( + 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], - tupleSize=TUPLE_SIZE, - topK=TOP_K, - vocabularySize=VOCAB_SIZE, - lookupKeys=LOOKUP_KEYS, - lookupValues=LOOKUP_VALUES, + includeTokenTypes=True, + tokenTypeLookup=[0, 0, 5, 1, 8, 3], layerName="tokenizer", ) - # A single layer tokenizes all inputs (the table is embedded once). - layer = transformer.get_keras_layer() - assert isinstance(layer, EventNgramLookupLayer) - assert layer.num_events_per_input == [2, 1] + 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() diff --git a/tests/kamae/spark/utils/test_ngram_utils.py b/tests/kamae/spark/utils/test_ngram_utils.py index 73aba6be..610601e0 100644 --- a/tests/kamae/spark/utils/test_ngram_utils.py +++ b/tests/kamae/spark/utils/test_ngram_utils.py @@ -53,6 +53,18 @@ def test_counts_across_all_given_columns(self, id_df): 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( @@ -206,21 +218,21 @@ def test_pads_when_fewer_matches_than_top_k(self): PAD_TOKEN_ID, ] - def test_all_zero_tuple_is_padding(self): - assert encode_tuple((0, 0), self.NGRAMS, top_k=3) == [PAD_TOKEN_ID] * 3 + 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_yields_one_unk_per_present_level(self): - # Two non-zero levels, none in the vocabulary -> two unks, then padding. - assert encode_tuple((7, 8), self.NGRAMS, top_k=4) == [ - UNK_TOKEN_ID, + 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, ] - def test_unk_count_never_exceeds_top_k(self): - assert encode_tuple((7, 8, 9), self.NGRAMS, top_k=2) == [UNK_TOKEN_ID] * 2 - class TestTokenizeEvents: LOOKUP = {(1, 2, 3, 4): [2, 3], (5, 6, 7, 8): [4, 5]} @@ -229,8 +241,9 @@ class TestTokenizeEvents: "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. - ([1, 2, 3, 4, 9, 9, 9, 9], [2, 3, UNK_TOKEN_ID, UNK_TOKEN_ID]), + # 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. @@ -242,7 +255,7 @@ class TestTokenizeEvents: # 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, UNK_TOKEN_ID, 0, 0]), + ([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): @@ -259,7 +272,7 @@ def test_raises_when_the_ids_are_not_a_whole_number_of_events(self): 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, UNK_TOKEN_ID] + ) == [UNK_TOKEN_ID, PAD_TOKEN_ID] class TestLogVocabularyByLength: From cba0aeb75b81b97231a4c6770d8e113c93cfa1aa Mon Sep 17 00:00:00 2001 From: Sbranikas Date: Mon, 28 Sep 2026 16:51:29 +0200 Subject: [PATCH 5/6] refactor: move fitted params to the transformer and tidy EventNgramLookup --- README.md | 2 +- .../tensorflow/layers/event_ngram_lookup.py | 39 +-- .../spark/estimators/event_ngram_lookup.py | 15 +- src/kamae/spark/params/shared.py | 134 +-------- .../spark/transformers/event_ngram_lookup.py | 275 ++++++++++++------ src/kamae/spark/utils/__init__.py | 1 + src/kamae/spark/utils/ngram_utils.py | 90 +++--- .../spark/utils/ngram_worker_functions.py | 24 +- .../layers/test_event_ngram_lookup.py | 5 + .../estimators/test_event_ngram_lookup.py | 8 + .../transformers/test_event_ngram_lookup.py | 46 ++- tests/kamae/spark/utils/test_ngram_utils.py | 33 ++- 12 files changed, 332 insertions(+), 340 deletions(-) diff --git a/README.md b/README.md index 1b636726..fd570db1 100644 --- a/README.md +++ b/README.md @@ -115,7 +115,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 | Tokenises 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) | +| 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/event_ngram_lookup.py b/src/kamae/keras/tensorflow/layers/event_ngram_lookup.py index e305344b..487ec5bd 100644 --- a/src/kamae/keras/tensorflow/layers/event_ngram_lookup.py +++ b/src/kamae/keras/tensorflow/layers/event_ngram_lookup.py @@ -29,19 +29,15 @@ ``[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, applied fully -vectorised (no per-element ``map_fn`` / ``cond`` and 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 from a ``.keras`` file and -exports cleanly to SavedModel. +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. Every unrecognised tuple takes this one - path, whether it was unseen while fitting or matched no n-gram, so the encoding - depends on the tuple rather than on whether it appeared in the fitting corpus. + single ```` followed by padding. """ from typing import Any, Dict, Iterable, List, Optional, Union @@ -105,9 +101,14 @@ class EventNgramLookupLayer(BaseLayer): 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 (e.g. - per listwise item) is preserved. The id axis is padded/truncated to - ``num_events * tuple_size`` before being split into tuples. + rank-3 ``(batch, list_size, num_events * tuple_size)``; the list dimension is + preserved. The id axis is zero-padded or truncated to ``num_events * tuple_size`` + before being split into tuples. + + Unlike the Spark transformer, which raises when a row's ids are not a whole number + of events, the layer zero-pads a partial trailing event and looks it up. This is + deliberate: it keeps assertions out of the serving graph, so a malformed request + is tokenized rather than failed. 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. @@ -229,14 +230,20 @@ def _tokenize(self, inputs: KerasTensor, num_events: int) -> KerasTensor: :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. :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}." + ) if input_rank == 3: batch_size = tf.shape(inputs)[0] - list_or_1 = tf.shape(inputs)[1] + list_size = tf.shape(inputs)[1] inputs_flat = tf.reshape(inputs, [-1, tf.shape(inputs)[2]]) restore_list = True else: @@ -244,8 +251,7 @@ def _tokenize(self, inputs: KerasTensor, num_events: int) -> KerasTensor: restore_list = False # Pad/truncate the id axis to num_events * tuple_size, then split into tuples. - # Padding by the full expected length before slicing covers both the short and - # the long case without branching on the input width. + # 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 @@ -269,13 +275,12 @@ def _tokenize(self, inputs: KerasTensor, num_events: int) -> KerasTensor: found = (indices > 0) & in_range gathered = tf.gather(self.values_tensor, tf.maximum(indices - 1, 0)) - # tf.where broadcasts, so the (top_k,) unk/pad constants apply as they are, - # with no per-tuple copy materialised. + # 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_or_1, output_length]) + return tf.reshape(tokens, [batch_size, list_size, output_length]) return tf.reshape(tokens, [-1, output_length]) def _call( diff --git a/src/kamae/spark/estimators/event_ngram_lookup.py b/src/kamae/spark/estimators/event_ngram_lookup.py index 6b35613f..68e3ce57 100644 --- a/src/kamae/spark/estimators/event_ngram_lookup.py +++ b/src/kamae/spark/estimators/event_ngram_lookup.py @@ -241,20 +241,17 @@ def _fit(self, dataset: DataFrame) -> EventNgramLookupTransformer: tuple_size = self.getTupleSize() top_k = self.getTopK() - min_ngram_freq = self.getMinNgramFreq() # 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=min_ngram_freq, + min_ngram_freq=self.getMinNgramFreq(), ) vocabulary = EventNgramVocabulary( ngrams=build_vocabulary( - ngram_counter=ngram_counter, - vocab_size=self.getVocabSize(), - min_ngram_freq=min_ngram_freq, + ngram_counter=ngram_counter, vocab_size=self.getVocabSize() ) ) @@ -273,17 +270,15 @@ def _fit(self, dataset: DataFrame) -> EventNgramLookupTransformer: vocabulary.build_type_lookup() if include_token_types else None ) - # Flatten the table into two parallel int lists so that it is JSON-serialisable - # and the fitted pipeline can be saved. One pass over the items keeps the keys - # and values positionally aligned. + # 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 now, rather than when the Keras layer is built, if the Keras layer could - # not pack the observed tuples into its int64 keys. + # 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( diff --git a/src/kamae/spark/params/shared.py b/src/kamae/spark/params/shared.py index f854796f..6fea8e1d 100644 --- a/src/kamae/spark/params/shared.py +++ b/src/kamae/spark/params/shared.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Dict, List, Optional, Tuple, Union +from typing import List, Union from pyspark.ml.param import Param, Params, TypeConverters @@ -1067,52 +1067,14 @@ class EventNgramLookupParams(Params): typeConverter=TypeConverters.toInt, ) - vocabularySize = Param( - Params._dummy(), - "vocabularySize", - "Number of token ids the fitted vocabulary actually produced, including the " - "reserved pad and unk tokens. This is the fitted result, not the requested " - "target (which is the estimator's vocabSize and is only an upper bound): it is " - "published so a downstream embedding table can be sized to the tokens that " - "were really learned.", - typeConverter=TypeConverters.toInt, - ) - - # The fitted lookup table is held as two parallel flat int arrays rather than as a - # dict, because Spark ML writes params to JSON metadata and a dict keyed by id - # tuples is not JSON-serialisable, which would make the fitted pipeline unsaveable. - # Same parallel-arrays shape as StringMapParams. - 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, - ) - includeTokenTypes = Param( Params._dummy(), "includeTokenTypes", - "Whether to also emit, per input column, a parallel '_types' column giving " - "each token's ID-level bitmask (a compact categorical feature for the model).", + "Whether to also emit, per input column, a parallel '_types' column " + "giving each token's ID-level bitmask.", typeConverter=TypeConverters.toBoolean, ) - 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 setNumEventsPerInput(self, value: List[int]) -> "EventNgramLookupParams": """ Sets the numEventsPerInput parameter. @@ -1176,79 +1138,6 @@ def getTopK(self) -> int: """ return self.getOrDefault(self.topK) - def setVocabularySize(self, value: int) -> "EventNgramLookupParams": - """ - Sets the vocabularySize parameter. - - :param value: Total number of unique tokens in the fitted vocabulary. - :returns: Instance of class mixed in. - """ - return self._set(vocabularySize=value) - - def getVocabularySize(self) -> int: - """ - Gets the vocabularySize parameter. - - :returns: Total number of unique tokens in the fitted vocabulary. - """ - return self.getOrDefault(self.vocabularySize) - - def setLookupKeys(self, value: List[int]) -> "EventNgramLookupParams": - """ - 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]) -> "EventNgramLookupParams": - """ - 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 getTupleToTokens(self) -> Dict[Tuple[int, ...], List[int]]: - """ - Rebuilds the ``event tuple -> top-k token list`` lookup table. - - Derived from the flat ``lookupKeys`` / ``lookupValues`` params, which are what - the fitted pipeline actually persists. - - :returns: Mapping from each event tuple to its list of token ids. - """ - keys = self.getLookupKeys() - values = self.getLookupValues() - if not keys or not values: - return {} - 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) - ) - } - def setIncludeTokenTypes(self, value: bool) -> "EventNgramLookupParams": """ Sets the includeTokenTypes parameter. @@ -1266,23 +1155,6 @@ def getIncludeTokenTypes(self) -> bool: """ return self.getOrDefault(self.includeTokenTypes) - def setTokenTypeLookup(self, value: List[int]) -> "EventNgramLookupParams": - """ - 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 getTokenTypeCols(self, token_cols: List[str]) -> List[str]: """ Gets the names of the derived token-type columns. diff --git a/src/kamae/spark/transformers/event_ngram_lookup.py b/src/kamae/spark/transformers/event_ngram_lookup.py index d694dfb5..2814310f 100644 --- a/src/kamae/spark/transformers/event_ngram_lookup.py +++ b/src/kamae/spark/transformers/event_ngram_lookup.py @@ -16,33 +16,29 @@ EventNgramLookupTransformer: tokenizes discrete IDs with a pre-computed lookup table. Applies the ``tuple -> top-k tokens`` table fitted by ``EventNgramLookupEstimator``. -Both the Spark path (``_transform``) and the TensorFlow path (``get_keras_layer``) are -plain ``O(1)`` lookups, so the same event tuple always maps to the same tokens and the -two paths produce identical output. +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 (a compact -categorical feature for the model). The type is a pure function of the token id, so it -is derived by gathering the fitted ``tokenTypeLookup`` at the token ids. +parallel ``_types`` column giving each token's ID-level bitmask. -Token id conventions (shared with the vocabulary and the TF layer): +Token id conventions (shared with the vocabulary and the Keras layer): - ``0`` = padding (````): emitted for all-zero / missing events. - - ``1`` = unknown (````): emitted for a non-zero tuple absent from the table. - -The table only holds the event tuples that matched at least one n-gram during fitting, -so any other non-zero tuple, including a new combination of individually known IDs, -maps to a single ```` followed by padding. + - ``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 typing import List, Optional, Tuple +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, @@ -58,6 +54,7 @@ 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, ) @@ -65,18 +62,146 @@ 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 (padding for all-zero events, ```` for non-zero tuples absent from - the table), producing a flat ``numEvents * topK`` array per input column. When + 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. """ @@ -113,7 +238,7 @@ def __init__( :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: Total number of unique tokens in the fitted vocabulary. + :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 @@ -150,12 +275,12 @@ def compatible_dtypes(self) -> Optional[List[DataType]]: def _transform(self, dataset: DataFrame) -> DataFrame: """ - Tokenizes each input column with a per-column UDF over ``tokenize_events``. + 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, a ``_types`` - column is added by gathering the type bitmask at each token id. + ``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 @@ -169,96 +294,77 @@ def _transform(self, dataset: DataFrame) -> DataFrame: validate_event_column_lengths(input_cols, output_cols, num_events_per_input) validate_event_id_columns(dataset, input_cols) - tuple_size = self.getTupleSize() - top_k = self.getTopK() - lookup_table = self.getTupleToTokens() - include_types = self.getIncludeTokenTypes() - type_lookup = self.getTokenTypeLookup() if include_types else None - type_cols = dict(zip(output_cols, self.getTokenTypeCols(output_cols))) - + table_kwargs = { + "lookup_table": self.getTupleToTokens(), + "tuple_size": self.getTupleSize(), + "top_k": self.getTopK(), + } token_array_type = ArrayType(IntegerType()) - for input_col, output_col, num_events in zip( - input_cols, output_cols, num_events_per_input - ): - # num_events is bound per column via the default argument so each UDF - # captures its own value rather than the last loop iteration's. - if not include_types: + 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( - lambda ids, num_events=num_events: tokenize_events( - ids, lookup_table, num_events, tuple_size, top_k - ), + partial(tokenize_events, num_events=num_events, **table_kwargs), token_array_type, ) dataset = dataset.withColumn(output_col, tokenize_udf(F.col(input_col))) - continue - - # A token's type is a pure function of its id, so both are produced by one - # UDF returning a struct, keeping the column to a single Python round-trip. - def _tokenize_with_types( - ids: Optional[List[int]], num_events: int = num_events - ) -> Tuple[List[int], List[int]]: - """Tokenizes one row and gathers each token's type bitmask. - - :param ids: One row's discrete ID values. - :param num_events: Number of events this column is sized to. - :returns: Tuple of (token ids, type bitmasks), same length. - """ - tokens = tokenize_events( - ids, lookup_table, num_events, tuple_size, top_k - ) - return tokens, [type_lookup[token] for token in tokens] - + 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( - _tokenize_with_types, - StructType( - [ - StructField("tokens", token_array_type), - StructField("types", token_array_type), - ] + 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 egress cast applied to the - # output columns does not reach them. Cast them here so that they carry the - # same dtype as the type tensors the Keras layer returns. + # 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_cols[output_col], casted_types) + .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), for in-graph tokenization on - TensorFlow (needed when reading raw files via TFParquet). The layer outputs one - token tensor per input, plus one type tensor per input when - ``includeTokenTypes`` is set. + 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``. """ - # The layer takes the table as nested key/value lists, so the flat params are - # reshaped back into one list per tuple. - tuple_size = self.getTupleSize() - top_k = self.getTopK() - keys = self.getLookupKeys() or [] - values = self.getLookupValues() or [] + lookup_table = self.getTupleToTokens() return EventNgramLookupLayer( num_events_per_input=self.getNumEventsPerInput(), - top_k=top_k, - tuple_size=tuple_size, - lookup_keys=[ - keys[i : i + tuple_size] for i in range(0, len(keys), tuple_size) - ], - lookup_values=[values[i : i + top_k] for i in range(0, len(values), top_k)], + 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 ), @@ -272,9 +378,8 @@ 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 - emits more outputs than there are output columns: it returns the per-input - type tensors after the token tensors, so the ``_types`` columns are - appended in that same order for the pipeline graph to zip them up correctly. + 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. """ diff --git a/src/kamae/spark/utils/__init__.py b/src/kamae/spark/utils/__init__.py index 82fa5abe..478d253c 100644 --- a/src/kamae/spark/utils/__init__.py +++ b/src/kamae/spark/utils/__init__.py @@ -39,6 +39,7 @@ build_vocabulary, collect_ngrams_from_dataframe, tokenize_events, + tokenize_events_with_types, validate_event_column_lengths, validate_event_id_columns, ) diff --git a/src/kamae/spark/utils/ngram_utils.py b/src/kamae/spark/utils/ngram_utils.py index 0a2c206e..055c5fef 100644 --- a/src/kamae/spark/utils/ngram_utils.py +++ b/src/kamae/spark/utils/ngram_utils.py @@ -28,11 +28,10 @@ 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, - PAD_TOKEN_ID, - UNK_TOKEN_ID, encode_tuple, extract_ngrams_from_column_worker, extract_tuples_from_column_worker, @@ -79,14 +78,9 @@ def build_type_lookup(self) -> List[int]: 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``). It is - compact categorical feature the model can embed alongside the token itself. The - reserved ```` and ```` tokens map to ``0``. - - The bitmask is parametric in the number of ID levels: the level integer is - parsed from each n-gram part, so it works for any ``tupleSize`` (and - multi-digit levels). - The type cardinality is ``2 ** tupleSize``. + 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. """ @@ -190,7 +184,7 @@ def collect_ngrams_from_dataframe( """ validate_event_id_columns(df, input_columns) - logger.info(f"Collecting n-grams from {len(input_columns)} columns in one pass...") + logger.info("Collecting n-grams from %d columns in one pass...", len(input_columns)) ngram_counts = ( df.select(*input_columns) @@ -209,8 +203,9 @@ def collect_ngrams_from_dataframe( all_ngrams = Counter(dict(ngram_counts)) logger.info( - f"Found {len(all_ngrams):,} unique n-grams with frequency >= " - f"{min_ngram_freq} across all columns" + "Found %d unique n-grams with frequency >= %d across all columns", + len(all_ngrams), + min_ngram_freq, ) return all_ngrams @@ -220,10 +215,8 @@ 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 tokenizer is mostly learning - single ID levels and the n-gram combinations are earning little, which is the - signal for tuning ``vocab_size`` and ``min_ngram_freq``. Counting is one pass - over the kept n-grams, so this is cheap enough to always report. + 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. @@ -238,36 +231,30 @@ def log_vocabulary_by_length(ngram_to_id: Dict[Tuple[str, ...], int]) -> None: f"({100.0 * by_length[length] / total:.1f}%)" for length in sorted(by_length) ] - logger.info(f"Vocabulary by n-gram length: {' | '.join(parts)}") + logger.info("Vocabulary by n-gram length: %s", " | ".join(parts)) def build_vocabulary( ngram_counter: Counter, vocab_size: int, - min_ngram_freq: int, ) -> Dict[Tuple[str, ...], int]: """ - Builds the vocabulary by filtering on frequency and keeping the top n-grams. + Builds the vocabulary by keeping the most frequent n-grams. - N-grams below ``min_ngram_freq`` are dropped, the rest are sorted by descending - frequency (ties broken by the n-gram itself, for determinism) and the top + 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. - :param min_ngram_freq: Minimum frequency for an n-gram to be included. :returns: Mapping from each kept n-gram tuple to its token id. """ - filtered_ngrams = { - ngram: freq for ngram, freq in ngram_counter.items() if freq >= min_ngram_freq - } - logger.info(f"After min_freq={min_ngram_freq}: {len(filtered_ngrams):,} n-grams") - # Sort by frequency descending, then by the n-gram itself so that equal-frequency # n-grams are ordered deterministically across runs. - top_ngrams = sorted(filtered_ngrams.items(), key=lambda x: (-x[1], x[0]))[ + top_ngrams = sorted(ngram_counter.items(), key=lambda x: (-x[1], x[0]))[ : max(vocab_size - NUM_RESERVED_TOKENS, 0) ] ngram_to_id = { @@ -276,8 +263,10 @@ def build_vocabulary( } logger.info( - f"Final vocabulary size: {len(ngram_to_id) + NUM_RESERVED_TOKENS:,} " - f"({len(ngram_to_id):,} n-grams + {NUM_RESERVED_TOKENS} reserved tokens)" + "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) @@ -300,8 +289,7 @@ def build_tuple_lookup_table( depend on how the work was partitioned. Tuples that matched no n-gram are left out of the table, so they resolve to - ```` at inference by the same rule as a tuple never seen during fitting. This - also keeps the table to the tuples that carry a token. + ```` 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. @@ -311,9 +299,6 @@ def build_tuple_lookup_table( :returns: Mapping from each token-bearing event tuple to its ``top_k`` token ids. """ logger.info("Building tuple->tokens lookup table (distributed)...") - - # The fitted n-gram vocabulary is bounded by vocab_size (not by the number of - # tuples), so it is captured by the encoding closure and shipped to the workers. ngrams = vocabulary.ngrams tuple_to_tokens = ( @@ -331,7 +316,7 @@ def build_tuple_lookup_table( .collectAsMap() ) - logger.info(f"Encoded {len(tuple_to_tokens):,} unique tuples") + logger.info("Encoded %d unique tuples", len(tuple_to_tokens)) return tuple_to_tokens @@ -348,11 +333,8 @@ def tokenize_events( 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 is unrecognised and yields a single ```` followed by padding. - One ```` marks that an event occurred but matched nothing, without giving an - unrecognised event more embedded tokens than a recognised one. This is the single - source of truth for the row-level tokenization used by the Spark transform UDF, and - is mirrored by the TensorFlow ``EventNgramLookupLayer``. + 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 @@ -383,3 +365,27 @@ def tokenize_events( 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 index f2e8439f..a434f58b 100644 --- a/src/kamae/spark/utils/ngram_worker_functions.py +++ b/src/kamae/spark/utils/ngram_worker_functions.py @@ -16,16 +16,14 @@ 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 (standard library only). Importing this module still imports the -``kamae`` package, and with it TensorFlow, so kamae must be installed on the workers, -as for kamae's other UDF-based transformers. +plain ints and tuples. """ from itertools import combinations from typing import Any, Dict, Iterator, List, Optional, Sequence, Tuple -PAD_TOKEN_ID = 0 -UNK_TOKEN_ID = 1 +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 @@ -63,9 +61,8 @@ def event_ngrams(event_ids: Sequence[int]) -> Iterator[Tuple[str, ...]]: 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. This is the single - definition of an n-gram, used both when counting them and when encoding a tuple - against a fitted vocabulary. + 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. @@ -127,14 +124,9 @@ def encode_tuple( smallest ids; ids are assigned by descending frequency, so a smaller id is a more frequent n-gram. - Returns `None` when the tuple has no token to contribute, i.e. when it is all-zero - or when every one of its n-grams missed the vocabulary. Such a tuple is left out of - the lookup table, so at inference it takes the same table-miss path as a tuple that - was never seen during fitting. Unknown is therefore decided in one place, by one - rule, rather than once here and again at inference. - - Runs inside a Spark ``map`` on the workers, and is the single source of truth for - per-tuple encoding used by ``build_tuple_lookup_table``. + 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. diff --git a/tests/kamae/keras/tensorflow/layers/test_event_ngram_lookup.py b/tests/kamae/keras/tensorflow/layers/test_event_ngram_lookup.py index 8dfb6340..040eac32 100644 --- a/tests/kamae/keras/tensorflow/layers/test_event_ngram_lookup.py +++ b/tests/kamae/keras/tensorflow/layers/test_event_ngram_lookup.py @@ -255,3 +255,8 @@ def test_raises_when_inputs_do_not_match_num_events_per_input(self): 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/spark/estimators/test_event_ngram_lookup.py b/tests/kamae/spark/estimators/test_event_ngram_lookup.py index c49efb59..365acdc0 100644 --- a/tests/kamae/spark/estimators/test_event_ngram_lookup.py +++ b/tests/kamae/spark/estimators/test_event_ngram_lookup.py @@ -63,6 +63,14 @@ def test_fit_returns_transformer_with_propagated_params(self, id_df): # 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) diff --git a/tests/kamae/spark/transformers/test_event_ngram_lookup.py b/tests/kamae/spark/transformers/test_event_ngram_lookup.py index 597a1b2c..355d39fa 100644 --- a/tests/kamae/spark/transformers/test_event_ngram_lookup.py +++ b/tests/kamae/spark/transformers/test_event_ngram_lookup.py @@ -27,7 +27,7 @@ # 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-serialisable int lists. +# 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 @@ -52,7 +52,7 @@ def _transformer(self, **overrides) -> EventNgramLookupTransformer: return EventNgramLookupTransformer(**params) @pytest.fixture - def search_level_df(self, spark_session): + 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 = [ @@ -63,13 +63,11 @@ def search_level_df(self, spark_session): schema = StructType([StructField("clicks_ids", ArrayType(IntegerType()), True)]) return spark_session.createDataFrame(data, schema) - def test_transform_known_unknown_padding(self, search_level_df): + def test_transform_known_unknown_padding(self, id_df): transformer = self._transformer() actual = [ row["clicks_tokens"] - for row in transformer.transform(search_level_df) - .select("clicks_tokens") - .collect() + for row in transformer.transform(id_df).select("clicks_tokens").collect() ] expected = [ [2, 3, 0, 4, 5, 0], # both tuples known @@ -78,13 +76,11 @@ def test_transform_known_unknown_padding(self, search_level_df): ] assert actual == expected - def test_spark_tf_parity_search_level(self, search_level_df): + def test_spark_tf_parity_search_level(self, id_df): transformer = self._transformer() spark_out = [ row["clicks_tokens"] - for row in transformer.transform(search_level_df) - .select("clicks_tokens") - .collect() + for row in transformer.transform(id_df).select("clicks_tokens").collect() ] layer = transformer.get_keras_layer() tf_in = tf.constant( @@ -136,21 +132,19 @@ def test_transform_rejects_string_input_column(self, spark_session): with pytest.raises(TypeError): self._transformer().transform(df).collect() - def test_transform_raises_when_column_lengths_differ(self, search_level_df): + 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(search_level_df) + transformer.transform(id_df) - def test_spark_tf_parity_with_list_dimension(self, search_level_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(search_level_df) - .select("clicks_tokens") - .collect() + 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. @@ -244,7 +238,7 @@ def test_getters(self): 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-serialisable shape for the transformer, and any pipeline holding it, to + # 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] @@ -259,25 +253,25 @@ def test_params_survive_save_and_load(self, spark_session, tmp_path): assert reloaded.getTokenTypeLookup() == [0, 0, 5, 1, 8, 3] def test_reloaded_transformer_tokenizes_identically( - self, spark_session, search_level_df, tmp_path + self, spark_session, id_df, tmp_path ): transformer = self._transformer() - expected = transformer.transform(search_level_df).collect() + expected = transformer.transform(id_df).collect() path = str(tmp_path / "transformer") transformer.write().overwrite().save(path) reloaded = EventNgramLookupTransformer.load(path) - assert reloaded.transform(search_level_df).collect() == expected + assert reloaded.transform(id_df).collect() == expected - def test_transform_emits_type_columns(self, search_level_df): + 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(search_level_df) + transformer.transform(id_df) .select("clicks_tokens", "clicks_tokens_types") .collect() ) @@ -287,7 +281,7 @@ def test_transform_emits_type_columns(self, search_level_df): type_lookup[t] for t in row["clicks_tokens"] ] - def test_output_dtype_applies_to_the_type_columns(self, search_level_df): + 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( @@ -296,7 +290,7 @@ def test_output_dtype_applies_to_the_type_columns(self, search_level_df): outputDtype="float", ) - schema = transformer.transform(search_level_df).schema + schema = transformer.transform(id_df).schema assert ( schema["clicks_tokens"].dataType == schema["clicks_tokens_types"].dataType @@ -311,14 +305,14 @@ def test_construct_layer_info_appends_type_outputs(self): assert info["inputs"] == ["clicks_ids"] assert info["outputs"] == ["clicks_tokens", "clicks_tokens_types"] - def test_spark_tf_parity_with_types(self, search_level_df): + 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(search_level_df) + transformer.transform(id_df) .select("clicks_tokens", "clicks_tokens_types") .collect() ): diff --git a/tests/kamae/spark/utils/test_ngram_utils.py b/tests/kamae/spark/utils/test_ngram_utils.py index 610601e0..f0fc572e 100644 --- a/tests/kamae/spark/utils/test_ngram_utils.py +++ b/tests/kamae/spark/utils/test_ngram_utils.py @@ -18,6 +18,7 @@ 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, @@ -25,10 +26,9 @@ collect_ngrams_from_dataframe, log_vocabulary_by_length, tokenize_events, + tokenize_events_with_types, ) from kamae.spark.utils.ngram_worker_functions import ( - PAD_TOKEN_ID, - UNK_TOKEN_ID, encode_tuple, extract_ngrams_from_column_worker, extract_tuples_from_column_worker, @@ -142,7 +142,7 @@ def test_extracts_non_zero_tuples(self, col_value, 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, min_ngram_freq=1) + 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 == { @@ -151,22 +151,17 @@ def test_assigns_ids_by_descending_frequency_from_two(self): ("L0_1", "L1_2"): 4, } - def test_drops_ngrams_below_min_frequency(self): - counter = Counter({("L0_1",): 100, ("L1_2",): 5}) - ngram_to_id = build_vocabulary(counter, vocab_size=100, min_ngram_freq=10) - assert ngram_to_id == {("L0_1",): 2} - 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, min_ngram_freq=1) + 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, min_ngram_freq=1) + first = build_vocabulary(counter, vocab_size=100) second = build_vocabulary( - Counter({("L0_1",): 10, ("L1_2",): 10}), vocab_size=100, min_ngram_freq=1 + Counter({("L0_1",): 10, ("L1_2",): 10}), vocab_size=100 ) assert first == second @@ -274,6 +269,20 @@ def test_empty_lookup_table_maps_every_event_to_unk(self): [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): @@ -305,5 +314,5 @@ def test_build_vocabulary_reports_the_distribution(self, caplog): {("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, min_ngram_freq=1) + build_vocabulary(ngram_counter=counter, vocab_size=100) assert "Vocabulary by n-gram length" in caplog.text From 19e4f0b4da61fa085e4e2726801ec2e3beffaf5d Mon Sep 17 00:00:00 2001 From: Sbranikas Date: Mon, 28 Sep 2026 22:50:44 +0200 Subject: [PATCH 6/6] refactor: reject an id axis that is not a whole number of events in EventNgramLookupLayer --- .../tensorflow/layers/event_ngram_lookup.py | 31 ++++++++++++++----- .../spark/estimators/event_ngram_lookup.py | 3 +- .../spark/transformers/event_ngram_lookup.py | 2 ++ .../layers/test_event_ngram_lookup.py | 20 ++++++++++++ 4 files changed, 48 insertions(+), 8 deletions(-) diff --git a/src/kamae/keras/tensorflow/layers/event_ngram_lookup.py b/src/kamae/keras/tensorflow/layers/event_ngram_lookup.py index 487ec5bd..b604b779 100644 --- a/src/kamae/keras/tensorflow/layers/event_ngram_lookup.py +++ b/src/kamae/keras/tensorflow/layers/event_ngram_lookup.py @@ -102,13 +102,12 @@ class EventNgramLookupLayer(BaseLayer): 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 zero-padded or truncated to ``num_events * tuple_size`` - before being split into tuples. + preserved. The id axis is padded with all-zero (padding) events or truncated to + ``num_events * tuple_size`` before being split into tuples. - Unlike the Spark transformer, which raises when a row's ids are not a whole number - of events, the layer zero-pads a partial trailing event and looks it up. This is - deliberate: it keeps assertions out of the serving graph, so a malformed request - is tokenized rather than failed. + 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. @@ -230,7 +229,10 @@ def _tokenize(self, inputs: KerasTensor, num_events: int) -> KerasTensor: :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. + :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. """ @@ -241,6 +243,21 @@ def _tokenize(self, inputs: KerasTensor, num_events: int) -> KerasTensor: 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] diff --git a/src/kamae/spark/estimators/event_ngram_lookup.py b/src/kamae/spark/estimators/event_ngram_lookup.py index 68e3ce57..c0908ddb 100644 --- a/src/kamae/spark/estimators/event_ngram_lookup.py +++ b/src/kamae/spark/estimators/event_ngram_lookup.py @@ -232,7 +232,8 @@ def _fit(self, dataset: DataFrame) -> EventNgramLookupTransformer: :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``). + 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() diff --git a/src/kamae/spark/transformers/event_ngram_lookup.py b/src/kamae/spark/transformers/event_ngram_lookup.py index 2814310f..f43a1059 100644 --- a/src/kamae/spark/transformers/event_ngram_lookup.py +++ b/src/kamae/spark/transformers/event_ngram_lookup.py @@ -204,6 +204,8 @@ class EventNgramLookupTransformer( 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 diff --git a/tests/kamae/keras/tensorflow/layers/test_event_ngram_lookup.py b/tests/kamae/keras/tensorflow/layers/test_event_ngram_lookup.py index 040eac32..fc981827 100644 --- a/tests/kamae/keras/tensorflow/layers/test_event_ngram_lookup.py +++ b/tests/kamae/keras/tensorflow/layers/test_event_ngram_lookup.py @@ -109,6 +109,26 @@ 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])