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

autogluon/multimodal/models/meta_transformer.py

@@ -0,0 +1,336 @@
+import logging
+from typing import Dict, Optional
+
+import torch
+from timm import create_model
+from timm.models.vision_transformer import Block
+from torch import nn
+
+from ..constants import (
+    CATEGORICAL,
+    FEATURES,
+    IMAGE,
+    IMAGE_VALID_NUM,
+    LABEL,
+    LOGITS,
+    NUMERICAL,
+    TEXT_SEGMENT_IDS,
+    TEXT_TOKEN_IDS,
+    TEXT_VALID_LENGTH,
+)
+from .custom_transformer import CLSToken
+from .ft_transformer import CategoricalFeatureTokenizer, NumEmbeddings
+from .utils import (
+    assign_layer_ids,
+    get_hf_config_and_model,
+    get_image_size_mean_std,
+    get_pretrained_tokenizer,
+    get_text_segment_num,
+    get_text_token_max_len,
+    init_weights,
+    replace_missing_images_with_learnable,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class MetaTransformer(nn.Module):
+    def __init__(
+        self,
+        prefix: str,
+        num_classes: int,
+        checkpoint_path: str,
+        model_version: str,
+        has_image: bool,
+        has_text: bool,
+        num_numerical_columns: int,
+        num_categories: Dict,
+        numerical_fill_values: Dict,
+        image_size: Optional[int] = None,
+        image_norm: Optional[str] = None,
+        image_chan_num: Optional[int] = 3,
+        use_learnable_image: Optional[bool] = False,
+        max_text_len: Optional[int] = None,
+        text_segment_num: Optional[int] = 1,
+    ):
+        super().__init__()
+        logger.debug(f"initializing {prefix} (MetaTransformer)")
+        self.prefix = prefix
+        self.checkpoint_name = checkpoint_path
+
+        if model_version == "base":
+            dim = 768
+            num_heads = 12
+            layer_num = 12
+        elif model_version == "large":
+            dim = 1024
+            num_heads = 16
+            layer_num = 24
+        else:
+            raise ValueError(f"Unsupported model version: {model_version}. Options are 'base' and 'large'.")
+
+        self.model = nn.Sequential(
+            *[
+                Block(
+                    dim=dim,
+                    num_heads=num_heads,
+                    mlp_ratio=4.0,
+                    qkv_bias=True,
+                    norm_layer=nn.LayerNorm,
+                    act_layer=nn.GELU,
+                )
+                for i in range(layer_num)
+            ]
+        )
+
+        checkpoint = torch.load(checkpoint_path)  # nosec B614
+        self.checkpoint_path = checkpoint_path
+        self.model.load_state_dict(checkpoint, strict=True)
+
+        self.head = nn.Linear(dim, num_classes) if num_classes else nn.Identity()
+
+        self.cls_token = CLSToken(token_dim=dim, initialization="uniform")
+        self.config = None
+
+        self.tokenizer = None
+        self.text_adaptor = None
+        self.image_tokenizer = None
+        self.image_adaptor = None
+        self.categorical_feature_tokenizer = None
+        self.categorical_adapter = None
+        self.numerical_feature_tokenizer = None
+        self.numerical_adapter = None
+
+        # if has_text or has_image:
+        #     clip_ckpt = "openai/clip-vit-base-patch32"
+        #     _, clip_model = get_hf_config_and_model(checkpoint_name=clip_ckpt, pretrained=True)
+
+        if has_text:
+            checkpoint_name = "microsoft/deberta-v3-base"
+            _, text_model = get_hf_config_and_model(checkpoint_name=checkpoint_name, pretrained=True)
+            self.text_config = text_model.config
+            # refer to https://github.com/invictus717/MetaTransformer/blob/master/Data2Seq/Data2Seq.py#L28
+            self.tokenizer = get_pretrained_tokenizer(
+                tokenizer_name="hf_auto",
+                checkpoint_name=checkpoint_name,
+            )
+            self.text_embed = text_model.embeddings
+            self.text_adaptor = nn.Linear(self.text_config.hidden_size, dim)
+            self.tokenizer_name = "hf_auto"
+            self.max_text_len = get_text_token_max_len(
+                provided_max_len=max_text_len,
+                config=self.text_config,
+                tokenizer=self.tokenizer,
+                checkpoint_name=checkpoint_name,
+            )
+            self.text_segment_num = get_text_segment_num(
+                config=self.text_config,
+                provided_segment_num=text_segment_num,
+                checkpoint_name=checkpoint_name,
+            )
+        if has_image:
+            image_model = create_model("timm/vit_base_patch16_224.mae", pretrained=True)
+            self.image_config = image_model.default_cfg
+            self.patch_embed = image_model.patch_embed
+            self.image_adaptor = nn.Linear(image_model.embed_dim, dim)
+            self.image_size, self.image_mean, self.image_std = get_image_size_mean_std(
+                model_name=self.prefix,
+                config=self.image_config,
+                provided_size=image_size,
+                provided_norm_type=image_norm,
+                support_variable_input_size=False,
+            )
+            self.use_learnable_image = use_learnable_image
+            if self.use_learnable_image:
+                self.learnable_image = nn.Parameter(torch.zeros(image_chan_num, self.image_size, self.image_size))
+                logger.debug("will use a learnable image to replace missing ones")
+
+        if num_categories:
+            self.num_categories = num_categories
+            self.categorical_feature_tokenizer = CategoricalFeatureTokenizer(
+                num_categories=list(num_categories.values()),
+                token_dim=dim,
+                bias=True,
+                initialization="normal",
+            )
+            self.categorical_adapter = nn.Linear(dim, dim)
+
+        if num_numerical_columns > 0:
+            self.num_numerical_columns = num_numerical_columns
+            self.numerical_fill_values = numerical_fill_values
+            self.numerical_feature_tokenizer = NumEmbeddings(
+                in_features=num_numerical_columns,
+                d_embedding=dim,
+                embedding_arch=["linear"],
+            )
+            self.numerical_adapter = nn.Linear(dim, dim)
+
+        self.out_features = dim
+
+        # init weights
+        self.head.apply(init_weights)
+        self.name_to_id = self.get_layer_ids()
+        self.head_layer_names = [n for n, layer_id in self.name_to_id.items() if layer_id == 0]
+
+    @property
+    def text_token_ids_key(self):
+        return f"{self.prefix}_{TEXT_TOKEN_IDS}"
+
+    @property
+    def text_valid_length_key(self):
+        return f"{self.prefix}_{TEXT_VALID_LENGTH}"
+
+    @property
+    def text_segment_ids_key(self):
+        return f"{self.prefix}_{TEXT_SEGMENT_IDS}"
+
+    @property
+    def image_key(self):
+        return f"{self.prefix}_{IMAGE}"
+
+    @property
+    def image_valid_num_key(self):
+        return f"{self.prefix}_{IMAGE_VALID_NUM}"
+
+    @property
+    def categorical_key(self):
+        return f"{self.prefix}_{CATEGORICAL}"
+
+    @property
+    def numerical_key(self):
+        return f"{self.prefix}_{NUMERICAL}"
+
+    @property
+    def label_key(self):
+        return f"{self.prefix}_{LABEL}"
+
+    def forward(
+        self,
+        batch: dict,
+    ):
+        multimodal_features = []
+        if self.image_tokenizer:
+            images = batch[self.image_key]
+            image_valid_num = batch[self.image_valid_num_key]
+            b, n, c, h, w = images.shape
+            steps = torch.arange(0, n).type_as(image_valid_num)
+            image_masks = steps.reshape((1, -1)) < image_valid_num.reshape((-1, 1))  # (b, n)
+            if self.use_learnable_image:
+                images = replace_missing_images_with_learnable(
+                    images=images,
+                    image_masks=image_masks,
+                    learnable_image=self.learnable_image,
+                )
+            image_embeddings = self.patch_embed(images.reshape((b * n, c, h, w)))  # (b*n, l, d)
+            assert image_embeddings.ndim == 3
+            image_embeddings = self.image_adaptor(image_embeddings)
+            multimodal_features.append(image_embeddings)
+        if self.text_adaptor:  # text tokenizer is used in text processor
+            text_token_ids = batch[self.text_token_ids_key]
+            text_valid_length = batch[self.text_valid_length_key]
+            steps = torch.arange(0, text_token_ids.shape[1]).type_as(text_valid_length)
+            text_masks = (steps.reshape((1, -1)) < text_valid_length.reshape((-1, 1))).type_as(text_token_ids)
+            # text_embeddings = self.text_embeddings(batch[self.text_token_ids_key])  # (b, l, d)
+            input_ids = text_token_ids
+            inputs_embeds = None
+            attention_mask = text_masks
+            position_ids = None
+            if "token_type_ids" in self.tokenizer.model_input_names:
+                token_type_ids = batch[self.text_segment_ids_key]
+            else:
+                token_type_ids = None
+
+            if input_ids is not None and inputs_embeds is not None:
+                raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+            elif input_ids is not None:
+                input_shape = input_ids.size()
+            elif inputs_embeds is not None:
+                input_shape = inputs_embeds.size()[:-1]
+            else:
+                raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+            device = input_ids.device if input_ids is not None else inputs_embeds.device
+
+            if attention_mask is None:
+                attention_mask = torch.ones(input_shape, device=device)
+            if token_type_ids is None:
+                token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
+
+            text_embeddings = self.text_embed(
+                input_ids=input_ids,
+                token_type_ids=token_type_ids,
+                position_ids=position_ids,
+                mask=attention_mask,
+                inputs_embeds=inputs_embeds,
+            )
+            text_embeddings = self.text_adaptor(text_embeddings)
+            assert text_embeddings.ndim == 3
+            multimodal_features.append(text_embeddings)
+        if self.categorical_feature_tokenizer:
+            categorical_inputs = []
+            for categorical_input in batch[self.categorical_key]:
+                categorical_inputs.append(categorical_input)
+            categorical_inputs = torch.stack(categorical_inputs, dim=1)
+
+            categorical_features = self.categorical_feature_tokenizer(categorical_inputs)
+            categorical_features = self.categorical_adapter(categorical_features)  # (b, l, d)
+            assert categorical_features.ndim == 3
+            multimodal_features.append(categorical_features)
+        if self.numerical_feature_tokenizer:
+            numerical_features = self.numerical_feature_tokenizer(batch[self.numerical_key])
+            numerical_features = self.numerical_adapter(numerical_features)
+            assert numerical_features.ndim == 3
+            multimodal_features.append(numerical_features)
+
+        multimodal_features = torch.cat(multimodal_features, dim=1)
+        multimodal_features = self.cls_token(multimodal_features)
+        features = self.model(multimodal_features)
+        pooled_features = features[:, -1, :]  # CLSToken append the cls token to the sequence tail
+        logits = self.head(pooled_features)
+        ret = {
+            LOGITS: logits,
+            FEATURES: pooled_features,
+        }
+        return {self.prefix: ret}
+
+    def get_layer_ids(self):
+        """
+        Assign an id to each layer. Layer ids will be used in layer-wise lr decay.
+        Basically, id gradually increases when going from the output end to
+        the input end. The layers defined in this class, e.g., head, have id 0.
+
+        In the AutoModel scenario, this function may not always return the correct result.
+        Thus, you can use "print(json.dumps(name_to_id, indent=2))" to manually check whether
+        the layer ids are reasonable.
+
+        Returns
+        -------
+        A dictionary mapping the layer names (keys) to their ids (values).
+        """
+        model_prefix = "model"
+        pre_encoder_patterns = (
+            "embeddings",
+            "LayerNorm",
+            "wte",
+            "wpe",
+            "shared.weight",
+            "encoder.conv.conv",
+            "relative_attention_bias",
+            "dummy_layer",
+        )
+        post_encoder_patterns = ("head", "pooler", "ln_f", "final_layer_norm")
+        names = [n for n, _ in self.named_parameters()]
+
+        name_to_id, names = assign_layer_ids(
+            names=names,
+            pre_encoder_patterns=pre_encoder_patterns,
+            post_encoder_patterns=post_encoder_patterns,
+            model_pre=model_prefix,
+        )
+        if len(names) > 0:
+            logger.debug(f"outer layers are treated as head: {names}")
+        for n in names:
+            assert n not in name_to_id
+            name_to_id[n] = 0
+
+        return name_to_id

autogluon/multimodal/models/mlp.py

@@ -51,7 +51,7 @@ class Unit(nn.Module):
         in_features: int,
         out_features: int,
         activation: str,
-        dropout_prob: float,
+        dropout: float,
     ):
         """
         Parameters
@@ -64,7 +64,7 @@ class Unit(nn.Module):
             Dimension of output features.
         activation
             Name of activation function.
-        dropout_prob
+        dropout
             Dropout probability.
         """
         super().__init__()
@@ -78,7 +78,7 @@ class Unit(nn.Module):
             raise ValueError(f"unknown normalization: {normalization}")
         self.fc = nn.Linear(in_features, out_features)
         self.act_fn = ALL_ACT_LAYERS[activation]()
-        self.dropout = nn.Dropout(dropout_prob)
+        self.dropout = nn.Dropout(dropout)
 
     def forward(self, x):
         # pre normalization
@@ -102,7 +102,7 @@ class MLP(nn.Module):
         out_features: Optional[int] = None,
         num_layers: Optional[int] = 1,
         activation: Optional[str] = "gelu",
-        dropout_prob: Optional[float] = 0.5,
+        dropout: Optional[float] = 0.5,
         normalization: Optional[str] = "layer_norm",
     ):
         """
@@ -118,7 +118,7 @@ class MLP(nn.Module):
             Number of layers.
         activation
             Name of activation function.
-        dropout_prob
+        dropout
             Dropout probability.
         normalization
             Name of normalization function.
@@ -134,7 +134,7 @@ class MLP(nn.Module):
                 in_features=in_features,
                 out_features=hidden_features,
                 activation=activation,
-                dropout_prob=dropout_prob,
+                dropout=dropout,
             )
             in_features = hidden_features
             layers.append(per_unit)

autogluon/multimodal/models/mmocr_text_detection.py

@@ -11,7 +11,7 @@ except ImportError:
     mmocr = None
 from torch import nn
 
-from ..constants import AUTOMM, BBOX, COLUMN, COLUMN_FEATURES, FEATURES, IMAGE, IMAGE_VALID_NUM, LABEL, LOGITS, MASKS
+from ..constants import BBOX, COLUMN, COLUMN_FEATURES, FEATURES, IMAGE, IMAGE_VALID_NUM, LABEL, LOGITS, MASKS
 from .utils import assign_layer_ids, get_column_features, get_mmocr_config_and_model, get_model_head
 
 logger = logging.getLogger(__name__)

autogluon/multimodal/models/mmocr_text_recognition.py

@@ -12,7 +12,6 @@ except ImportError:
 from torch import nn
 
 from ..constants import (
-    AUTOMM,
     COLUMN,
     COLUMN_FEATURES,
     FEATURES,

autogluon/multimodal/models/ner_text.py

@@ -3,25 +3,18 @@ from typing import Dict, List, Optional, Tuple
 
 import torch
 import torch.nn.functional as F
-from torch import nn
 from transformers import logging as hf_logging
 
 from ..constants import (
-    AUTOMM,
-    COLUMN,
     COLUMN_FEATURES,
     FEATURES,
-    LABEL,
     LOGITS,
     MASKS,
     NER_ANNOTATION,
-    TEXT_SEGMENT_IDS,
-    TEXT_TOKEN_IDS,
-    TEXT_VALID_LENGTH,
     TOKEN_WORD_MAPPING,
     WORD_OFFSETS,
 )
-from .huggingface_text import HFAutoModelForTextPrediction
+from .hf_text import HFAutoModelForTextPrediction
 from .utils import assign_layer_ids, get_column_features, get_pretrained_tokenizer
 
 hf_logging.set_verbosity_error()

autogluon/multimodal/models/numerical_mlp.py

@@ -1,4 +1,5 @@
-from typing import List, Optional
+import logging
+from typing import Dict, List, Optional
 
 from torch import nn
 
@@ -7,6 +8,8 @@ from .ft_transformer import NumEmbeddings
 from .mlp import MLP
 from .utils import init_weights
 
+logger = logging.getLogger(__name__)
+
 
 class NumericalMLP(nn.Module):
     """
@@ -21,11 +24,12 @@ class NumericalMLP(nn.Module):
         out_features: Optional[int] = None,
         num_layers: Optional[int] = 1,
         activation: Optional[str] = "leaky_relu",
-        dropout_prob: Optional[float] = 0.5,
+        dropout: Optional[float] = 0.5,
         normalization: Optional[str] = "layer_norm",
         num_classes: Optional[int] = 0,
-        d_token: Optional[int] = 8,
+        token_dim: Optional[int] = 8,
         embedding_arch: Optional[List[str]] = None,
+        numerical_fill_values: Optional[Dict] = None,
     ):
         """
         Parameters
@@ -42,13 +46,13 @@ class NumericalMLP(nn.Module):
             Number of MLP layers.
         activation
             Name of activation function.
-        dropout_prob
+        dropout
             Dropout probability.
         normalization
             Name of normalization function.
         num_classes
             Number of classes. 1 for a regression task.
-        d_token
+        token_dim
             The size of one token for `NumericalEmbedding`.
         embedding_arch
             A list containing the names of embedding layers.
@@ -56,19 +60,21 @@ class NumericalMLP(nn.Module):
             {'linear', 'shared_linear', 'autodis', 'positional', 'relu', 'layernorm'}
         """
         super().__init__()
+        logger.debug(f"initializing {prefix} (NumericalMLP)")
         self.out_features = out_features
+        self.numerical_fill_values = numerical_fill_values
 
         self.numerical_feature_tokenizer = (
             NumEmbeddings(
                 in_features=in_features,
-                d_embedding=d_token,
+                d_embedding=token_dim,
                 embedding_arch=embedding_arch,
             )
             if embedding_arch is not None
             else nn.Identity()
         )
 
-        in_features = in_features * d_token if embedding_arch is not None else in_features
+        in_features = in_features * token_dim if embedding_arch is not None else in_features
 
         self.mlp = MLP(
             in_features=in_features,
@@ -76,7 +82,7 @@ class NumericalMLP(nn.Module):
             out_features=out_features,
             num_layers=num_layers,
             activation=activation,
-            dropout_prob=dropout_prob,
+            dropout=dropout,
             normalization=normalization,
         )
         self.head = nn.Linear(out_features, num_classes) if num_classes > 0 else nn.Identity()

autogluon/multimodal/models/sam.py

@@ -3,14 +3,13 @@ from typing import Dict, List, Optional, Tuple
 
 import torch
 import torch.nn.functional as F
-from omegaconf import DictConfig
 from torch import nn
 from transformers import SamConfig
 
 from ..constants import CLASS_LABEL, CLASS_LOGITS, COLUMN, IMAGE, IMAGE_VALID_NUM, LABEL, LOGITS, MASK_LABEL, MOE_LOSS
 from .adaptation_layers import ConvLoRALinear
 from .custom_hf_models.modeling_sam_for_conv_lora import SamImageSegmentationOutput, SamModel
-from .utils import assign_layer_ids, freeze_model_layers
+from .utils import assign_layer_ids, freeze_model_layers, image_mean_std
 
 logger = logging.getLogger(__name__)
 
@@ -269,6 +268,7 @@ class SAMForSemanticSegmentation(nn.Module):
         pretrained: Optional[bool] = True,
         frozen_layers: Optional[list] = None,
         num_mask_tokens: int = 1,
+        image_norm: Optional[str] = None,
     ):
         """
         Load a pretrained Segment Anything Model (SAM).
@@ -287,6 +287,15 @@ class SAMForSemanticSegmentation(nn.Module):
             A list of substrings of frozen layers' names.
         num_mask_tokens
             The number of mask proposals.
+        image_norm
+            How to normalize an image. We now support:
+            - inception
+                Normalize image by IMAGENET_INCEPTION_MEAN and IMAGENET_INCEPTION_STD from timm
+            - imagenet
+                Normalize image by IMAGENET_DEFAULT_MEAN and IMAGENET_DEFAULT_STD from timm
+            - clip
+                Normalize image by mean (0.48145466, 0.4578275, 0.40821073) and
+                std (0.26862954, 0.26130258, 0.27577711), used for CLIP.
         """
 
         super().__init__()
@@ -305,6 +314,7 @@ class SAMForSemanticSegmentation(nn.Module):
 
         self.image_size = self.model.vision_encoder.image_size
         self.config = self.model.config
+        self.image_mean, self.image_std = image_mean_std(image_norm)
 
         self.model.mask_decoder.num_mask_tokens = num_mask_tokens
         mask_token_data = self.model.mask_decoder.mask_tokens.weight.data[0]

autogluon/multimodal/models/t_few.py

@@ -1,7 +1,4 @@
-import collections
 import logging
-import os
-import random
 from functools import lru_cache
 from typing import Dict, List, Optional, Tuple
 
@@ -12,7 +9,6 @@ from transformers import AutoConfig, AutoModelForSeq2SeqLM
 from transformers import logging as hf_logging
 
 from ..constants import (
-    AUTOMM,
     CHOICES_IDS,
     COLUMN,
     COLUMN_FEATURES,
@@ -26,7 +22,14 @@ from ..constants import (
     TEXT_TOKEN_IDS,
     TEXT_VALID_LENGTH,
 )
-from .utils import DummyLayer, assign_layer_ids, get_column_features, get_pretrained_tokenizer
+from .utils import (
+    DummyLayer,
+    assign_layer_ids,
+    get_column_features,
+    get_pretrained_tokenizer,
+    get_text_segment_num,
+    get_text_token_max_len,
+)
 
 hf_logging.set_verbosity_error()
 
@@ -56,6 +59,8 @@ class TFewModel(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,
     ):
         """
         Load a pretrained T5-based text transformer backbone.
@@ -106,6 +111,17 @@ class TFewModel(nn.Module):
             tokenizer_name=self.tokenizer_name,
             checkpoint_name=self.checkpoint_name,
         )
+        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,
+        )
         self.eos_token = self.tokenizer.eos_token
         self.out_features = (
             self.model.config.hidden_size