| """Configuration for the self-contained Jolia Atlas backbone.""" |
|
|
| from __future__ import annotations |
|
|
| from transformers import PretrainedConfig |
|
|
|
|
| class JoliaConfig(PretrainedConfig): |
| """Config for :class:`~modeling_jolia.JoliaModel`. |
| |
| Jolia wraps a ``MultiModalAtlas`` 3D vision backbone and, when the source |
| run was trained with a ParallelOrganCLIP objective, a per-scale bank of |
| organ-query cross-attention pooling heads. |
| |
| Args: |
| embed_dim: Pooled (CLS-equivalent) embedding size. |
| multiscale_feats: Whether the backbone concatenates multi-scale pooled |
| features (always ``True`` for the released checkpoint). |
| atlas_config: Nested dict describing the Atlas modalities / stages. |
| Passed straight through to ``MultiModalAtlasConfig``. |
| patch_embed_dim: Per-scale organ-query dimension (``dim`` of the Atlas |
| stages). Total ``query_dim`` is ``num_scales * patch_embed_dim``. |
| num_organs: Number of learnable organ-query slots per scale. |
| num_scales: Number of multi-scale levels with an organ-query head. |
| num_heads: Attention heads in each organ-query pooler. |
| use_cls_residual: Add the CLS vector to every organ query (matches the |
| training-time ParallelOrganCLIP setting). |
| organ_slot_names: Ordered names for the organ-query slots. Length may be |
| ``< num_organs`` — trailing slots are unused padding. Enables the |
| named :meth:`JoliaModel.encode_organs` API. |
| text_embed_dim: Hidden size of the paired text encoder. Set non-zero to |
| build the linear text-projection head used for CLIP-style zero-shot |
| (maps text features into the same ``embed_dim`` space as the |
| vision CLS). ``0`` -> no text head. |
| text_encoder_id: HuggingFace id of the text encoder Jolia was trained |
| against (e.g. ``"Qwen/Qwen3-Embedding-8B"``). Informational; loading |
| it is opt-in via :class:`text_encoder_jolia.JoliaTextEncoder`. |
| text_context_length: Tokenizer ``max_length`` to match training. |
| """ |
|
|
| model_type = "jolia" |
|
|
| def __init__( |
| self, |
| embed_dim: int = 576, |
| multiscale_feats: bool = True, |
| atlas_config: dict | None = None, |
| patch_embed_dim: int = 0, |
| num_organs: int = 0, |
| num_scales: int = 0, |
| num_heads: int = 1, |
| use_cls_residual: bool = False, |
| organ_slot_names: list[str] | None = None, |
| text_embed_dim: int = 0, |
| text_encoder_id: str = "", |
| text_context_length: int = 512, |
| **kwargs: object, |
| ) -> None: |
| self.embed_dim = embed_dim |
| self.multiscale_feats = multiscale_feats |
| self.atlas_config = atlas_config or {} |
| self.patch_embed_dim = patch_embed_dim |
| self.num_organs = num_organs |
| self.num_scales = num_scales |
| self.num_heads = num_heads |
| self.use_cls_residual = use_cls_residual |
| self.organ_slot_names = list(organ_slot_names) if organ_slot_names else [] |
| self.text_embed_dim = text_embed_dim |
| self.text_encoder_id = text_encoder_id |
| self.text_context_length = text_context_length |
| super().__init__(**kwargs) |
|
|
| @property |
| def has_text_head(self) -> bool: |
| return self.text_embed_dim > 0 |
|
|
| @property |
| def query_dim(self) -> int: |
| """Total organ-query dimension (concatenated across scales).""" |
| return int(self.num_scales * self.patch_embed_dim) |
|
|
| @property |
| def has_queries(self) -> bool: |
| return self.num_organs > 0 and self.num_scales > 0 |
|
|
