from typing import List, Tuple from transformers import PretrainedConfig, AutoTokenizer class MolmoVisionConfig(PretrainedConfig): def __init__( self, image_default_input_size: Tuple[int, int] = (336, 336), image_patch_size: int = 14, image_pos_patch_size: int = 14, image_emb_dim: int = 1024, image_num_heads: int = 16, image_num_key_value_heads: int = 16, image_num_layers: int = 23, image_head_dim: int = 64, image_mlp_dim: int = 4096, image_mlp_activations: str = "quick_gelu", residual_dropout: float = 0, image_num_pos: int = 577, image_norm_eps: float = 1e-5, float32_attention: bool = True, attention_type: str = "spda", **kwargs ): super().__init__(**kwargs) self.image_default_input_size = image_default_input_size self.image_patch_size = image_patch_size self.image_pos_patch_size = image_pos_patch_size self.image_emb_dim = image_emb_dim self.image_num_heads = image_num_heads self.image_num_key_value_heads = image_num_key_value_heads self.image_num_layers = image_num_layers self.image_head_dim = image_head_dim self.image_mlp_dim = image_mlp_dim self.image_mlp_activations = image_mlp_activations self.residual_dropout = residual_dropout self.image_num_pos = image_num_pos self.image_norm_eps = image_norm_eps self.float32_attention = float32_attention @property def image_num_patch(self): h, w = self.image_default_input_size return h // self.image_patch_size, w // self.image_patch_size class MolmoConfig(PretrainedConfig): model_type = "molmo" keys_to_ignore_at_inference = ["past_key_values"] def __init__( self, vocab_size=50304, embedding_size=50304, hidden_size=4096, intermediate_size=11008, num_hidden_layers=32, num_attention_heads=32, num_key_value_heads=None, float32_attention=True, max_position_embeddings=2048, initializer_range=0.02, use_cache=True, layer_norm_eps: float = 1e-5, rope_theta=10000.0, clip_qkv=None, activation_type="silu", qkv_bias: bool = False, weight_tying: bool = False, use_position_ids: bool=True, tie_word_embeddings: bool=True, bias_for_layer_norm: bool=False, qk_layer_norm: bool=False, norm_after: bool = False, layer_norm_type: str="rms", vision_config: MolmoVisionConfig=None, vit_layers=(-2, -9), residual_dropout: float=0.0, embedding_dropout: float=0.0, attention_dropout: float=0.0, image_feature_dropout: float=0.0, additional_vocab_size=128, attention_type: str = "sdpa", image_padding_embed="pad_and_partial_pad", moe_num_experts=None, moe_top_k=None, normalize_input_embeds: bool=False, scale_logits: bool=False, **kwargs, ): if isinstance(vision_config, dict): self.vision_config = MolmoVisionConfig(**vision_config) elif vision_config is None: self.vision_config = MolmoVisionConfig() else: self.vision_config = vision_config self.vocab_size = vocab_size self.embedding_size = embedding_size self.max_position_embeddings = max_position_embeddings self.hidden_size = hidden_size self.intermediate_size = intermediate_size self.num_hidden_layers = num_hidden_layers self.num_attention_heads = num_attention_heads self.layer_norm_eps = layer_norm_eps self.weight_tying = weight_tying self.use_position_ids = use_position_ids self.qk_layer_norm = qk_layer_norm self.num_key_value_heads = num_key_value_heads self.float32_attention= float32_attention self.initializer_range = initializer_range self.use_cache = use_cache self.rope_theta = rope_theta self.clip_qkv = clip_qkv self.activation_type = activation_type self.qkv_bias = qkv_bias self.norm_after = norm_after self.tie_word_embeddings = tie_word_embeddings self.layer_norm_type = layer_norm_type self.moe_num_experts = moe_num_experts self.moe_top_k = moe_top_k self.vit_layers = vit_layers self.residual_dropout = residual_dropout self.embedding_dropout = embedding_dropout self.attention_dropout = attention_dropout self.image_feature_dropout = image_feature_dropout self.image_padding_embed = image_padding_embed self.bias_for_layer_norm = bias_for_layer_norm self.additional_vocab_size = additional_vocab_size self.attention_type = attention_type self.normalize_input_embeds = normalize_input_embeds self.scale_logits = scale_logits super().__init__( tie_word_embeddings=tie_word_embeddings, **kwargs, ) @property def effective_num_key_value_heads(self) -> int: if self.num_key_value_heads is None: return self.num_attention_heads else: return self.num_key_value_heads @property def image_num_patch(self): assert self.vision_config is not None return self.vision_config.image_num_patch MolmoVisionConfig.register_for_auto_class() MolmoConfig.register_for_auto_class()