Spaces:
Sleeping
Sleeping
| # -*- coding: utf-8 -*- | |
| r""" | |
| BERT Encoder | |
| ============== | |
| Pretrained BERT encoder from Hugging Face. | |
| """ | |
| from argparse import Namespace | |
| from typing import Dict | |
| import torch | |
| from transformers import AutoModel | |
| from polos.models.encoders.encoder_base import Encoder | |
| from polos.tokenizers_ import HFTextEncoder | |
| from torchnlp.utils import lengths_to_mask | |
| class BERTEncoder(Encoder): | |
| """BERT encoder. | |
| :param tokenizer: BERT text encoder. | |
| :param hparams: ArgumentParser. | |
| """ | |
| def __init__( | |
| self, | |
| tokenizer: HFTextEncoder, | |
| hparams: Namespace, | |
| ) -> None: | |
| super().__init__(tokenizer) | |
| self.model = AutoModel.from_pretrained(hparams.pretrained_model) | |
| self.model.encoder.output_hidden_states = True | |
| self._output_units = self.model.config.hidden_size | |
| self._n_layers = self.model.config.num_hidden_layers + 1 | |
| self._max_pos = self.model.config.max_position_embeddings | |
| def from_pretrained(cls, hparams: Namespace) -> Encoder: | |
| """Function that loads a pretrained encoder from Hugging Face. | |
| :param hparams: Namespace. | |
| :return: Encoder model | |
| """ | |
| tokenizer = HFTextEncoder(model=hparams.pretrained_model) | |
| model = BERTEncoder(tokenizer=tokenizer, hparams=hparams) | |
| return model | |
| def freeze_embeddings(self) -> None: | |
| """ Frezees the embedding layer of the network to save some memory while training. """ | |
| for param in self.model.embeddings.parameters(): | |
| param.requires_grad = False | |
| def layerwise_lr(self, lr: float, decay: float): | |
| """ | |
| :return: List with grouped model parameters with layer-wise decaying learning rate | |
| """ | |
| # Embedding Layer | |
| opt_parameters = [ | |
| { | |
| "params": self.model.embeddings.parameters(), | |
| "lr": lr * decay ** (self.num_layers), | |
| } | |
| ] | |
| # All layers | |
| opt_parameters += [ | |
| { | |
| "params": self.model.encoder.layer[l].parameters(), | |
| "lr": lr * decay ** l, | |
| } | |
| for l in range(self.num_layers - 2, 0, -1) | |
| ] | |
| return opt_parameters | |
| def forward( | |
| self, tokens: torch.Tensor, lengths: torch.Tensor | |
| ) -> Dict[str, torch.Tensor]: | |
| """ | |
| Encodes a batch of sequences. | |
| :param tokens: Torch tensor with the input sequences [batch_size x seq_len]. | |
| :param lengths: Torch tensor with the lenght of each sequence [seq_len]. | |
| :return: Dictionary with `sentemb` (tensor with dims [batch_size x output_units]), `wordemb` | |
| (tensor with dims [batch_size x seq_len x output_units]), `mask` (input mask), | |
| `all_layers` (List with word_embeddings from all layers), `extra` (tuple with the | |
| last_hidden_state, the pooler_output representing the entire sentence and the word | |
| embeddings for all BERT layers). | |
| """ | |
| mask = lengths_to_mask(lengths, device=tokens.device) | |
| last_hidden_states, pooler_output, all_layers = self.model( | |
| tokens, mask, output_hidden_states=True, return_dict=False | |
| ) | |
| return { | |
| "sentemb": pooler_output, | |
| "wordemb": last_hidden_states, | |
| "all_layers": all_layers, | |
| "mask": mask, | |
| "extra": (last_hidden_states, pooler_output, all_layers), | |
| } | |