from transformers import PretrainedConfig


class SaMerConfig(PretrainedConfig):
    """Configuration for SaMer adapter-only checkpoints."""

    model_type = "samer"

    def __init__(
        self,
        base_model_name_or_path="vidore/colpali-v1.3-hf",
        base_model_revision=None,
        base_model_dtype="bfloat16",
        num_regions=64,
        cluster_iters=3,
        spatial_weight=0.1,
        assignment_temperature=0.07,
        image_token_count=1024,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.base_model_name_or_path = base_model_name_or_path
        self.base_model_revision = base_model_revision
        self.base_model_dtype = base_model_dtype
        self.num_regions = num_regions
        self.cluster_iters = cluster_iters
        self.spatial_weight = spatial_weight
        self.assignment_temperature = assignment_temperature
        self.image_token_count = image_token_count