Spaces:
Running
Running
| import torch.nn as nn | |
| from bert import BERT | |
| class BERTSM(nn.Module): | |
| """ | |
| BERT Sequence Model | |
| Masked Sequence Model | |
| """ | |
| def __init__(self, bert: BERT, vocab_size): | |
| """ | |
| :param bert: BERT model which should be trained | |
| :param vocab_size: total vocab size for masked_lm | |
| """ | |
| super().__init__() | |
| self.bert = bert | |
| self.mask_lm = MaskedSequenceModel(self.bert.hidden, vocab_size) | |
| self.same_student = SameStudentPrediction(self.bert.hidden) | |
| def forward(self, x, segment_label, pred=False): | |
| x = self.bert(x, segment_label) | |
| # torch.Size([32, 200, 512]) | |
| # print("???????????? ",x.shape) | |
| if pred: | |
| return x[:, 0], self.mask_lm(x), self.same_student(x) | |
| else: | |
| return x[:, 0], self.mask_lm(x) | |
| class MaskedSequenceModel(nn.Module): | |
| """ | |
| predicting origin token from masked input sequence | |
| n-class classification problem, n-class = vocab_size | |
| """ | |
| def __init__(self, hidden, vocab_size): | |
| """ | |
| :param hidden: output size of BERT model | |
| :param vocab_size: total vocab size | |
| """ | |
| super().__init__() | |
| self.linear = nn.Linear(hidden, vocab_size) | |
| self.softmax = nn.LogSoftmax(dim=-1) | |
| def forward(self, x): | |
| return self.softmax(self.linear(x)) | |
| class SameStudentPrediction(nn.Module): | |
| def __init__(self, hidden): | |
| """ | |
| :param hidden: BERT model output size | |
| """ | |
| super().__init__() | |
| self.linear = nn.Linear(hidden, 2) | |
| self.softmax = nn.LogSoftmax(dim=-1) | |
| def forward(self, x): | |
| return self.softmax(self.linear(x[:, 0])) | |