code: def preprocess_keras_inputs(
      context_keras_inputs: Dict[
          str, Union[tf.Tensor, tf.RaggedTensor]
      ],
      example_keras_inputs: Dict[
          str, Union[tf.Tensor, tf.RaggedTensor]
      ],
      mask: tf.Tensor,
  ) -> Tuple[Dict[str, tf.Tensor], Dict[str, tf.Tensor]]:
      """Preprocess context and example inputs.
  
      Args:
        context_keras_inputs: Mapping context feature names to tf.keras.Input.
        example_keras_inputs: Mapping example feature names to tf.keras.Input.
        mask: 2-D mask for valid examples of shape [batch_size, list_size].
  
      Returns:
        preprocessed_context_features: Mapping from context feature names
          to preprocessed dense 2-D tensors of shape [batch_size, ...].
        preprocessed_example_features: Mapping from example feature names
          to preprocessed dense 3-D tensors of shape [batch_size, list_size, ...].
      """
      ragged_to_dense_layer = tf.keras.layers.Lambda(
          lambda x: x.to_tensor()
          if isinstance(x, tf.RaggedTensor)
          else x
      )
  
      # Create shared embedding layer for context and example features.
      shared_embedding_layer = tf.keras.layers.Embedding(
          input_dim=FLAGS.vocab_size,
          output_dim=FLAGS.embedding_dimension,
          embeddings_initializer=None,
          embeddings_constraint=None,
      )
      id_lookup_layer = (
          tf.keras.layers.experimental.preprocessing.StringLookup(
              max_tokens=FLAGS.vocab_size,
              vocabulary=FLAGS.vocab_file_path,
              num_oov_indices=0,
              oov_token="[UNK#]",
              mask_token=None,
          )
      )
  
      def _embedding(x):
          ids = id_lookup_layer(x)
          embedded_tensor = shared_embedding_layer(ids)
          # Using mean reduction here. User can specify any kind of reduction,
          # e.g. sum. The reduction is along the ragged axis for variable number of
          # tokens. This reduction transforms a ragged feature of dense shape:
          # [batch_size, list_size, num_tokens, embedding_size] ->
          # [batch_size, list_size, embedding_size].
          embedded_tensor = tf.reduce_mean(embedded_tensor, axis=-2)
          return ragged_to_dense_layer(embedded_tensor)
  
      preprocessed_context_features, preprocessed_example_features = (
          {},
          {},
      )
      preprocessed_context_features.update(
          {
              k: _embedding(context_keras_inputs[k])
              for k in _CONTEXT_RAGGED_FEATURE_KEYS
          }
      )
      preprocessed_example_features.update(
          {
              k: _embedding(example_keras_inputs[k])
              for k in _EXAMPLE_RAGGED_FEATURE_KEYS
          }
      )
  
      # Document interaction attention layer.
      if FLAGS.use_document_interaction:
          concat_tensor = tfr.keras.layers.ConcatFeatures()(
              inputs=(
                  preprocessed_context_features,
                  preprocessed_example_features,
                  mask,
              )
          )
          din_layer = tfr.keras.layers.DocumentInteractionAttention(
              num_heads=FLAGS.num_attention_heads,
              head_size=FLAGS.head_size,
              num_layers=FLAGS.num_attention_layers,
              dropout=FLAGS.dropout_rate,
          )
          preprocessed_example_features[
              "document_interaction_embedding"
          ] = din_layer(inputs=(concat_tensor, mask))
  
      return (
          preprocessed_context_features,
          preprocessed_example_features,
      )
  
file_path: ranking/tensorflow_ranking/examples/keras/antique_kpl_din.py
imports: import os
  from typing import Dict, Optional, Tuple, Union
  from absl import app
  from absl import flags
  from absl import logging
  import tensorflow
  import tensorflow_ranking
label_desc: Use image node_id 4 for calling a global function with example usage: preprocess_keras_inputs(context_keras_inputs, example_keras_inputs, mask) and returns: preprocessed_context_features,  preprocessed_example_features, ragged_to_dense_layer