| 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 |
|
|
| |
| |
| |
| |
| |
| |
|
|
|
|
| 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 |
|
|
