modeling_molparser_mobile.py

"""MolParser Mobile model code for Hugging Face Hub remote loading."""

from __future__ import annotations

import torch
import torch.nn as nn
import timm
from transformers import (
    AutoConfig,
    AutoModel,
    PretrainedConfig,
    PreTrainedModel,
    VisionEncoderDecoderConfig,
    VisionEncoderDecoderModel,
)
from transformers.modeling_outputs import BaseModelOutput


MOBILE_TIMM_MODEL_NAME = "vit_tiny_r_s16_p8_224.augreg_in21k_ft_in1k"
MOBILE_IMAGE_SIZE = 224
MOBILE_HIDDEN_SIZE = 192
MOBILE_PIXEL_UNSHUFFLE = 1


class CustomEncoderConfig(PretrainedConfig):
    model_type = "custom_timm_encoder"

    def __init__(
        self,
        timm_model_name: str = MOBILE_TIMM_MODEL_NAME,
        timm_pretrained: bool = False,
        timm_output_dim: int = MOBILE_HIDDEN_SIZE,
        pixel_unshuffle: int = MOBILE_PIXEL_UNSHUFFLE,
        hidden_size: int = MOBILE_HIDDEN_SIZE,
        initializer_range: float = 0.02,
        model_input_size: int = MOBILE_IMAGE_SIZE,
        **kwargs,
    ):
        self.timm_model_name = timm_model_name
        self.timm_pretrained = timm_pretrained
        self.timm_output_dim = timm_output_dim
        self.pixel_unshuffle = pixel_unshuffle
        self.hidden_size = hidden_size
        self.initializer_range = initializer_range
        self.model_input_size = model_input_size
        super().__init__(**kwargs)
        if getattr(self, "auto_map", None) is None:
            self.auto_map = {
                "AutoConfig": "modeling_molparser_mobile.CustomEncoderConfig",
                "AutoModel": "modeling_molparser_mobile.CustomTimmEncoder",
            }


class CustomTimmEncoder(PreTrainedModel):
    config_class = CustomEncoderConfig
    main_input_name = "pixel_values"

    def __init__(self, config: CustomEncoderConfig):
        super().__init__(config)
        self.config = config
        self.pixel_unshuffle_factor = int(config.pixel_unshuffle)
        self.unshuffle = (
            nn.PixelUnshuffle(self.pixel_unshuffle_factor)
            if self.pixel_unshuffle_factor > 1
            else nn.Identity()
        )

        timm_kwargs = {
            "pretrained": bool(config.timm_pretrained),
            "features_only": True,
            "num_classes": 0,
            "global_pool": "",
        }
        if getattr(config, "model_input_size", None):
            timm_kwargs["img_size"] = int(config.model_input_size)
        self.timm_model = timm.create_model(config.timm_model_name, **timm_kwargs)

        in_channels = int(config.timm_output_dim) * self.pixel_unshuffle_factor**2
        self.use_projection = not (
            self.pixel_unshuffle_factor == 1 and in_channels == int(config.hidden_size)
        )
        if self.use_projection:
            self.projection = nn.Sequential(
                nn.Conv2d(in_channels=in_channels, out_channels=config.hidden_size, kernel_size=1),
                nn.GELU(),
            )
        else:
            self.projection = nn.Identity()

    def forward(self, pixel_values: torch.Tensor, **kwargs):
        encoder_features = self.timm_model(pixel_values)[-1]
        if encoder_features.ndim != 4:
            raise ValueError(f"Expected 4D feature map, got shape={tuple(encoder_features.shape)}")
        if encoder_features.shape[1] == self.config.timm_output_dim:
            pass
        elif encoder_features.shape[-1] == self.config.timm_output_dim:
            encoder_features = encoder_features.permute(0, 3, 1, 2).contiguous()
        else:
            raise ValueError(
                "Unexpected timm feature shape "
                f"{tuple(encoder_features.shape)} for timm_output_dim={self.config.timm_output_dim}"
            )

        encoder_features = self.unshuffle(encoder_features)
        encoder_features = self.projection(encoder_features)
        _, channels, _, _ = encoder_features.shape
        if channels != self.config.hidden_size:
            raise ValueError(
                f"Unexpected encoder channels={channels}, expected hidden_size={self.config.hidden_size}."
            )
        encoder_hidden_states = encoder_features.flatten(2).transpose(1, 2)
        return BaseModelOutput(last_hidden_state=encoder_hidden_states)


try:
    AutoConfig.register("custom_timm_encoder", CustomEncoderConfig)
except ValueError:
    pass
try:
    AutoModel.register(CustomEncoderConfig, CustomTimmEncoder)
except ValueError:
    pass

CustomEncoderConfig.register_for_auto_class()
CustomTimmEncoder.register_for_auto_class("AutoModel")


class MolParserVisionEncoderDecoderConfig(VisionEncoderDecoderConfig):
    model_type = "molparser_vision_encoder_decoder"


class MolParserVisionEncoderDecoderModel(VisionEncoderDecoderModel):
    config_class = MolParserVisionEncoderDecoderConfig

    def __init__(self, config=None, encoder=None, decoder=None):
        if config is not None and not isinstance(config, self.config_class):
            config = self.config_class.from_dict(config.to_dict())
        super().__init__(config=config, encoder=encoder, decoder=decoder)

    @classmethod
    def get_init_context(cls, dtype, is_quantized, _is_ds_init_called, allow_all_kernels):
        contexts = super().get_init_context(dtype, is_quantized, _is_ds_init_called, allow_all_kernels)
        if is_quantized:
            return contexts
        # timm ViT initialization calls tensor.item(), which cannot run on meta tensors.
        return [ctx for ctx in contexts if not (isinstance(ctx, torch.device) and ctx.type == "meta")]


try:
    AutoConfig.register("molparser_vision_encoder_decoder", MolParserVisionEncoderDecoderConfig)
except ValueError:
    pass

MolParserVisionEncoderDecoderConfig.register_for_auto_class()
MolParserVisionEncoderDecoderModel.register_for_auto_class("AutoModelForImageTextToText")


def assert_mobile_config(config: MolParserVisionEncoderDecoderConfig) -> None:
    encoder = getattr(config, "encoder", None)
    decoder = getattr(config, "decoder", None)
    if encoder is None or decoder is None:
        raise ValueError("MolParser Mobile config must contain encoder and decoder sub-configs.")

    checks = {
        "encoder.timm_model_name": getattr(encoder, "timm_model_name", None) == MOBILE_TIMM_MODEL_NAME,
        "encoder.model_input_size": int(getattr(encoder, "model_input_size", 0) or 0) == MOBILE_IMAGE_SIZE,
        "encoder.hidden_size": int(getattr(encoder, "hidden_size", 0) or 0) == MOBILE_HIDDEN_SIZE,
        "encoder.pixel_unshuffle": int(getattr(encoder, "pixel_unshuffle", 0) or 0) == MOBILE_PIXEL_UNSHUFFLE,
        "decoder.d_model": int(getattr(decoder, "d_model", 0) or 0) == MOBILE_HIDDEN_SIZE,
        "decoder.decoder_layers": int(getattr(decoder, "decoder_layers", 0) or 0) == 6,
        "decoder.decoder_attention_heads": int(getattr(decoder, "decoder_attention_heads", 0) or 0) == 4,
    }
    failed = [name for name, ok in checks.items() if not ok]
    if failed:
        raise ValueError(
            "This Hub package is mobile-only; non-mobile config values found: "
            + ", ".join(failed)
        )


def load_molparser_mobile_model(checkpoint_path: str) -> MolParserVisionEncoderDecoderModel:
    config = MolParserVisionEncoderDecoderConfig.from_pretrained(
        checkpoint_path,
        trust_remote_code=True,
    )
    assert_mobile_config(config)
    if getattr(config, "encoder", None) is not None:
        config.encoder.timm_pretrained = False
    model = MolParserVisionEncoderDecoderModel.from_pretrained(
        checkpoint_path,
        config=config,
        trust_remote_code=True,
    )
    return model


__all__ = [
    "CustomEncoderConfig",
    "CustomTimmEncoder",
    "MolParserVisionEncoderDecoderConfig",
    "MolParserVisionEncoderDecoderModel",
    "assert_mobile_config",
    "load_molparser_mobile_model",
]