File size: 1,713 Bytes
6a34fd4 |
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 |
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]))
|