| from transformers import PretrainedConfig |
| from src.regression.PL import EncoderPL, DecoderPL |
| from typing import List |
|
|
|
|
| class FullModelConfigHF(PretrainedConfig): |
|
|
| model_type = "full_model" |
|
|
| def __init__( |
| self, |
| tokenizer_ckpt: str = "", |
| bert_ckpt: str = "", |
| decoder_ckpt: str = "", |
| layer_norm: bool = True, |
| nontext_features: List[str] = ["aov"], |
| **kwargs, |
| ): |
|
|
| self.tokenizer_ckpt = tokenizer_ckpt |
| self.bert_ckpt = bert_ckpt |
| self.decoder_ckpt = decoder_ckpt |
| self.nontext_features = nontext_features |
| self.layer_norm = layer_norm |
| super().__init__(**kwargs) |
|
|
