"""Configuration for the Scripture Genome head (the base is separate)."""

from transformers import PretrainedConfig


class ScriptureGenomeConfig(PretrainedConfig):
    model_type = "scripture_genome_head"

    def __init__(
        self,
        input_dim=1024,
        hidden_dim=4096,
        output_dim=568,
        encoder_name="BAAI/bge-m3",
        encoder_revision="5617a9f61b028005a4858fdac845db406aefb181",
        max_length=512,
        checkpoint_sha256=None,
        gene_ids_sha256=None,
        track="question",
        **kwargs,
    ):
        super().__init__(**kwargs)
        if (input_dim, hidden_dim, output_dim) != (1024, 4096, 568):
            raise ValueError("This release supports only the selected 1024 -> 4096 -> 568 head")
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.output_dim = output_dim
        self.encoder_name = encoder_name
        self.encoder_revision = encoder_revision
        self.max_length = max_length
        self.checkpoint_sha256 = checkpoint_sha256
        self.gene_ids_sha256 = gene_ids_sha256
        self.track = track
        self.pooling = "cls_l2_normalized"


ScriptureGenomeConfig.register_for_auto_class()