from transformers import PretrainedConfig

class ViConBERTConfig(PretrainedConfig):
    model_type = "viconbert"

    def __init__(
        self,
        base_model="vinai/phobert-base",
        base_model_cache_dir="embeddings/base_models",
        hidden_dim=512,
        out_dim=768,
        dropout=0.1,
        num_layers=1,
        num_head=3,
        encoder_type="attentive",
        context_window_size=3,
        **kwargs
    ):
        super().__init__(**kwargs)
        self.base_model = base_model
        self.base_model_cache_dir = base_model_cache_dir
        self.hidden_dim = hidden_dim
        self.out_dim = out_dim
        self.dropout = dropout
        self.num_layers = num_layers
        self.num_head = num_head
        self.encoder_type = encoder_type
        self.context_window_size = context_window_size