configuration_embedder.py

from typing import Any

from transformers import AutoConfig, PretrainedConfig


class EmbedderConfig(PretrainedConfig):
    model_type = "embedder"

    def __init__(
        self,
        base_model_name: str = "nomic-ai/nomic-embed-text-v1.5",
        num_blocks: int = 2,
        dropout: float = 0.0,
        encoder_only: bool = False,
        **kwargs: Any,
    ):
        super().__init__(**kwargs)
        self.base_model_name = base_model_name
        self.num_blocks = num_blocks
        self.dropout = dropout
        self.encoder_only = encoder_only
        self.encoder_config = AutoConfig.from_pretrained(  # nosec B615
            base_model_name,
            trust_remote_code=True,
        )
        self.auto_map = {
            "AutoConfig": "configuration_embedder.EmbedderConfig",
            "AutoModel": "modeling_embedder.EmbedderModel",
        }

    def __getattr__(self, name: str) -> Any:
        if name != "encoder_config" and hasattr(self.encoder_config, name):
            return getattr(self.encoder_config, name)
        raise AttributeError(name)


EmbedderConfig.register_for_auto_class()