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