configuration_colvec1.py

"""
Configuration for ColVec1 retrieval model.
"""

from transformers.configuration_utils import PretrainedConfig


class ColVec1Config(PretrainedConfig):
    """Configuration for the ColVec1 retrieval wrapper."""

    model_type = "colvec1"

    def __init__(
        self,
        embed_dim: int = 128,
        text_hidden_size: int = 2560,
        padding_side: str = "left",
        initializer_range: float = 0.02,
        base_model_name_or_path: str = None,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.embed_dim = embed_dim
        self.text_hidden_size = text_hidden_size
        self.padding_side = padding_side
        self.initializer_range = initializer_range
        self.base_model_name_or_path = base_model_name_or_path


__all__ = ["ColVec1Config"]