leaf-base / configuration_leaf.py

Upload model

8e32295 verified over 2 years ago

462 Bytes

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)

	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)