File size: 669 Bytes
82c0c38 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 |
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)
|