File size: 7,801 Bytes
29d8d1e c81207b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 | """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",
]
|