from transformers import BertConfig


class BertMultiTaskConfig(BertConfig):
    model_type = "bert"

    def __init__(self, tasks: dict[str, int] | None = None, **kwargs):
        self.tasks = tasks
        super().__init__(**kwargs)