from transformers import PretrainedConfig


# From https://huggingface.co/klue/bert-base
_default_config = {
    "architectures": ["BertForMaskedLM"],
    "attention_probs_dropout_prob": 0.1,
    "hidden_act": "gelu",
    "hidden_dropout_prob": 0.1,
    "hidden_size": 768,
    "initializer_range": 0.02,
    "intermediate_size": 3072,
    "layer_norm_eps": 1e-12,
    "max_position_embeddings": 512,
    "model_type": "bert",
    "num_attention_heads": 12,
    "num_hidden_layers": 12,
    "pad_token_id": 0,
    "type_vocab_size": 2,
    "vocab_size": 32000
}

class DprConfig(PretrainedConfig):
    model_type = "dpr"

    def __init__(self, qst_config=_default_config, ctx_config=_default_config, **kwargs):
        self.qst_config = qst_config
        self.ctx_config = ctx_config
        super().__init__(**kwargs)