File size: 462 Bytes
8e32295 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 | 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)
|