| import torch.nn as nn |
|
|
| from modules.token_embedders.bert_encoder import BertEncoder |
| from utils.nn_utils import batched_index_select, gelu |
|
|
|
|
| class BertEmbedModel(nn.Module): |
| """This class acts as an embeddding layer with bert model |
| """ |
| def __init__(self, cfg, vocab, rel_mlp=False): |
| """This function constructs `BertEmbedModel` components and |
| sets `BertEmbedModel` parameters |
| |
| Arguments: |
| cfg {dict} -- config parameters for constructing multiple models |
| vocab {Vocabulary} -- vocabulary |
| """ |
|
|
| super().__init__() |
| self.rel_mlp = rel_mlp |
| self.activation = gelu |
| self.bert_encoder = BertEncoder(bert_model_name=cfg.bert_model_name, |
| trainable=cfg.fine_tune, |
| output_size=cfg.bert_output_size, |
| activation=self.activation, |
| dropout=cfg.bert_dropout) |
| self.encoder_output_size = self.bert_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_bert_encoder_repr, batch_cls_repr = self.bert_encoder( |
| batch_inputs['wordpiece_tokens'], batch_inputs['wordpiece_segment_ids']) |
| else: |
| batch_seq_bert_encoder_repr, batch_cls_repr = self.bert_encoder( |
| batch_inputs['wordpiece_tokens']) |
| |
| if not self.rel_mlp: |
| batch_seq_tokens_encoder_repr = batched_index_select(batch_seq_bert_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_bert_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 |
|
|
