from typing import Literal

from transformers import PretrainedConfig


class LeafConfig(PretrainedConfig):
    model_type = "leaf"

    def __init__(
            self,
            num_classes: int = 2097,
            model_name: str = Literal["BAAI/bge-m3", "sentence-transformers/distiluse-base-multilingual-cased-v2"],
            **kwargs,
    ):
        self.num_classes = num_classes
        self.model_name = model_name
        super().__init__(**kwargs)