from typing import Literal

from transformers import PretrainedConfig


class DSIRConfig(PretrainedConfig):
    model_type = "dsir"
    is_composition = False

    def __init__(
        self,
        n: int = 2,
        num_buckets: int = 10_000,
        laplace_smoothing: float = 1e-4,
        **kwargs
    ):
        super().__init__(**kwargs)
        self.n = n
        self.num_buckets = num_buckets
        self.laplace_smoothing = laplace_smoothing