"""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", ]