autogluon.multimodal 1.2.1b20250303__py3-none-any.whl → 1.2.1b20250304__py3-none-any.whl

autogluon/multimodal/models/ft_transformer.py

@@ -1,6 +1,7 @@
+import logging
 import os
 import tempfile
-from typing import List, Optional
+from typing import Dict, List, Optional
 
 import torch
 from torch import Tensor, nn
@@ -9,6 +10,8 @@ from ..constants import CATEGORICAL, FEATURES, LABEL, LOGITS, NUMERICAL
 from .custom_transformer import CLSToken, Custom_Transformer, _TokenInitialization
 from .utils import init_weights
 
+logger = logging.getLogger(__name__)
+
 
 class CategoricalFeatureTokenizer(nn.Module):
     """
@@ -21,7 +24,7 @@ class CategoricalFeatureTokenizer(nn.Module):
     def __init__(
         self,
         num_categories: List[int],
-        d_token: int,
+        token_dim: int,
         bias: Optional[bool] = True,
         initialization: Optional[str] = "normal",
     ) -> None:
@@ -30,7 +33,7 @@ class CategoricalFeatureTokenizer(nn.Module):
         ----------
         num_categories:
             A list of integers. Each one is the number of categories in one categorical column.
-        d_token:
+        token_dim:
             The size of one token.
         bias:
             If `True`, for each feature, an additional trainable vector will be added to the
@@ -51,21 +54,21 @@ class CategoricalFeatureTokenizer(nn.Module):
         category_offsets = torch.tensor([0] + num_categories[:-1]).cumsum(0)
 
         self.register_buffer("category_offsets", category_offsets, persistent=False)
-        self.embeddings = nn.Embedding(sum(num_categories), d_token)
-        self.bias = nn.Parameter(Tensor(len(num_categories), d_token)) if bias else None
+        self.embeddings = nn.Embedding(sum(num_categories), token_dim)
+        self.bias = nn.Parameter(Tensor(len(num_categories), token_dim)) if bias else None
         initialization_ = _TokenInitialization.from_str(initialization)
 
         for parameter in [self.embeddings.weight, self.bias]:
             if parameter is not None:
-                initialization_.apply(parameter, d_token)
+                initialization_.apply(parameter, token_dim)
 
     @property
-    def n_tokens(self) -> int:
+    def num_tokens(self) -> int:
         """The number of tokens."""
         return len(self.num_categories)
 
     @property
-    def d_token(self) -> int:
+    def token_dim(self) -> int:
         """The size of one token."""
         return self.embeddings.embedding_dim
 
@@ -190,7 +193,7 @@ class NumericalFeatureTokenizer(nn.Module):
     def __init__(
         self,
         in_features: int,
-        d_token: int,
+        token_dim: int,
         bias: Optional[bool] = True,
         initialization: Optional[str] = "normal",
     ):
@@ -199,7 +202,7 @@ class NumericalFeatureTokenizer(nn.Module):
         ----------
         in_features:
             Dimension of input features i.e. the number of continuous (scalar) features
-        d_token:
+        token_dim:
             The size of one token.
         bias:
             If `True`, for each feature, an additional trainable vector will be added to the
@@ -216,19 +219,19 @@ class NumericalFeatureTokenizer(nn.Module):
         super().__init__()
 
         initialization_ = _TokenInitialization.from_str(initialization)
-        self.weight = nn.Parameter(Tensor(in_features, d_token))
-        self.bias = nn.Parameter(Tensor(in_features, d_token)) if bias else None
+        self.weight = nn.Parameter(Tensor(in_features, token_dim))
+        self.bias = nn.Parameter(Tensor(in_features, token_dim)) if bias else None
         for parameter in [self.weight, self.bias]:
             if parameter is not None:
-                initialization_.apply(parameter, d_token)
+                initialization_.apply(parameter, token_dim)
 
     @property
-    def n_tokens(self) -> int:
+    def num_tokens(self) -> int:
         """The number of tokens."""
         return len(self.weight)
 
     @property
-    def d_token(self) -> int:
+    def token_dim(self) -> int:
         """The size of one token."""
         return self.weight.shape[1]
 
@@ -266,7 +269,7 @@ class AutoDis(nn.Module):
         super().__init__()
         self.first_layer = NumericalFeatureTokenizer(
             in_features=in_features,
-            d_token=n_meta_embeddings,
+            token_dim=n_meta_embeddings,
             bias=False,
             initialization="uniform",
         )
@@ -357,7 +360,7 @@ class NumEmbeddings(nn.Module):
         if embedding_arch[0] == "linear":
             layers.append(
                 NumericalFeatureTokenizer(
-                    in_features=in_features, d_token=d_embedding, bias=True, initialization="normal"
+                    in_features=in_features, token_dim=d_embedding, bias=True, initialization="normal"
                 )
             )
         elif embedding_arch[0] == "positional":
@@ -406,13 +409,13 @@ class NumEmbeddings(nn.Module):
         self.layers = nn.Sequential(*layers)
 
     @property
-    def n_tokens(self) -> int:
+    def num_tokens(self) -> int:
         """The number of tokens."""
         y = self.forward(torch.ones(1, self.in_features))
         return y.shape[1]
 
     @property
-    def d_token(self) -> int:
+    def token_dim(self) -> int:
         """The size of one token."""
         y = self.forward(torch.ones(1, self.in_features))
         return y.shape[-1]
@@ -432,7 +435,8 @@ class FT_Transformer(nn.Module):
         self,
         prefix: str,
         num_numerical_columns: int,
-        num_categories: List[int],
+        num_categories: Dict,
+        numerical_fill_values: Dict,
         embedding_arch: List[str],
         token_dim: int,
         hidden_size: Optional[int] = 192,
@@ -441,7 +445,7 @@ class FT_Transformer(nn.Module):
         token_bias: Optional[bool] = True,
         token_initialization: Optional[str] = "normal",
         num_blocks: Optional[int] = 0,
-        attention_n_heads: Optional[int] = 8,
+        attention_num_heads: Optional[int] = 8,
         attention_initialization: Optional[str] = "kaiming",
         attention_normalization: Optional[str] = "layer_norm",
         attention_dropout: Optional[str] = 0.2,
@@ -485,7 +489,7 @@ class FT_Transformer(nn.Module):
             Must be one of `['uniform', 'normal']`.
         num_blocks
             Number of the `FT_Transformer` blocks, which should be non-negative.
-        attention_n_heads
+        attention_num_heads
             Number of attention heads in each `FT_Transformer` block, which should be positive.
         attention_initialization
             Weights initialization scheme for Multi Headed Attention module.
@@ -527,11 +531,11 @@ class FT_Transformer(nn.Module):
         """
 
         super().__init__()
-
+        logger.debug(f"initializing {prefix} (FT_Transformer)")
         assert num_categories or num_numerical_columns > 0, "there must be categorical columns or numerical columns"
-        assert token_dim > 0, "d_token must be positive"
-        assert num_blocks >= 0, "n_blocks must be non-negative"
-        assert attention_n_heads > 0, "attention_n_heads must be positive"
+        assert token_dim > 0, "token_dim must be positive"
+        assert num_blocks >= 0, "num_blocks must be non-negative"
+        assert attention_num_heads > 0, "attention_num_heads must be positive"
         assert token_initialization in ["uniform", "normal"], "initialization must be uniform or normal"
 
         self.prefix = prefix
@@ -544,14 +548,15 @@ class FT_Transformer(nn.Module):
         if num_categories:
             self.num_categories = num_categories
             self.categorical_feature_tokenizer = CategoricalFeatureTokenizer(
-                num_categories=num_categories,
-                d_token=token_dim,
+                num_categories=list(num_categories.values()),
+                token_dim=token_dim,
                 bias=token_bias,
                 initialization=token_initialization,
             )
             self.categorical_adapter = nn.Linear(token_dim, hidden_size)
 
         if num_numerical_columns > 0:
+            self.numerical_fill_values = numerical_fill_values
             self.numerical_feature_tokenizer = NumEmbeddings(
                 in_features=num_numerical_columns,
                 d_embedding=token_dim,
@@ -560,13 +565,13 @@ class FT_Transformer(nn.Module):
             self.numerical_adapter = nn.Linear(token_dim, hidden_size)
 
         self.transformer = Custom_Transformer(
-            d_token=hidden_size,
-            n_blocks=num_blocks,
-            attention_n_heads=attention_n_heads,
+            token_dim=hidden_size,
+            num_blocks=num_blocks,
+            attention_num_heads=attention_num_heads,
             attention_dropout=attention_dropout,
             attention_initialization=attention_initialization,
             attention_normalization=attention_normalization,
-            ffn_d_hidden=ffn_hidden_size,
+            ffn_hidden_size=ffn_hidden_size,
             ffn_dropout=ffn_dropout,
             ffn_activation=ffn_activation,
             ffn_normalization=ffn_normalization,
@@ -574,7 +579,7 @@ class FT_Transformer(nn.Module):
             prenormalization=prenormalization,
             first_prenormalization=first_prenormalization,
             last_layer_query_idx=None,
-            n_tokens=None,
+            num_tokens=None,
             kv_compression_ratio=kv_compression_ratio,
             kv_compression_sharing=kv_compression_sharing,
             head_activation=head_activation,
@@ -594,7 +599,7 @@ class FT_Transformer(nn.Module):
         )
 
         self.cls_token = CLSToken(
-            d_token=hidden_size,
+            token_dim=hidden_size,
             initialization="uniform",
         )
 
@@ -605,7 +610,7 @@ class FT_Transformer(nn.Module):
             self.categorical_adapter.apply(init_weights)
         self.head.apply(init_weights)
         # init transformer backbone from provided checkpoint
-        from autogluon.multimodal.utils.download import download
+        from ..utils.download import download
 
         if pretrained and checkpoint_name:
             if os.path.exists(checkpoint_name):

autogluon/multimodal/models/fusion/base.py

@@ -4,8 +4,6 @@ from typing import Optional
 
 from torch import nn
 
-from ...constants import AUTOMM, LABEL
-
 logger = logging.getLogger(__name__)
 
 
@@ -18,12 +16,12 @@ class AbstractMultimodalFusionModel(ABC, nn.Module):
         self,
         prefix: str,
         models: list,
-        loss_weight: Optional[float] = None,
+        aux_loss_weight: Optional[float] = None,
     ):
         super().__init__()
 
         self.prefix = prefix
-        self.loss_weight = loss_weight
+        self.aux_loss_weight = aux_loss_weight
         self.model = nn.ModuleList(models)
 
     @property

autogluon/multimodal/models/fusion/fusion_mlp.py

@@ -4,7 +4,19 @@ from typing import List, Optional
 import torch
 from torch import nn
 
-from ...constants import AUTOMM, FEATURES, LABEL, LOGITS, WEIGHT
+from ...constants import (
+    AUG_LOGITS,
+    FEATURES,
+    LABEL,
+    LOGITS,
+    MULTIMODAL_FEATURES,
+    MULTIMODAL_FEATURES_POST_AUG,
+    MULTIMODAL_FEATURES_PRE_AUG,
+    ORI_LOGITS,
+    VAE_MEAN,
+    VAE_VAR,
+    WEIGHT,
+)
 from ..mlp import MLP
 from ..utils import init_weights, run_model
 from .base import AbstractMultimodalFusionModel
@@ -27,9 +39,9 @@ class MultimodalFusionMLP(AbstractMultimodalFusionModel):
         num_classes: int,
         adapt_in_features: Optional[str] = None,
         activation: Optional[str] = "gelu",
-        dropout_prob: Optional[float] = 0.5,
+        dropout: Optional[float] = 0.5,
         normalization: Optional[str] = "layer_norm",
-        loss_weight: Optional[float] = None,
+        aux_loss_weight: Optional[float] = None,
     ):
         """
         Parameters
@@ -56,24 +68,26 @@ class MultimodalFusionMLP(AbstractMultimodalFusionModel):
                 dimension 768.
         activation
             Name of activation function.
-        dropout_prob
+        dropout
             Dropout probability.
         normalization
             Name of normalization function.
-        loss_weight
+        aux_loss_weight
             The weight of individual models. For example, if we fuse the features of ViT, CLIP, and BERT,
-            The loss will be computed as "loss = fusion_loss + loss_weight(vit_loss + clip_loss + bert_loss)".
+            The loss will be computed as "loss = fusion_loss + aux_loss_weight(vit_loss + clip_loss + bert_loss)".
             Basically, it supports adding an auxiliary loss for each individual model.
         """
         super().__init__(
             prefix=prefix,
             models=models,
-            loss_weight=loss_weight,
+            aux_loss_weight=aux_loss_weight,
         )
-        logger.debug("initializing MultimodalFusionMLP")
-        if loss_weight is not None:
-            assert loss_weight > 0
+        logger.debug(f"initializing {prefix} (MultimodalFusionMLP)")
+        if aux_loss_weight is not None:
+            assert aux_loss_weight >= 0
+            logger.debug(f"auxiliary loss weight: {aux_loss_weight}")
         self.num_classes = num_classes
+        self.augmenter = None
 
         raw_in_features = [per_model.out_features for per_model in models]
         if adapt_in_features is not None:
@@ -92,6 +106,7 @@ class MultimodalFusionMLP(AbstractMultimodalFusionModel):
             in_features = sum(raw_in_features)
 
         assert len(self.adapter) == len(self.model)
+        self.augmenter_in_features = in_features
 
         fusion_mlp = []
         for per_hidden_features in hidden_features:
@@ -102,7 +117,7 @@ class MultimodalFusionMLP(AbstractMultimodalFusionModel):
                     out_features=per_hidden_features,
                     num_layers=1,
                     activation=activation,
-                    dropout_prob=dropout_prob,
+                    dropout=dropout,
                     normalization=normalization,
                 )
             )
@@ -146,12 +161,16 @@ class MultimodalFusionMLP(AbstractMultimodalFusionModel):
 
         Returns
         -------
-        If "loss_weight" is None, it returns dictionary containing the fusion model's logits and
+        If "aux_loss_weight" is None, it returns dictionary containing the fusion model's logits and
         features. Otherwise, it returns a list of dictionaries collecting all the models' output,
         including the fusion model's.
         """
         multimodal_features = []
         multimodal_logits = []
+        multimodal_features_pre_aug = None
+        multimodal_features_post_aug = None
+        vae_mean = None
+        vae_var = None
         offset = 0
         for per_model, per_adapter in zip(self.model, self.adapter):
             per_model_args = args[offset : offset + len(per_model.input_keys)]
@@ -163,23 +182,68 @@ class MultimodalFusionMLP(AbstractMultimodalFusionModel):
             multimodal_logits.append(per_output[per_model.prefix][LOGITS])
             offset += len(per_model.input_keys)
 
-        features = self.fusion_mlp(torch.cat(multimodal_features, dim=1))
+        # make sure the returned multimodal features contain unimodal encoder features
+        multimodal_features_ret = multimodal_features
+        multimodal_features = torch.cat(multimodal_features, dim=1)
+        batch_size = multimodal_features.shape[0]
+        if self.training and self.augmenter is not None:
+            multimodal_features_pre_aug = multimodal_features.detach().clone()  # [bs, dim]
+            multimodal_features_post_aug, vae_mean, vae_var = self.augmenter(multimodal_features_pre_aug)
+            multimodal_features_post_aug_clone = multimodal_features_post_aug.clone()
+            multimodal_features_post_aug_clone.register_hook(lambda grad: -grad * self.augmenter.adv_weight)
+            multimodal_features = torch.cat([multimodal_features, multimodal_features_post_aug_clone], dim=0)
+
+        features = self.fusion_mlp(multimodal_features)
         logits = self.head(features)
+        ori_logits = logits[:batch_size].detach()  # detach the original logits when computing the consistency loss
+        aug_logits = logits[batch_size:]
+
+        return (
+            features,
+            logits,
+            multimodal_logits,
+            multimodal_features_ret,
+            multimodal_features_pre_aug,
+            multimodal_features_post_aug,
+            ori_logits,
+            aug_logits,
+            vae_mean,
+            vae_var,
+        )
 
-        return features, logits, multimodal_logits
-
-    def get_output_dict(self, features: torch.Tensor, logits: torch.Tensor, multimodal_logits: List[torch.Tensor]):
+    def get_output_dict(
+        self,
+        features: torch.Tensor,
+        logits: torch.Tensor,
+        multimodal_logits: List[torch.Tensor],
+        multimodal_features: List[torch.Tensor],
+        multimodal_features_pre_aug: torch.Tensor,
+        multimodal_features_post_aug: torch.Tensor,
+        ori_logits: torch.Tensor,
+        aug_logits: torch.Tensor,
+        vae_mean: torch.Tensor,
+        vae_var: torch.Tensor,
+    ):
         fusion_output = {
             self.prefix: {
                 LOGITS: logits,
                 FEATURES: features,
+                MULTIMODAL_FEATURES: multimodal_features,
+                MULTIMODAL_FEATURES_PRE_AUG: multimodal_features_pre_aug,
+                MULTIMODAL_FEATURES_POST_AUG: multimodal_features_post_aug,
+                ORI_LOGITS: ori_logits,
+                AUG_LOGITS: aug_logits,
+                VAE_MEAN: vae_mean,
+                VAE_VAR: vae_var,
             }
         }
-        if self.loss_weight is not None:
+        # filter out None
+        fusion_output = {self.prefix: {k: v for k, v in fusion_output[self.prefix].items() if v is not None}}
+        if self.aux_loss_weight is not None:
             output = {}
             for per_model, per_logits in zip(self.model, multimodal_logits):
                 per_output = {per_model.prefix: {}}
-                per_output[per_model.prefix][WEIGHT] = torch.tensor(self.loss_weight).to(per_logits.dtype)
+                per_output[per_model.prefix][WEIGHT] = torch.tensor(self.aux_loss_weight).to(per_logits.dtype)
                 per_output[per_model.prefix][LOGITS] = per_logits
                 output.update(per_output)
             fusion_output[self.prefix].update({WEIGHT: torch.tensor(1.0).to(logits)})

autogluon/multimodal/models/fusion/fusion_ner.py

@@ -5,7 +5,7 @@ import torch
 import torch.nn.functional as F
 from torch import nn
 
-from ...constants import AUTOMM, FEATURES, LABEL, LOGITS, NER_ANNOTATION, NER_TEXT, TOKEN_WORD_MAPPING, WORD_OFFSETS
+from ...constants import FEATURES, LABEL, LOGITS, NER_ANNOTATION, NER_TEXT, TOKEN_WORD_MAPPING, WORD_OFFSETS
 from ..mlp import MLP
 from ..utils import run_model
 from .base import AbstractMultimodalFusionModel

autogluon/multimodal/models/fusion/fusion_transformer.py

@@ -4,7 +4,7 @@ from typing import Optional
 import torch
 from torch import nn
 
-from ...constants import AUTOMM, FEATURES, LABEL, LOGITS, WEIGHT
+from ...constants import FEATURES, LABEL, LOGITS, WEIGHT
 from ..custom_transformer import CLSToken, Custom_Transformer
 from ..utils import init_weights, run_model
 from .base import AbstractMultimodalFusionModel
@@ -25,15 +25,15 @@ class MultimodalFusionTransformer(AbstractMultimodalFusionModel):
         models: list,
         hidden_features: int,
         num_classes: int,
-        n_blocks: Optional[int] = 0,
-        attention_n_heads: Optional[int] = 8,
+        num_blocks: Optional[int] = 0,
+        attention_num_heads: Optional[int] = 8,
         attention_initialization: Optional[str] = "kaiming",
         attention_normalization: Optional[str] = "layer_norm",
         attention_dropout: Optional[str] = 0.2,
         residual_dropout: Optional[str] = 0.0,
         ffn_activation: Optional[str] = "reglu",
         ffn_normalization: Optional[str] = "layer_norm",
-        ffn_d_hidden: Optional[str] = 192,
+        ffn_hidden_size: Optional[str] = 192,
         ffn_dropout: Optional[str] = 0.0,
         prenormalization: Optional[bool] = True,
         first_prenormalization: Optional[bool] = False,
@@ -42,7 +42,7 @@ class MultimodalFusionTransformer(AbstractMultimodalFusionModel):
         head_activation: Optional[str] = "relu",
         head_normalization: Optional[str] = "layer_norm",
         adapt_in_features: Optional[str] = None,
-        loss_weight: Optional[float] = None,
+        aux_loss_weight: Optional[float] = None,
         additive_attention: Optional[bool] = False,
         share_qv_weights: Optional[bool] = False,
     ):
@@ -59,9 +59,9 @@ class MultimodalFusionTransformer(AbstractMultimodalFusionModel):
             feature dimensions.
         num_classes
             The number of classes.
-        n_blocks
+        num_blocks
             Number of the `FT_Transformer` blocks, which should be non-negative.
-        attention_n_heads
+        attention_num_heads
             Number of attention heads in each `FT_Transformer` block, which should be positive.
         attention_dropout
             Dropout ratio for the Multi Headed Attention module.
@@ -71,7 +71,7 @@ class MultimodalFusionTransformer(AbstractMultimodalFusionModel):
             Normalization policy for attention layers. "layer_norm" is a good default.
         residual_dropout
             Dropout ratio for the linear layers in FT_Transformer block.
-        ffn_d_hidden
+        ffn_hidden_size
             Number of the hidden nodes of the linear layers in the Feed-Forward Network module.
         ffn_dropout
             Dropout ratio of the hidden nodes of the linear layers in the Feed-Forward Network module.
@@ -99,9 +99,9 @@ class MultimodalFusionTransformer(AbstractMultimodalFusionModel):
                 Adapt all features to the maximum dimension. For example, if three models have
                 feature dimensions are [512, 768, 64], it will linearly map all the features to
                 dimension 768.
-        loss_weight
+        aux_loss_weight
             The weight of individual models. For example, if we fuse the features of ViT, CLIP, and BERT,
-            The loss will be computed as "loss = fusion_loss + loss_weight(vit_loss + clip_loss + bert_loss)".
+            The loss will be computed as "loss = fusion_loss + aux_loss_weight(vit_loss + clip_loss + bert_loss)".
             Basically, it supports adding an auxiliary loss for each individual model.
         additive_attention
             If 'true' the transformer will use additive attention with linear complexity to sequence length.
@@ -111,11 +111,11 @@ class MultimodalFusionTransformer(AbstractMultimodalFusionModel):
         super().__init__(
             prefix=prefix,
             models=models,
-            loss_weight=loss_weight,
+            aux_loss_weight=aux_loss_weight,
         )
-        logger.debug("initializing MultimodalFusionTransformer")
-        if loss_weight is not None:
-            assert loss_weight > 0
+        logger.debug(f"initializing {prefix} (MultimodalFusionTransformer)")
+        if aux_loss_weight is not None:
+            assert aux_loss_weight >= 0
 
         raw_in_features = [per_model.out_features for per_model in models]
 
@@ -133,13 +133,13 @@ class MultimodalFusionTransformer(AbstractMultimodalFusionModel):
         assert len(self.adapter) == len(self.model)
 
         self.fusion_transformer = Custom_Transformer(
-            d_token=in_features,
-            n_blocks=n_blocks,
-            attention_n_heads=attention_n_heads,
+            token_dim=in_features,
+            num_blocks=num_blocks,
+            attention_num_heads=attention_num_heads,
             attention_dropout=attention_dropout,
             attention_initialization=attention_initialization,
             attention_normalization=attention_normalization,
-            ffn_d_hidden=ffn_d_hidden,
+            ffn_hidden_size=ffn_hidden_size,
             ffn_dropout=ffn_dropout,
             ffn_activation=ffn_activation,
             ffn_normalization=ffn_normalization,
@@ -147,7 +147,7 @@ class MultimodalFusionTransformer(AbstractMultimodalFusionModel):
             prenormalization=prenormalization,
             first_prenormalization=first_prenormalization,
             last_layer_query_idx=None,
-            n_tokens=None,
+            num_tokens=None,
             kv_compression_ratio=kv_compression_ratio,
             kv_compression_sharing=kv_compression_sharing,
             head_activation=head_activation,
@@ -167,7 +167,7 @@ class MultimodalFusionTransformer(AbstractMultimodalFusionModel):
         )
 
         self.cls_token = CLSToken(
-            d_token=in_features,
+            token_dim=in_features,
             initialization="uniform",
         )
 
@@ -196,9 +196,9 @@ class MultimodalFusionTransformer(AbstractMultimodalFusionModel):
                 multimodal_feature = torch.unsqueeze(multimodal_feature, dim=1)
             multimodal_features.append(multimodal_feature)
 
-            if self.loss_weight is not None:
+            if self.aux_loss_weight is not None:
                 per_output[per_model.prefix].update(
-                    {WEIGHT: torch.tensor(self.loss_weight).to(multimodal_features[0])}
+                    {WEIGHT: torch.tensor(self.aux_loss_weight).to(multimodal_features[0])}
                 )
                 output.update(per_output)
 
@@ -213,7 +213,7 @@ class MultimodalFusionTransformer(AbstractMultimodalFusionModel):
                 FEATURES: features,
             }
         }
-        if self.loss_weight is not None:
+        if self.aux_loss_weight is not None:
             fusion_output[self.prefix].update({WEIGHT: torch.tensor(1.0).to(logits)})
             output.update(fusion_output)
             return output

autogluon/multimodal/models/{huggingface_text.py → hf_text.py}

@@ -7,7 +7,6 @@ from transformers import logging as hf_logging
 from transformers.models.t5 import T5PreTrainedModel
 
 from ..constants import (
-    AUTOMM,
     COLUMN,
     COLUMN_FEATURES,
     FEATURES,
@@ -24,6 +23,8 @@ from .utils import (
     get_column_features,
     get_hf_config_and_model,
     get_pretrained_tokenizer,
+    get_text_segment_num,
+    get_text_token_max_len,
     init_weights,
 )
 
@@ -48,6 +49,8 @@ class HFAutoModelForTextPrediction(nn.Module):
         low_cpu_mem_usage: Optional[bool] = False,
         pretrained: Optional[bool] = True,
         tokenizer_name: Optional[str] = "hf_auto",
+        max_text_len: Optional[int] = None,
+        text_segment_num: Optional[int] = 1,
         use_fast: Optional[bool] = True,
     ):
         """
@@ -82,13 +85,18 @@ class HFAutoModelForTextPrediction(nn.Module):
             Whether using the pretrained weights. If pretrained=True, download the pretrained model.
         tokenizer_name
             Name of the huggingface tokenizer type.
+        max_text_len
+            The maximum length of text tokens.
+        text_segment_num
+            The number of text segments.
         use_fast
             Use a fast Rust-based tokenizer if it is supported for a given model.
             If a fast tokenizer is not available for a given model, a normal Python-based tokenizer is returned instead.
             See: https://huggingface.co/docs/transformers/model_doc/auto#transformers.AutoTokenizer.from_pretrained.use_fast
         """
         super().__init__()
-        logger.debug(f"initializing {checkpoint_name}")
+        logger.debug(f"initializing {prefix} (HFAutoModelForTextPrediction)")
+        logger.debug(f"model checkpoint: {checkpoint_name}")
         self.checkpoint_name = checkpoint_name
         self.num_classes = num_classes
 
@@ -101,6 +109,17 @@ class HFAutoModelForTextPrediction(nn.Module):
             checkpoint_name=self.checkpoint_name,
             use_fast=use_fast,
         )
+        self.max_text_len = get_text_token_max_len(
+            provided_max_len=max_text_len,
+            config=self.config,
+            tokenizer=self.tokenizer,
+            checkpoint_name=self.checkpoint_name,
+        )
+        self.text_segment_num = get_text_segment_num(
+            config=self.config,
+            provided_segment_num=text_segment_num,
+            checkpoint_name=self.checkpoint_name,
+        )
 
         if isinstance(self.model, T5PreTrainedModel):
             self.is_t5 = True