from transformers import PretrainedConfig


class MHuBERTIPACTCFTConfig(PretrainedConfig):
    model_type = "mhubert_ipa_ctc_ft"

    def __init__(
        self,
        base_model="utter-project/mHuBERT-147",
        backbone_id=None,
        task="end-to-end CTC IPA phone recognition",
        label_type="IPA",
        label_map="ipa_map.json",
        ctc_blank_token="<blank>",
        blank_position="last",
        top_encoder_layers_finetuned=4,
        architecture=None,
        evaluation=None,
        backbone_config=None,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.base_model = base_model
        self.backbone_id = backbone_id or base_model
        self.task = task
        self.label_type = label_type
        self.label_map = label_map
        self.ctc_blank_token = ctc_blank_token
        self.blank_position = blank_position
        self.top_encoder_layers_finetuned = int(top_encoder_layers_finetuned)
        self.architecture = architecture or {
            "input_dim": 768,
            "proj_dim": 256,
            "lstm_hidden": 256,
            "lstm_layers": 2,
            "lstm_bidirectional": True,
            "dropout": 0.3,
            "n_phones": 45,
            "output_dim": 46,
            "blank_id": 45,
        }
        self.blank_id = int(self.architecture["blank_id"])
        self.output_dim = int(self.architecture["output_dim"])
        self.n_phones = int(self.architecture["n_phones"])
        self.evaluation = evaluation or {}
        self.backbone_config = backbone_config or {}