text: from functools import partial
  
  import tensorflow as tf
  
  from detext.utils.parsing_utils import InputFtrType, TaskType, InternalFtrType
  
  DEEP_FTR_TYPES = {
      InputFtrType.QUERY_COLUMN_NAME,
      InputFtrType.DOC_TEXT_COLUMN_NAMES,
      InputFtrType.DOC_ID_COLUMN_NAMES,
      InputFtrType.USER_TEXT_COLUMN_NAMES,
      InputFtrType.USER_ID_COLUMN_NAMES
  }
  
  WIDE_FTR_TYPES = {
      InputFtrType.DENSE_FTRS_COLUMN_NAMES,
      InputFtrType.SPARSE_FTRS_COLUMN_NAMES,
      InputFtrType.SHALLOW_TOWER_SPARSE_FTRS_COLUMN_NAMES
  }
  
  MULTITASK_FTR_TYPES = {
      InputFtrType.TASK_ID_COLUMN_NAME
  }
  
  DOC_FTR_TYPES = {
      InputFtrType.DOC_TEXT_COLUMN_NAMES,
      InputFtrType.DOC_ID_COLUMN_NAMES,
      InputFtrType.DENSE_FTRS_COLUMN_NAMES,
      InputFtrType.SPARSE_FTRS_COLUMN_NAMES,
      InputFtrType.SHALLOW_TOWER_SPARSE_FTRS_COLUMN_NAMES
  }
  
  
  def get_wide_features_inputs_ranking(features, feature_type2name):
      """ Returns wide ranking features by parsing features from data input_fn
  
      :param features: dict containing features in data
      :param feature_type2name: dict mapping feature types to feature names
      """
      result = {}
      for ftr_type, ftr_name_lst in feature_type2name.items():
          if ftr_type not in WIDE_FTR_TYPES:
              continue
          assert isinstance(ftr_name_lst, list), "Wide features must accept a list of feature names"
          result[ftr_type] = [features[ftr_name] for ftr_name in ftr_name_lst]
  
      return {k: v for k, v in result.items() if v is not None}
  
  
  def get_deep_features_inputs_ranking(features, feature_type2name):
      """ Returns deep ranking features by parsing features from data input_fn
  
      :param features: dict containing the features in data
      :param feature_type2name: dict mapping feature types to feature names
      """
      result = {}
      for ftr_type, ftr_name_lst in feature_type2name.items():
          if ftr_type not in DEEP_FTR_TYPES:
              continue
          if isinstance(ftr_name_lst, list):
              # Feature type accepting a list of feature names. E.g., doc_text, doc_id, user_text, user_id, etc.
              result[ftr_type] = [features[ftr_name] for ftr_name in ftr_name_lst]
          else:
              # Feature type accepting only one feature name. E.g., query, label, weight
              result[ftr_type] = features[ftr_name_lst]
  
      return {k: v for k, v in result.items() if v is not None}
  
  
  def get_multitask_features_inputs(features, feature_type2name):
      """ Returns deep ranking features by parsing features from data input_fn
  
      :param features: dict containing the features in data
      :param feature_type2name: dict mapping feature types to feature names
      """
      result = {}
      for ftr_type, ftr_name_lst in feature_type2name.items():
          if ftr_type not in MULTITASK_FTR_TYPES:
              continue
          # Feature type accepting only one feature name
          result[ftr_type] = features[ftr_name_lst]
  
      return {k: v for k, v in result.items() if v is not None}
  
  
  def get_wide_features_inputs_classification(features, feature_type2name):
      """ Returns wide classification features by parsing features from data input_fn
  
      :param features: dict containing features in data
      :param feature_type2name: dict mapping feature types to feature names
      """
      result = {}
      for ftr_type, ftr_name_lst in feature_type2name.items():
          if ftr_type not in WIDE_FTR_TYPES:
              continue
          result[ftr_type] = []
          assert isinstance(ftr_name_lst, list), "Wide features must accept a list of feature names"
  
          for ftr_name in ftr_name_lst:
              feat = features[ftr_name]
              if ftr_type in DOC_FTR_TYPES:  # Expand the list dimension (axis=1) for classification
                  if ftr_type in [InputFtrType.SPARSE_FTRS_COLUMN_NAMES, InputFtrType.SHALLOW_TOWER_SPARSE_FTRS_COLUMN_NAMES]:
                      feat = tf.sparse.expand_dims(feat, axis=1, name=ftr_name)
                  else:
                      feat = tf.expand_dims(feat, axis=1, name=ftr_name)
              result[ftr_type].append(feat)
  
      return {k: v for k, v in result.items() if v is not None}
  
  
  def get_deep_features_inputs_classification(features, feature_type2name):
      """Returns deep classification features by parsing features from data input_fn
  
      :param features: dict containing features in data
      :param feature_type2name: dict mapping feature types to feature names
      """
      result = {}
      for ftr_type, ftr_name_lst in feature_type2name.items():
          if ftr_type not in DEEP_FTR_TYPES:
              continue
          result[ftr_type] = []
          if isinstance(ftr_name_lst, list):
              # Feature type accepting a list of feature names. E.g., doc_text, doc_id, user_text, user_id, etc.
              for ftr_name in ftr_name_lst:
                  feat = features[ftr_name]
                  if ftr_type in DOC_FTR_TYPES:  # Expand the list dimension (axis=1) for classification
                      feat = tf.expand_dims(feat, axis=1, name=ftr_name)
                  result[ftr_type].append(feat)
          else:
              # Feature type accepting only one feature name. E.g., query, label, weight
              feat = features[ftr_name_lst]
              result[ftr_type] = feat
  
      return {k: v for k, v in result.items() if v is not None}
  
  
  TASK_TYPE_TO_GET_WIDE_FEATURES_FN = {
      TaskType.RANKING: get_wide_features_inputs_ranking,
      TaskType.CLASSIFICATION: get_wide_features_inputs_classification,
      TaskType.BINARY_CLASSIFICATION: get_wide_features_inputs_classification
  }
  
  TASK_TYPE_TO_GET_DEEP_FEATURES_FN = {
      TaskType.RANKING: get_deep_features_inputs_ranking,
      TaskType.CLASSIFICATION: get_deep_features_inputs_classification,
      TaskType.BINARY_CLASSIFICATION: get_deep_features_inputs_classification
  }
  
  
  class FeatureNameTypeConverter(tf.keras.layers.Layer):
      """Converter that converts feature_name->tensor mapping to feature_type->tensor mapping
  
      By using feature type (static and known internally in DeText) instead of feature names (dynamic), the Detext layer design becomes independent of user
          feature naming and free from the trouble of keeping notes about feature_type->feature_names
      """
  
      def __init__(self, task_type, feature_type2name, **kwargs):
          super(FeatureNameTypeConverter, self).__init__(**kwargs)
          self._task_type = task_type
          self._feature_type2name = feature_type2name
  
          self.converters = {
              InternalFtrType.WIDE_FTR_BAG: partial(TASK_TYPE_TO_GET_WIDE_FEATURES_FN[task_type], feature_type2name=feature_type2name),
              InternalFtrType.DEEP_FTR_BAG: partial(TASK_TYPE_TO_GET_DEEP_FEATURES_FN[task_type], feature_type2name=feature_type2name),
              InternalFtrType.MULTITASK_FTR_BAG: partial(get_multitask_features_inputs, feature_type2name=feature_type2name)
          }
  
      def call(self, inputs, **kwargs):
          return {k: self.converters[k](v) for k, v in inputs.items()}
  