from typing import Literal

from huggingface_hub.dataclasses import strict
from transformers import AutoConfig
from transformers.configuration_utils import PreTrainedConfig, SubConfigSpec


@strict
class UltraVBertConfig(PreTrainedConfig):
    model_type = "ultravbert"
    sub_configs_defaults = {
        "text_config": SubConfigSpec(config_class=AutoConfig, model_type="ultrabert"),
        "vision_config": SubConfigSpec(config_class=AutoConfig, model_type="siglip_vision_model"),
    }

    text_config: PreTrainedConfig | dict | None = None
    vision_config: PreTrainedConfig | dict | None = None

    image_token_id: int = 49191
    pixel_shuffle_factor: int = 4
    initializer_range: float = 0.02
    initializer_cutoff_factor: float = 2.0
    classifier_pooling: Literal["cls", "mean"] = "cls"
    classifier_dropout: float | int = 0.0
    tie_word_embeddings: bool = False
    connector_norm: bool = True


__all__ = ["UltraVBertConfig"]