Spaces:
Sleeping
Sleeping
| #from transformers import BertPreTrainedModel, BertForSequenceClassification, BertModel | |
| from transformers import AutoModel, PreTrainedModel | |
| from transformers.modeling_outputs import TokenClassifierOutput | |
| from torch import nn | |
| from torch.nn import CrossEntropyLoss | |
| import torch | |
| from .layers import CRF | |
| from itertools import islice | |
| NUM_PER_LAYER = 16 | |
| class BERTLstmCRF(PreTrainedModel): | |
| _keys_to_ignore_on_load_unexpected = [r"pooler"] | |
| def __init__(self, config): | |
| super().__init__(config) | |
| print(config) | |
| self.num_labels = config.num_labels | |
| self.bert = AutoModel.from_pretrained(config._name_or_path, config=config, add_pooling_layer=False) | |
| classifier_dropout = (config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob) | |
| self.dropout = nn.Dropout(config.hidden_dropout_prob) | |
| self.bilstm = nn.LSTM(config.hidden_size, (config.hidden_size) // 2, batch_first=True, bidirectional=True) | |
| self.classifier = nn.Linear(config.hidden_size, config.num_labels) | |
| self.crf = CRF(num_tags=config.num_labels, batch_first=True) | |
| if self.config.freeze == True: | |
| self.manage_freezing() | |
| #self.bert.init_weights() # load pretrained weights | |
| def manage_freezing(self): | |
| for _, param in self.bert.embeddings.named_parameters(): | |
| param.requires_grad = False | |
| num_encoders_to_freeze = self.config.num_frozen_encoder | |
| if num_encoders_to_freeze > 0: | |
| for _, param in islice(self.bert.encoder.named_parameters(), num_encoders_to_freeze*NUM_PER_LAYER): | |
| param.requires_grad = False | |
| def forward(self, | |
| input_ids=None, | |
| attention_mask=None, | |
| token_type_ids=None, | |
| position_ids=None, | |
| head_mask=None, | |
| inputs_embeds=None, | |
| labels=None, | |
| output_attentions=None, | |
| output_hidden_states=None, | |
| return_dict=None | |
| ): | |
| # Default `model.config.use_return_dict´ is `True´ | |
| return_dict = return_dict if return_dict is not None else self.config.use_return_dict | |
| outputs = self.bert(input_ids, | |
| attention_mask=attention_mask, | |
| token_type_ids=token_type_ids, | |
| position_ids=position_ids, | |
| head_mask=head_mask, | |
| inputs_embeds=inputs_embeds, | |
| output_attentions=output_attentions, | |
| output_hidden_states=output_hidden_states, | |
| return_dict=return_dict) | |
| sequence_output = outputs[0] | |
| sequence_output = self.dropout(sequence_output) | |
| lstm_output, hc = self.bilstm(sequence_output) | |
| logits = self.classifier(lstm_output) | |
| loss = None | |
| if labels is not None: | |
| # During train/test as we don't pass labels during inference | |
| loss = -1 * self.crf(logits, labels) | |
| tags = torch.Tensor(self.crf.decode(logits)) | |
| return loss, tags | |