|
|
from transformers import PreTrainedModel |
|
|
from transformers import AutoModelForMaskedLM, AutoTokenizer |
|
|
from pytorch_lightning.loggers import WandbLogger |
|
|
|
|
|
from src.regression.PL import FullModelPL, EncoderPL, DecoderPL |
|
|
from src.regression.HF.configs import FullModelConfigHF |
|
|
|
|
|
from config import DEVICE |
|
|
|
|
|
|
|
|
class FullModelHF(PreTrainedModel): |
|
|
config_class = FullModelConfigHF |
|
|
|
|
|
def __init__(self, config): |
|
|
|
|
|
super().__init__(config) |
|
|
|
|
|
self.tokenizer = AutoTokenizer.from_pretrained(config.tokenizer_ckpt) |
|
|
mlm_bert = AutoModelForMaskedLM.from_pretrained(config.bert_ckpt) |
|
|
self.bert = mlm_bert.distilbert |
|
|
|
|
|
encoder = EncoderPL(tokenizer=self.tokenizer, bert=self.bert).to(DEVICE) |
|
|
|
|
|
wandb_logger = WandbLogger( |
|
|
project="transformers", |
|
|
entity="sanjin_juric_fot", |
|
|
|
|
|
|
|
|
) |
|
|
|
|
|
artifact = wandb_logger.use_artifact(config.decoder_ckpt) |
|
|
artifact_dir = artifact.download() |
|
|
decoder = DecoderPL.load_from_checkpoint(artifact_dir + "/" + "model.ckpt").to(DEVICE) |
|
|
|
|
|
self.model = FullModelPL( |
|
|
encoder=encoder, |
|
|
decoder=decoder, |
|
|
layer_norm=config.layer_norm, |
|
|
nontext_features=config.nontext_features, |
|
|
).to(DEVICE) |
|
|
|
|
|
def forward(self, input): |
|
|
return self.model._get_loss(input) |
|
|
|
