| from transformers.configuration_utils import PretrainedConfig |
|
|
|
|
| class SubConfig(PretrainedConfig): |
| |
| def to_dict(self): |
| output = super().to_dict() |
| output["model_type"] = self.model_type |
| return output |
|
|
|
|
| class VisionConfig(SubConfig): |
| def __init__( |
| self, |
| **kwargs, |
| ): |
| super().__init__(**kwargs) |
|
|
| @staticmethod |
| def from_exp_config(vision_config: dict): |
| model_type = vision_config["model_type"] |
|
|
| if model_type == "vjepa2.1": |
| from .vjepa21_encoder import ARCH_MAP, VIT_EMBED_DIMS |
|
|
| arch = vision_config.get("arch", "vjepa2_1_vit_large_384") |
| factory_name = ARCH_MAP[arch][0].__name__ |
| embed_dim = VIT_EMBED_DIMS.get(factory_name, 1024) |
| use_hierarchical = vision_config.get("use_hierarchical", False) |
| n_levels = ARCH_MAP[arch][1] |
| vision_config["hidden_size"] = embed_dim * n_levels if use_hierarchical else embed_dim |
| elif model_type not in ["vjepa2", "dinov2", "dinov3_vit"]: |
| raise NotImplementedError() |
|
|
| vision_config = VisionConfig(**vision_config) |
|
|
| return vision_config |
|
|
|
|
| class TextConfig(SubConfig): |
| def __init__( |
| self, |
| model_type, |
| **kwargs, |
| ): |
| super().__init__(**kwargs) |
| self.model_type = model_type |
|
|
| @staticmethod |
| def from_exp_config( |
| text_config: dict, |
| ): |
| if text_config["model_type"] not in ["mpnet", "modernbert", "cxr-bert"]: |
| raise NotImplementedError() |
|
|
| return TextConfig(**text_config) |
|
|
|
|
| class AlignLayerConfig(SubConfig): |
| def __init__( |
| self, |
| model_type: str = "align_transformer", |
| projector_config=None, |
| **kwargs, |
| ): |
| super().__init__(**kwargs) |
| self.model_type = model_type |
| self.projector_config = projector_config |
|
|
| @staticmethod |
| def from_exp_config( |
| align_layer_config: dict, |
| ): |
| if align_layer_config["model_type"] == "align_transformer": |
| from transformers.models.dinov2.configuration_dinov2 import Dinov2Config |
|
|
| config = Dinov2Config(**align_layer_config) |
| config = config.to_dict() |
| align_layer_config.update(config) |
| align_layer_config["model_type"] = "align_transformer" |
|
|
| align_layer_config = AlignLayerConfig( |
| **align_layer_config, |
| ) |
|
|
| return align_layer_config |
|
|
|
|
| class ProjectorConfig(SubConfig): |
| def __init__(self, **kwargs): |
| super().__init__(**kwargs) |
|
|
|
|
| class DFRConfig(SubConfig): |
| def __init__( |
| self, |
| vfm_teacher_config: dict = None, |
| projector_config: dict = None, |
| layer_mapping: dict = None, |
| **kwargs, |
| ): |
| super().__init__(**kwargs) |
| self.vfm_teacher_config = vfm_teacher_config |
| self.projector_config = projector_config |
| self.layer_mapping = layer_mapping |
|
|
| @staticmethod |
| def from_exp_config(dfr_config: dict): |
| if "vfm_teacher_config" in dfr_config and dfr_config["vfm_teacher_config"] is not None: |
| dfr_config["vfm_teacher_config"] = VisionConfig.from_exp_config(dfr_config["vfm_teacher_config"]) |
|
|
| if "projector_config" in dfr_config and dfr_config["projector_config"] is not None: |
| dfr_config["projector_config"] = ProjectorConfig(**dfr_config["projector_config"]) |
|
|
| return DFRConfig(**dfr_config) |
|
|
|
|
| class GLINTConfig(PretrainedConfig): |
| is_composition = True |
|
|
| def __init__( |
| self, |
| vision_config: dict = None, |
| text_config: dict = None, |
| align_layer_config: dict = None, |
| predictor_config: dict = None, |
| prototype_head_config: dict = None, |
| dfr_config: dict = None, |
| vision_only: bool = False, |
| **kwargs, |
| ): |
| super().__init__(**kwargs) |
| self.vision_only = vision_only |
|
|
| if vision_config is not None: |
| pretrained_dir = kwargs.get("pretrained_dir") |
| if pretrained_dir and "pretrained_dir" not in vision_config: |
| vision_config["pretrained_dir"] = pretrained_dir |
| self.vision_config = VisionConfig.from_exp_config(vision_config) |
|
|
| if text_config is not None: |
| self.text_config = TextConfig.from_exp_config(text_config) |
|
|
| if align_layer_config is not None: |
| self.align_layer_config = AlignLayerConfig.from_exp_config(align_layer_config) |
|
|
| if dfr_config is not None: |
| pretrained_dir = kwargs.get("pretrained_dir") |
| teacher = dfr_config.get("vfm_teacher_config") |
| if pretrained_dir and teacher is not None: |
| teacher.setdefault("pretrained_dir", pretrained_dir) |
| self.dfr_config = DFRConfig.from_exp_config(dfr_config) |
|
|
| self.kwargs = kwargs |
|
|
