"""
© Battelle Memorial Institute 2023
Made available under the GNU General Public License v 2.0

BECAUSE THE PROGRAM IS LICENSED FREE OF CHARGE, THERE IS NO WARRANTY
FOR THE PROGRAM, TO THE EXTENT PERMITTED BY APPLICABLE LAW.  EXCEPT WHEN
OTHERWISE STATED IN WRITING THE COPYRIGHT HOLDERS AND/OR OTHER PARTIES
PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY OF ANY KIND, EITHER EXPRESSED
OR IMPLIED, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE.  THE ENTIRE RISK AS
TO THE QUALITY AND PERFORMANCE OF THE PROGRAM IS WITH YOU.  SHOULD THE
PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF ALL NECESSARY SERVICING,
REPAIR OR CORRECTION.
"""

from transformers import PretrainedConfig


class FupBERTConfig(PretrainedConfig):
    model_type = "FupBERT"

    def __init__(
            self,
            ntoken=608,
            ninp=768,
            nhead=12,
            nhid=3072,
            nlayers=12,
            token_reduction='mean',
            padding_idx=0,
            cls_idx=1,
            edge_idx=2,
            num_out=1,
            dropout=0.1,
            **kwargs):

        # Store the input parameters
        self.ntoken = ntoken
        self.ninp = ninp
        self.nhead = nhead
        self.nhid = nhid
        self.nlayers = nlayers
        self.token_reduction = token_reduction
        self.padding_idx = padding_idx
        self.cls_idx = cls_idx
        self.edge_idx = edge_idx
        self.num_out = num_out
        self.dropout = dropout
        super().__init__(**kwargs)