Spaces:
Sleeping
Sleeping
File size: 3,252 Bytes
749745d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 |
from copy import deepcopy
import numpy as np
import torch
from torch import nn
# from pytorch_pretrained_bert.modeling import BertModel
from transformers import BertConfig, RobertaConfig, RobertaModel, BertModel
class BertEncoder(nn.Module):
def __init__(self, cfg):
super(BertEncoder, self).__init__()
self.cfg = cfg
self.bert_name = cfg.MODEL.LANGUAGE_BACKBONE.MODEL_TYPE
print("LANGUAGE BACKBONE USE GRADIENT CHECKPOINTING: ", self.cfg.MODEL.LANGUAGE_BACKBONE.USE_CHECKPOINT)
if self.bert_name == "bert-base-uncased":
config = BertConfig.from_pretrained(self.bert_name)
config.gradient_checkpointing = self.cfg.MODEL.LANGUAGE_BACKBONE.USE_CHECKPOINT
self.model = BertModel.from_pretrained(self.bert_name, add_pooling_layer=False, config=config)
self.language_dim = 768
elif self.bert_name == "roberta-base":
config = RobertaConfig.from_pretrained(self.bert_name)
config.gradient_checkpointing = self.cfg.MODEL.LANGUAGE_BACKBONE.USE_CHECKPOINT
self.model = RobertaModel.from_pretrained(self.bert_name, add_pooling_layer=False, config=config)
self.language_dim = 768
else:
raise NotImplementedError
self.num_layers = cfg.MODEL.LANGUAGE_BACKBONE.N_LAYERS
def forward(self, x):
input = x["input_ids"]
mask = x["attention_mask"]
if self.cfg.MODEL.DYHEAD.FUSE_CONFIG.USE_DOT_PRODUCT_TOKEN_LOSS:
# with padding, always 256
outputs = self.model(
input_ids=input,
attention_mask=mask,
output_hidden_states=True,
)
# outputs has 13 layers, 1 input layer and 12 hidden layers
encoded_layers = outputs.hidden_states[1:]
features = None
features = torch.stack(encoded_layers[-self.num_layers :], 1).mean(1)
# language embedding has shape [len(phrase), seq_len, language_dim]
features = features / self.num_layers
embedded = features * mask.unsqueeze(-1).float()
aggregate = embedded.sum(1) / (mask.sum(-1).unsqueeze(-1).float())
else:
# without padding, only consider positive_tokens
max_len = (input != 0).sum(1).max().item()
outputs = self.model(
input_ids=input[:, :max_len],
attention_mask=mask[:, :max_len],
output_hidden_states=True,
)
# outputs has 13 layers, 1 input layer and 12 hidden layers
encoded_layers = outputs.hidden_states[1:]
features = None
features = torch.stack(encoded_layers[-self.num_layers :], 1).mean(1)
# language embedding has shape [len(phrase), seq_len, language_dim]
features = features / self.num_layers
embedded = features * mask[:, :max_len].unsqueeze(-1).float()
aggregate = embedded.sum(1) / (mask.sum(-1).unsqueeze(-1).float())
ret = {"aggregate": aggregate, "embedded": embedded, "masks": mask, "hidden": encoded_layers[-1]}
return ret
|