from typing import Tuple

from transformers import BertConfig


class TTCompressedBertConfig(BertConfig):
    """Class TTCompressedBertConfig defines a configuration for TT-compressed
    BERT. Here, we split shape to input and output shape in order to serialize
    them to different fields in JSON.
    """

    def __init__(self, *args, shape_in: Tuple[int] = (),
                 shape_out: Tuple[int] = (), rank: int = 128, **kwargs):
        super().__init__(*args, **kwargs)
        self.shape_in = shape_in
        self.shape_out = shape_out
        self.rank = rank


TTCompressedBertConfig.register_for_auto_class()