File size: 2,656 Bytes
4fb0bd1 |
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 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
import torch.nn as nn
from modules.token_embedders.pretrained_encoder import PretrainedEncoder
from utils.nn_utils import batched_index_select, gelu
class PretrainedEmbedModel(nn.Module):
"""This class acts as an embeddding layer with pre-trained model
"""
def __init__(self, cfg, vocab, rel_mlp=False):
"""This function constructs `PretrainedEmbedModel` components and
sets `PretrainedEmbedModel` parameters
Arguments:
cfg {dict} -- config parameters for constructing multiple models
vocab {Vocabulary} -- vocabulary
"""
super().__init__()
self.activation = gelu
self.pretrained_encoder = PretrainedEncoder(pretrained_model_name=cfg.pretrained_model_name,
trainable=cfg.fine_tune,
output_size=cfg.bert_output_size,
activation=self.activation,
dropout=cfg.bert_dropout)
self.encoder_output_size = self.pretrained_encoder.get_output_dims()
def forward(self, batch_inputs):
"""This function propagetes forwardly
Arguments:
batch_inputs {dict} -- batch input data
"""
if 'wordpiece_segment_ids' in batch_inputs:
batch_seq_pretrained_encoder_repr, batch_cls_repr = self.pretrained_encoder(
batch_inputs['wordpiece_tokens'], batch_inputs['wordpiece_segment_ids'])
else:
batch_seq_pretrained_encoder_repr, batch_cls_repr = self.pretrained_encoder(
batch_inputs['wordpiece_tokens'])
batch_seq_tokens_encoder_repr = batched_index_select(batch_seq_pretrained_encoder_repr,
batch_inputs['wordpiece_tokens_index'])
batch_inputs['seq_encoder_reprs'] = batch_seq_tokens_encoder_repr
# if not self.rel_mlp:
# batch_seq_tokens_encoder_repr = batched_index_select(batch_seq_pretrained_encoder_repr,
# batch_inputs['wordpiece_tokens_index'])
# batch_inputs['seq_encoder_reprs'] = batch_seq_tokens_encoder_repr
# else:
# batch_inputs['seq_encoder_reprs'] = batch_seq_pretrained_encoder_repr
batch_inputs['seq_cls_repr'] = batch_cls_repr
def get_hidden_size(self):
"""This function returns embedding dimensions
Returns:
int -- embedding dimensitons
"""
return self.encoder_output_size
|