Spaces:
Running
Running
| from typing import Dict, Any, Union | |
| from transformers import PretrainedConfig, AutoConfig | |
| class TinyDocVLMConfig(PretrainedConfig): | |
| model_type = "tinydoc_vlm" | |
| is_composition = True | |
| def __init__( | |
| self, | |
| vision_config: Union[Dict[str, Any], PretrainedConfig] = None, | |
| decoder_config: Union[Dict[str, Any], PretrainedConfig] = None, | |
| pixel_shuffle_scale: int = 3, | |
| image_size: int = 384, | |
| patch_size: int = 16, | |
| **kwargs, | |
| ): | |
| super().__init__(**kwargs) | |
| # Set defaults if not provided | |
| if vision_config is None: | |
| # Default SigLIP-B/16-like configuration (approx 93M parameters) | |
| vision_config = { | |
| "model_type": "siglip_vision_model", | |
| "hidden_size": 768, | |
| "intermediate_size": 3072, | |
| "num_hidden_layers": 12, | |
| "num_attention_heads": 12, | |
| "patch_size": 16, | |
| "image_size": 384, | |
| "num_channels": 3, | |
| "layer_norm_eps": 1e-6, | |
| } | |
| if decoder_config is None: | |
| # Default SmolLM2-135M-like configuration (approx 135M parameters) | |
| decoder_config = { | |
| "model_type": "llama", | |
| "vocab_size": 49152, | |
| "hidden_size": 576, | |
| "intermediate_size": 1536, | |
| "num_hidden_layers": 30, | |
| "num_attention_heads": 9, | |
| "num_key_value_heads": 3, | |
| "max_position_embeddings": 8192, | |
| "rms_norm_eps": 1e-5, | |
| "rope_theta": 273000.0, | |
| "attention_bias": False, | |
| } | |
| # Initialize config objects | |
| if isinstance(vision_config, dict): | |
| vision_config_copy = vision_config.copy() | |
| vision_model_type = vision_config_copy.pop("model_type", "siglip_vision_model") | |
| self.vision_config = AutoConfig.for_model(vision_model_type, **vision_config_copy) | |
| else: | |
| self.vision_config = vision_config | |
| if isinstance(decoder_config, dict): | |
| decoder_config_copy = decoder_config.copy() | |
| decoder_model_type = decoder_config_copy.pop("model_type", "llama") | |
| self.decoder_config = AutoConfig.for_model(decoder_model_type, **decoder_config_copy) | |
| else: | |
| self.decoder_config = decoder_config | |
| self.pixel_shuffle_scale = pixel_shuffle_scale | |
| self.image_size = image_size | |
| self.patch_size = patch_size | |
| def __getattr__(self, name): | |
| if name in ('decoder_config', 'vision_config'): | |
| raise AttributeError(name) | |
| if 'decoder_config' in self.__dict__: | |
| try: | |
| return getattr(self.decoder_config, name) | |
| except AttributeError: | |
| pass | |
| raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'") | |
| def to_dict(self) -> Dict[str, Any]: | |
| output = super().to_dict() | |
| output["vision_config"] = self.vision_config.to_dict() | |
| output["decoder_config"] = self.decoder_config.to_dict() | |
| return output | |