autogluon.multimodal 1.2.1b20250303__tar.gz → 1.2.1b20250305__tar.gz

{autogluon.multimodal-1.2.1b20250303 → autogluon.multimodal-1.2.1b20250305}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: autogluon.multimodal
-Version: 1.2.1b20250303
+Version: 1.2.1b20250305
 Summary: Fast and Accurate ML in 3 Lines of Code
 Home-page: https://github.com/autogluon/autogluon
 Author: AutoGluon Community

{autogluon.multimodal-1.2.1b20250303 → autogluon.multimodal-1.2.1b20250305}/setup.py

@@ -61,7 +61,7 @@ install_requires = ag.get_dependency_version_ranges(install_requires)
 
 tests_require = [
     "ruff",
-    "datasets>=2.10.0,<2.15.0",
+    "datasets>=2.16.0,<2.20.0",
     "onnx>=1.13.0,<1.16.2;platform_system=='Windows'",  # cap at 1.16.1 for issue https://github.com/onnx/onnx/issues/6267
     "onnx>=1.13.0,<1.18.0;platform_system!='Windows'",
     "onnxruntime>=1.17.0,<1.20.0",  # install for gpu system due to https://github.com/autogluon/autogluon/issues/3804
@@ -78,8 +78,8 @@ if __name__ == "__main__":
     setup_args["package_data"]["autogluon.multimodal"] = [
         "configs/data/*.yaml",
         "configs/model/*.yaml",
-        "configs/optimization/*.yaml",
-        "configs/environment/*.yaml",
+        "configs/optim/*.yaml",
+        "configs/env/*.yaml",
         "configs/distiller/*.yaml",
         "configs/matcher/*.yaml",
     ]

autogluon.multimodal-1.2.1b20250305/src/autogluon/multimodal/__init__.py

@@ -0,0 +1,10 @@
+from autogluon.common.utils.log_utils import _add_stream_handler
+
+try:
+    from .version import __version__
+except ImportError:
+    pass
+
+from .predictor import MultiModalPredictor
+
+_add_stream_handler()

{autogluon.multimodal-1.2.1b20250303 → autogluon.multimodal-1.2.1b20250305}/src/autogluon/multimodal/constants.py

@@ -59,12 +59,17 @@ SEMANTIC_SEGMENTATION_GT = "semantic_segmentation_gt"
 
 # Output keys
 LOGITS = "logits"
+ORI_LOGITS = "ori_logits"
+AUG_LOGITS = "aug_logits"
 TEMPLATE_LOGITS = "template_logits"
 LM_TARGET = "lm_target"
 LOSS = "loss"
 OUTPUT = "output"
 WEIGHT = "weight"
 FEATURES = "features"
+MULTIMODAL_FEATURES = "multimodal_features"  # used for the adapted multimodal features before the fusion module
+MULTIMODAL_FEATURES_PRE_AUG = "multimodal_features_pre_aug"
+MULTIMODAL_FEATURES_POST_AUG = "multimodal_features_post_aug"
 RAW_FEATURES = "raw_features"
 MASKS = "masks"
 PROBABILITY = "probability"
@@ -73,6 +78,8 @@ BBOX = "bbox"
 ROIS = "rois"
 SCORE = "score"
 LOGIT_SCALE = "logit_scale"
+VAE_MEAN = "vae_mean"
+VAE_VAR = "vae_var"
 
 # Loss
 MOE_LOSS = "moe_loss"
@@ -142,6 +149,7 @@ FM = "fm"
 MAE = "mae"
 BER = "ber"
 IOU = "iou"
+COVERAGE = "coverage"
 RETRIEVAL_METRICS = [NDCG, PRECISION, RECALL, MRR]
 METRIC_MODE_MAP = {
     ACC: MAX,
@@ -168,6 +176,7 @@ METRIC_MODE_MAP = {
     SM: MAX,
     IOU: MAX,
     BER: MIN,
+    COVERAGE: MAX,
 }
 
 MATCHING_METRICS = {
@@ -179,7 +188,7 @@ MATCHING_METRICS_WITHOUT_PROBLEM_TYPE = [RECALL, NDCG]
 
 EVALUATION_METRICS = {
     # Use evaluation metrics from METRICS for these types
-    BINARY: METRICS[BINARY].keys(),
+    BINARY: list(METRICS[BINARY].keys()) + [COVERAGE],
     MULTICLASS: METRICS[MULTICLASS].keys(),
     REGRESSION: METRICS[REGRESSION].keys(),
     OBJECT_DETECTION: DETECTION_METRICS,
@@ -197,6 +206,7 @@ VALIDATION_METRICS = {
 # Training status
 TRAIN = "train"
 VALIDATE = "validate"
+VAL = "val"
 TEST = "test"
 PREDICT = "predict"
 
@@ -217,11 +227,11 @@ Y_TRUE = "y_true"
 # Configuration keys
 MODEL = "model"
 DATA = "data"
-OPTIMIZATION = "optimization"
-ENVIRONMENT = "environment"
+OPTIM = "optim"
+ENV = "env"
 DISTILLER = "distiller"
 MATCHER = "matcher"
-VALID_CONFIG_KEYS = [MODEL, DATA, OPTIMIZATION, ENVIRONMENT, DISTILLER, MATCHER]
+VALID_CONFIG_KEYS = [MODEL, DATA, OPTIM, ENV, DISTILLER, MATCHER]
 
 # Image normalization mean and std. This is only to normalize images for the CLIP model.
 CLIP_IMAGE_MEAN = (0.48145466, 0.4578275, 0.40821073)
@@ -275,7 +285,7 @@ PEFT_STRATEGIES = list(set(PEFT_ADDITIVE_STRATEGIES) | set(PEFT_NON_ADDITIVE_STR
 # DeepSpeed constants
 DEEPSPEED_OFFLOADING = "deepspeed_stage_3_offload"
 DEEPSPEED_STRATEGY = "deepspeed"
-DEEPSPEED_MODULE = "autogluon.multimodal.optimization.deepspeed"
+DEEPSPEED_MODULE = "autogluon.multimodal.optim.deepspeed"
 DEEPSPEED_MIN_PL_VERSION = "1.7.1"
 
 # registered model keys. TODO: document how to add new models.
@@ -298,6 +308,7 @@ DOCUMENT_TRANSFORMER = "document_transformer"
 HF_MODELS = (HF_TEXT, T_FEW, CLIP, NER_TEXT, DOCUMENT_TRANSFORMER)
 MMLAB_MODELS = (MMDET_IMAGE, MMOCR_TEXT_DET, MMOCR_TEXT_RECOG)
 SAM = "sam"
+META_TRANSFORMER = "meta_transformer"
 
 # matcher loss type
 CONTRASTIVE_LOSS = "contrastive_loss"

{autogluon.multimodal-1.2.1b20250303 → autogluon.multimodal-1.2.1b20250305}/src/autogluon/multimodal/data/__init__.py

@@ -1,4 +1,3 @@
-from . import collator, infer_types, randaug, utils
 from .datamodule import BaseDataModule
 from .dataset import BaseDataset
 from .dataset_mmlab import MultiImageMixDataset
@@ -9,8 +8,9 @@ from .infer_types import (
     infer_rois_column_type,
     is_image_column,
 )
-from .label_encoder import CustomLabelEncoder, NerLabelEncoder
 from .mixup import MixupModule
+from .infer_types import infer_column_types, infer_output_shape, infer_problem_type, is_image_column, infer_ner_column_type
+from .label_encoder import CustomLabelEncoder, NerLabelEncoder
 from .preprocess_dataframe import MultiModalFeaturePreprocessor
 from .process_categorical import CategoricalProcessor
 from .process_document import DocumentProcessor
@@ -21,3 +21,15 @@ from .process_ner import NerProcessor
 from .process_numerical import NumericalProcessor
 from .process_semantic_seg_img import SemanticSegImageProcessor
 from .process_text import TextProcessor
+from .utils import (
+    create_data_processor,
+    create_fusion_data_processors,
+    data_to_df,
+    get_detected_data_types,
+    get_mixup,
+    infer_dtypes_by_model_names,
+    infer_scarcity_mode_by_data_size,
+    init_df_preprocessor,
+    split_train_tuning_data,
+    turn_on_off_feature_column_info,
+)

{autogluon.multimodal-1.2.1b20250303 → autogluon.multimodal-1.2.1b20250305}/src/autogluon/multimodal/data/dataset.py

@@ -4,7 +4,7 @@ from typing import Dict, List, Optional, Union
 import pandas as pd
 import torch
 
-from ..constants import AUTOMM, GET_ITEM_ERROR_RETRY
+from ..constants import GET_ITEM_ERROR_RETRY
 from .preprocess_dataframe import MultiModalFeaturePreprocessor
 from .utils import apply_data_processor, apply_df_preprocessor, get_per_sample_features
 
@@ -100,7 +100,7 @@ class BaseDataset(torch.utils.data.Dataset):
                 per_ret = apply_data_processor(
                     per_sample_features=per_sample_features,
                     data_processors=per_processors_group,
-                    feature_modalities=getattr(self, f"modality_types_{group_id}"),
+                    data_types=getattr(self, f"modality_types_{group_id}"),
                     is_training=self.is_training,
                 )
                 ret.update(per_ret)

{autogluon.multimodal-1.2.1b20250303 → autogluon.multimodal-1.2.1b20250305}/src/autogluon/multimodal/data/infer_types.py

@@ -19,7 +19,6 @@ from ..constants import (
     DOCUMENT_IMAGE,
     DOCUMENT_PDF,
     IDENTIFIER,
-    IMAGE,
     IMAGE_BASE64_STR,
     IMAGE_BYTEARRAY,
     IMAGE_PATH,
@@ -37,7 +36,6 @@ from ..constants import (
     TEXT,
     TEXT_NER,
 )
-from .utils import is_rois_input
 
 logger = logging.getLogger(__name__)
 
@@ -114,6 +112,22 @@ def is_categorical_column(
             return False
 
 
+def is_rois_input(sample):
+    """
+    check if a sample is rois for object detection
+
+    Parameters
+    ----------
+    sample
+        The sampled data.
+
+    Returns
+    -------
+    bool, whether a sample is rois for object detection
+    """
+    return isinstance(sample, list) and len(sample) and isinstance(sample[0], list) and len(sample[0]) == 5
+
+
 def is_rois_column(data: pd.Series) -> bool:
     """
     Identify if a column is one rois column.

{autogluon.multimodal-1.2.1b20250303 → autogluon.multimodal-1.2.1b20250305}/src/autogluon/multimodal/data/label_encoder.py

@@ -9,7 +9,7 @@ import pandas as pd
 from omegaconf import DictConfig, OmegaConf
 from sklearn.preprocessing import LabelEncoder
 
-from ..constants import AUTOMM, END_OFFSET, ENTITY_GROUP, NER_ANNOTATION, PROBABILITY, START_OFFSET
+from ..constants import END_OFFSET, ENTITY_GROUP, PROBABILITY, START_OFFSET
 
 logger = logging.getLogger(__name__)
 
@@ -137,12 +137,12 @@ class NerLabelEncoder:
         transformed_y
             A list of word level annotations.
         """
-        from .utils import process_ner_annotations
+        from .process_ner import NerProcessor
 
         all_annotations, _ = self.extract_ner_annotations(y)
         transformed_y = []
         for annotation, text_snippet in zip(all_annotations, x.items()):
-            word_label, _, _, _ = process_ner_annotations(
+            word_label, _, _, _ = NerProcessor.process_ner_annotations(
                 annotation, text_snippet[-1], self.entity_map, tokenizer, is_eval=True
             )
             word_label_invers = []

{autogluon.multimodal-1.2.1b20250303/src/autogluon/multimodal/utils → autogluon.multimodal-1.2.1b20250305/src/autogluon/multimodal/data}/nlpaug.py

@@ -78,14 +78,14 @@ class InsertPunctuation(Augmenter):
         new = " ".join(new)
         return new
 
-    @classmethod
-    def clean(cls, data):
+    @staticmethod
+    def clean(data):
         if isinstance(data, list):
             return [d.strip() if d else d for d in data]
         return data.strip()
 
-    @classmethod
-    def is_duplicate(cls, dataset, data):
+    @staticmethod
+    def is_duplicate(dataset, data):
         for d in dataset:
             if d == data:
                 return True

{autogluon.multimodal-1.2.1b20250303 → autogluon.multimodal-1.2.1b20250305}/src/autogluon/multimodal/data/preprocess_dataframe.py

@@ -14,17 +14,14 @@ from sklearn.preprocessing import MinMaxScaler, StandardScaler
 from autogluon.features import CategoryFeatureGenerator
 
 from ..constants import (
-    AUTOMM,
     CATEGORICAL,
     DOCUMENT,
-    DOCUMENT_IMAGE,
     IDENTIFIER,
     IMAGE,
     IMAGE_BASE64_STR,
     IMAGE_BYTEARRAY,
     IMAGE_PATH,
     LABEL,
-    NER,
     NER_ANNOTATION,
     NULL,
     NUMERICAL,
@@ -73,19 +70,17 @@ class MultiModalFeaturePreprocessor(TransformerMixin, BaseEstimator):
 
         if label_column:
             if label_generator is None:
-                self._label_generator = CustomLabelEncoder(
-                    positive_class=OmegaConf.select(config, "pos_label", default=None)
-                )
+                self._label_generator = CustomLabelEncoder(positive_class=config.pos_label)
             else:
                 self._label_generator = label_generator
 
             # Scaler used for numerical labels
-            numerical_label_preprocessing = OmegaConf.select(config, "label.numerical_label_preprocessing")
+            numerical_label_preprocessing = config.label.numerical_preprocessing
             if numerical_label_preprocessing == "minmaxscaler":
                 self._label_scaler = MinMaxScaler()
             elif numerical_label_preprocessing == "standardscaler":
                 self._label_scaler = StandardScaler()
-            elif numerical_label_preprocessing is None or numerical_label_preprocessing.lower() == "none":
+            elif numerical_label_preprocessing is None:
                 self._label_scaler = StandardScaler(with_mean=False, with_std=False)
             else:
                 raise ValueError(
@@ -135,8 +130,7 @@ class MultiModalFeaturePreprocessor(TransformerMixin, BaseEstimator):
         # Some columns will be ignored
         self._ignore_columns_set = set()
         self._text_feature_names = []
-        self._categorical_feature_names = []
-        self._categorical_num_categories = []
+        self._categorical_num_categories = dict()
         self._numerical_feature_names = []
         self._image_feature_names = []
         self._rois_feature_names = []
@@ -154,10 +148,7 @@ class MultiModalFeaturePreprocessor(TransformerMixin, BaseEstimator):
 
     @property
     def image_path_names(self):
-        if hasattr(self, "_image_path_names"):
-            return self._image_path_names
-        else:
-            return [col_name for col_name in self._image_feature_names if self._column_types[col_name] == IMAGE_PATH]
+        return [col_name for col_name in self._image_feature_names if self._column_types[col_name] == IMAGE_PATH]
 
     @property
     def rois_feature_names(self):
@@ -173,7 +164,7 @@ class MultiModalFeaturePreprocessor(TransformerMixin, BaseEstimator):
 
     @property
     def image_feature_names(self):
-        return self._image_path_names if hasattr(self, "_image_path_names") else self._image_feature_names
+        return self._image_feature_names
 
     @property
     def text_feature_names(self):
@@ -181,12 +172,21 @@ class MultiModalFeaturePreprocessor(TransformerMixin, BaseEstimator):
 
     @property
     def categorical_feature_names(self):
-        return self._categorical_feature_names
+        return list(self.categorical_num_categories.keys())
 
     @property
     def numerical_feature_names(self):
         return self._numerical_feature_names
 
+    @property
+    def numerical_fill_values(self):
+        ret = dict()
+        for col_name in self._numerical_feature_names:
+            generator = self._feature_generators[col_name]
+            ret[col_name] = generator.transform(np.full([1, 1], np.nan))[:, 0][0]
+
+        return ret
+
     @property
     def document_feature_names(self):
         # Added for backward compatibility.
@@ -216,17 +216,12 @@ class MultiModalFeaturePreprocessor(TransformerMixin, BaseEstimator):
 
     @property
     def required_feature_names(self):
-        image_feature_names = (
-            self._image_path_names if hasattr(self, "_image_path_names") else self._image_feature_names
-        )
-        rois_feature_names = self._rois_feature_names if hasattr(self, "_rois_feature_names") else []
-
         return (
-            image_feature_names
+            self._image_feature_names
             + self._text_feature_names
             + self._numerical_feature_names
-            + self._categorical_feature_names
-            + rois_feature_names
+            + self.categorical_feature_names
+            + self._rois_feature_names
         )
 
     @property
@@ -268,16 +263,13 @@ class MultiModalFeaturePreprocessor(TransformerMixin, BaseEstimator):
 
     def get_column_names(self, modality: str):
         if modality.startswith(IMAGE):
-            if hasattr(self, "_image_path_names"):
-                return self._image_path_names
-            else:
-                return self._image_feature_names
+            return self._image_feature_names
         elif modality == ROIS:
             return self._rois_feature_names
         elif modality == TEXT:
             return self._text_feature_names
         elif modality == CATEGORICAL:
-            return self._categorical_feature_names
+            return self.categorical_feature_names
         elif modality == NUMERICAL:
             return self._numerical_feature_names
         elif modality.startswith(DOCUMENT):
@@ -344,8 +336,7 @@ class MultiModalFeaturePreprocessor(TransformerMixin, BaseEstimator):
                         continue
                     num_categories = len(generator.category_map[col_name])
                     # Add one unknown category
-                    self._categorical_num_categories.append(num_categories + 1)
-                    self._categorical_feature_names.append(col_name)
+                    self._categorical_num_categories[col_name] = num_categories + 1
             elif col_type == NUMERICAL:
                 processed_data = pd.to_numeric(col_value)
                 if len(processed_data.unique()) == 1:
@@ -392,7 +383,7 @@ class MultiModalFeaturePreprocessor(TransformerMixin, BaseEstimator):
         elif self.label_type == NUMERICAL:
             y = pd.to_numeric(y).to_numpy()
             self._label_scaler.fit(np.expand_dims(y, axis=-1))
-        elif self.label_type == ROIS or self.label_type == SEMANTIC_SEGMENTATION_GT:
+        elif self.label_type in [ROIS, SEMANTIC_SEGMENTATION_GT]:
             pass  # Do nothing. TODO: Shall we call fit here?
         elif self.label_type == NER_ANNOTATION:
             # If there are ner annotations and text columns but no NER feature columns,
@@ -426,6 +417,24 @@ class MultiModalFeaturePreprocessor(TransformerMixin, BaseEstimator):
         if y is not None:
             self._fit_y(y=y, X=X)
 
+    @staticmethod
+    def convert_categorical_to_text(col_value: pd.Series, template: str, col_name: str):
+        # TODO: do we need to consider whether categorical values are valid text?
+        col_value = col_value.astype("object")
+        if template == "direct":
+            processed_data = col_value.apply(lambda ele: "" if pd.isnull(ele) else str(ele))
+        elif template == "list":
+            processed_data = col_value.apply(lambda ele: "" if pd.isnull(ele) else col_name + ": " + str(ele))
+        elif template == "text":
+            processed_data = col_value.apply(lambda ele: "" if pd.isnull(ele) else col_name + " is " + str(ele))
+        elif template == "latex":
+            processed_data = col_value.apply(lambda ele: "" if pd.isnull(ele) else str(ele) + " & ")
+        else:
+            raise ValueError(
+                f"Unsupported template {template} for converting categorical data into text. Select one from: ['direct', 'list', 'text', 'latex']."
+            )
+        return processed_data
+
     def transform_text(
         self,
         df: pd.DataFrame,
@@ -455,10 +464,15 @@ class MultiModalFeaturePreprocessor(TransformerMixin, BaseEstimator):
         for col_name in self._text_feature_names:
             col_value = df[col_name]
             col_type = self._column_types[col_name]
-            if col_type == TEXT or col_type == CATEGORICAL:
-                # TODO: do we need to consider whether categorical values are valid text?
+            if col_type == TEXT:
                 col_value = col_value.astype("object")
                 processed_data = col_value.apply(lambda ele: "" if pd.isnull(ele) else str(ele))
+            elif col_type == CATEGORICAL:
+                processed_data = self.convert_categorical_to_text(
+                    col_value=col_value,
+                    template=self._config.categorical.convert_to_text_template,
+                    col_name=col_name,
+                )
             elif col_type == NUMERICAL:
                 processed_data = pd.to_numeric(col_value).apply("{:.3f}".format)
             elif col_type == f"{TEXT}_{IDENTIFIER}":
@@ -710,7 +724,7 @@ class MultiModalFeaturePreprocessor(TransformerMixin, BaseEstimator):
             self._fit_called or self._fit_x_called
         ), "You will need to first call preprocessor.fit before calling preprocessor.transform_categorical."
         categorical_features = {}
-        for col_name, num_category in zip(self._categorical_feature_names, self._categorical_num_categories):
+        for col_name, num_category in self._categorical_num_categories.items():
             col_value = df[col_name]
             processed_data = col_value.astype("category")
             generator = self._feature_generators[col_name]
@@ -757,7 +771,7 @@ class MultiModalFeaturePreprocessor(TransformerMixin, BaseEstimator):
         elif self.label_type == NUMERICAL:
             y = pd.to_numeric(y_df).to_numpy()
             y = self._label_scaler.transform(np.expand_dims(y, axis=-1))[:, 0].astype(np.float32)
-        elif self.label_type == ROIS or self.label_type == SEMANTIC_SEGMENTATION_GT:
+        elif self.label_type in [ROIS, SEMANTIC_SEGMENTATION_GT]:
             y = y_df.to_list()
         elif self.label_type == NER_ANNOTATION:
             y = self._label_generator.transform(y_df)
@@ -866,8 +880,11 @@ class MultiModalFeaturePreprocessor(TransformerMixin, BaseEstimator):
         ), "You will need to first call preprocessor.fit_y() before calling preprocessor.transform_prediction."
 
         if self.label_type == CATEGORICAL:
-            assert y_pred.shape[1] >= 2
-            y_pred = y_pred.argmax(axis=1)
+            assert len(y_pred.shape) <= 2
+            if len(y_pred.shape) == 2 and y_pred.shape[1] >= 2:
+                y_pred = y_pred.argmax(axis=1)
+            else:
+                y_pred = (y_pred > 0.5).astype(int)
             # Transform the predicted label back to the original space (e.g., string values)
             if inverse_categorical:
                 y_pred = self._label_generator.inverse_transform(y_pred)

{autogluon.multimodal-1.2.1b20250303 → autogluon.multimodal-1.2.1b20250305}/src/autogluon/multimodal/data/process_categorical.py

@@ -1,11 +1,14 @@
+import logging
+import random
 from typing import Any, Dict, List, Optional, Union
 
-import numpy as np
 from torch import nn
 
 from ..constants import CATEGORICAL, COLUMN
 from .collator import StackCollator, TupleCollator
 
+logger = logging.getLogger(__name__)
+
 
 class CategoricalProcessor:
     """
@@ -18,6 +21,7 @@ class CategoricalProcessor:
         self,
         model: nn.Module,
         requires_column_info: bool = False,
+        dropout: Optional[float] = 0,
     ):
         """
         Parameters
@@ -27,8 +31,16 @@ class CategoricalProcessor:
         requires_column_info
             Whether to require feature column information in dataloader.
         """
+        logger.debug(f"initializing categorical processor for model {model.prefix}")
         self.prefix = model.prefix
         self.requires_column_info = requires_column_info
+        self.num_categories = model.num_categories
+        self.dropout = dropout
+        assert 0 <= self.dropout <= 1
+        if self.dropout > 0:
+            logger.debug(f"categorical value dropout probability: {self.dropout}")
+            fill_values = {k: v - 1 for k, v in self.num_categories.items()}
+            logger.debug(f"dropped values will be replaced by {fill_values}")
 
     @property
     def categorical_key(self):
@@ -60,6 +72,7 @@ class CategoricalProcessor:
     def process_one_sample(
         self,
         categorical_features: Dict[str, int],
+        is_training: bool,
     ) -> Dict:
         """
         Process one sample's categorical features. Assume the categorical features
@@ -69,6 +82,8 @@ class CategoricalProcessor:
         ----------
         categorical_features
             Categorical features of one sample.
+        is_training
+            Whether to do processing in the training mode.
 
         Returns
         -------
@@ -80,6 +95,17 @@ class CategoricalProcessor:
             for i, col_name in enumerate(categorical_features.keys()):
                 ret[f"{self.categorical_column_prefix}_{col_name}"] = i
 
+        if is_training and self.dropout > 0:
+            categorical_features_copy = dict()
+            for k, v in categorical_features.items():
+                if random.uniform(0, 1) <= self.dropout:
+                    categorical_features_copy[k] = self.num_categories[k] - 1
+                else:
+                    categorical_features_copy[k] = v
+            categorical_features = categorical_features_copy
+
+        # make sure keys are in the same order
+        assert list(categorical_features.keys()) == list(self.num_categories.keys())
         ret[self.categorical_key] = list(categorical_features.values())
 
         return ret
@@ -87,7 +113,7 @@ class CategoricalProcessor:
     def __call__(
         self,
         categorical_features: Dict[str, int],
-        feature_modalities: Dict[str, Union[int, float, list]],
+        sub_dtypes: Dict[str, str],
         is_training: bool,
     ) -> Dict:
         """
@@ -97,13 +123,16 @@ class CategoricalProcessor:
         ----------
         categorical_features
             Categorical features of one sample.
-        feature_modalities
-            The modality of the feature columns.
+        sub_dtypes
+            The sub data types of all categorical columns.
         is_training
-            Whether to do processing in the training mode. This unused flag is for the API compatibility.
+            Whether to do processing in the training mode.
 
         Returns
         -------
         A dictionary containing one sample's processed categorical features.
         """
-        return self.process_one_sample(categorical_features)
+        return self.process_one_sample(
+            categorical_features=categorical_features,
+            is_training=is_training,
+        )