feat(KDP): global embedding for numeric features option added

Author: Gandalfdore

kdp/custom_layers.py

@@ -2046,17 +2046,22 @@ def build(self, input_shape):
             init_min_tensor = tf.fill([self.num_features], init_min_tensor)
         if init_max_tensor.shape.ndims == 0:
             init_max_tensor = tf.fill([self.num_features], init_max_tensor)
-        # Convert tensors to numpy arrays, which are acceptable by tf.constant_initializer.
-        init_min_value = (
-            init_min_tensor.numpy()
-            if hasattr(init_min_tensor, "numpy")
-            else init_min_tensor
-        )
-        init_max_value = (
-            init_max_tensor.numpy()
-            if hasattr(init_max_tensor, "numpy")
-            else init_max_tensor
-        )
+
+        if tf.executing_eagerly():
+            init_min_value = init_min_tensor.numpy()
+            init_max_value = init_max_tensor.numpy()
+        else:
+            # Fallback: if not executing eagerly, force conversion to list
+            init_min_value = (
+                init_min_tensor.numpy().tolist()
+                if hasattr(init_min_tensor, "numpy")
+                else self.init_min
+            )
+            init_max_value = (
+                init_max_tensor.numpy().tolist()
+                if hasattr(init_max_tensor, "numpy")
+                else self.init_max
+            )
 
         self.learned_min = self.add_weight(
             name="learned_min",
@@ -2133,3 +2138,106 @@ def get_config(self):
             }
         )
         return config
+
+
+class GlobalAdvancedNumericalEmbedding(tf.keras.layers.Layer):
+    """
+    Global AdvancedNumericalEmbedding processes concatenated numeric features.
+    It applies an inner AdvancedNumericalEmbedding over the flattened input and then
+    performs global pooling (average or max) to produce a compact representation.
+    """
+
+    def __init__(
+        self,
+        global_embedding_dim: int,
+        global_mlp_hidden_units: int,
+        global_num_bins: int,
+        global_init_min,
+        global_init_max,
+        global_dropout_rate: float,
+        global_use_batch_norm: bool,
+        global_pooling: str = "average",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.global_embedding_dim = global_embedding_dim
+        self.global_mlp_hidden_units = global_mlp_hidden_units
+        self.global_num_bins = global_num_bins
+
+        # Ensure initializer parameters are Python scalars, lists, or numpy arrays.
+        if not isinstance(global_init_min, (list, tuple, np.ndarray)):
+            try:
+                global_init_min = float(global_init_min)
+            except Exception:
+                raise ValueError(
+                    "init_min must be a Python scalar, list, tuple or numpy array"
+                )
+        if not isinstance(global_init_max, (list, tuple, np.ndarray)):
+            try:
+                global_init_max = float(global_init_max)
+            except Exception:
+                raise ValueError(
+                    "init_max must be a Python scalar, list, tuple or numpy array"
+                )
+        self.global_init_min = global_init_min
+        self.global_init_max = global_init_max
+        self.global_dropout_rate = global_dropout_rate
+        self.global_use_batch_norm = global_use_batch_norm
+        self.global_pooling = global_pooling
+
+        # Use the existing advanced numerical embedding block
+        self.inner_embedding = AdvancedNumericalEmbedding(
+            embedding_dim=self.global_embedding_dim,
+            mlp_hidden_units=self.global_mlp_hidden_units,
+            num_bins=self.global_num_bins,
+            init_min=self.global_init_min,
+            init_max=self.global_init_max,
+            dropout_rate=self.global_dropout_rate,
+            use_batch_norm=self.global_use_batch_norm,
+            name="global_numeric_emebedding",
+        )
+        if self.global_pooling == "average":
+            self.global_pooling_layer = tf.keras.layers.GlobalAveragePooling1D(
+                name="global_avg_pool"
+            )
+        elif self.global_pooling == "max":
+            self.global_pooling_layer = tf.keras.layers.GlobalMaxPooling1D(
+                name="global_max_pool"
+            )
+        else:
+            raise ValueError(f"Unsupported pooling method: {self.global_pooling}")
+
+    def call(self, inputs: tf.Tensor, training: bool = False) -> tf.Tensor:
+        """
+        Expects inputs with shape (batch, ...) and flattens them (except for the batch dim).
+        Then, the inner embedding produces a 3D output (batch, num_features, embedding_dim),
+        which is finally pooled to yield (batch, embedding_dim).
+        """
+        # If inputs have more than 2 dimensions, flatten them (except for batch dimension).
+        if len(inputs.shape) > 2:
+            inputs = tf.reshape(inputs, (tf.shape(inputs)[0], -1))
+        # Pass through the inner advanced embedding.
+        x_embedded = self.inner_embedding(inputs, training=training)
+        # Global pooling over numeric features axis.
+        x_pooled = self.global_pooling_layer(x_embedded)
+        return x_pooled
+
+    def compute_output_shape(self, input_shape):
+        # Regardless of the input shape, the output shape is (batch_size, embedding_dim)
+        return (input_shape[0], self.global_embedding_dim)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "global_embedding_dim": self.global_embedding_dim,
+                "global_mlp_hidden_units": self.global_mlp_hidden_units,
+                "global_num_bins": self.global_num_bins,
+                "global_init_min": self.global_init_min,
+                "global_init_max": self.global_init_max,
+                "global_dropout_rate": self.global_dropout_rate,
+                "global_use_batch_norm": self.global_use_batch_norm,
+                "global_pooling": self.global_pooling,
+            }
+        )
+        return config

kdp/processor.py

@@ -12,6 +12,7 @@
 import tensorflow as tf
 from loguru import logger
 
+from kdp.custom_layers import GlobalAdvancedNumericalEmbedding
 from kdp.features import (
     CategoricalFeature,
     CategoryEncodingOptions,
@@ -200,6 +201,15 @@ def __init__(
         init_max: float = 3.0,
         dropout_rate: float = 0.1,
         use_batch_norm: bool = True,
+        use_global_numerical_embedding: bool = False,
+        global_embedding_dim: int = 8,
+        global_mlp_hidden_units: int = 16,
+        global_num_bins: int = 10,
+        global_init_min: float = -3.0,
+        global_init_max: float = 3.0,
+        global_dropout_rate: float = 0.1,
+        global_use_batch_norm: bool = True,
+        global_pooling: str = "average",
     ) -> None:
         """Initialize a preprocessing model.
 
@@ -282,6 +292,17 @@ def __init__(
         self.dropout_rate = dropout_rate
         self.use_batch_norm = use_batch_norm
 
+        # advanced global numerical embedding control
+        self.use_global_numerical_embedding = use_global_numerical_embedding
+        self.global_embedding_dim = global_embedding_dim
+        self.global_mlp_hidden_units = global_mlp_hidden_units
+        self.global_num_bins = global_num_bins
+        self.global_init_min = global_init_min
+        self.global_init_max = global_init_max
+        self.global_dropout_rate = global_dropout_rate
+        self.global_use_batch_norm = global_use_batch_norm
+        self.global_pooling = global_pooling
+
         # PLACEHOLDERS
         self.preprocessors = {}
         self.inputs = {}
@@ -1046,6 +1067,17 @@ def _prepare_outputs(self) -> None:
                     name="ConcatenateNumeric",
                     axis=-1,
                 )(numeric_features)
+                if self.use_global_numerical_embedding:
+                    concat_num = GlobalAdvancedNumericalEmbedding(
+                        global_embedding_dim=self.global_embedding_dim,
+                        global_mlp_hidden_units=self.global_mlp_hidden_units,
+                        global_num_bins=self.global_num_bins,
+                        global_init_min=self.global_init_min,
+                        global_init_max=self.global_init_max,
+                        global_dropout_rate=self.global_dropout_rate,
+                        global_use_batch_norm=self.global_use_batch_norm,
+                        global_pooling=self.global_pooling,
+                    )(concat_num)
             else:
                 concat_num = None