File size: 4,929 Bytes
c5a246d | 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 | from transformers.configuration_utils import PretrainedConfig
class SubConfig(PretrainedConfig):
# PretrainedConfig.to_dict() emits the class-level model_type (""), breaking save/load round-trips
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
|