| from transformers import AutoModel | |
| from torch import nn | |
| class BERTClassifier(nn.Module): | |
| def __init__(self, bert_path="cointegrated/rubert-tiny2"): | |
| super().__init__() | |
| self.bert = AutoModel.from_pretrained(bert_path) | |
| for param in self.bert.parameters(): | |
| param.requires_grad = False | |
| self.linear = nn.Sequential( | |
| nn.Linear(312, 150), | |
| nn.Dropout(0.1), | |
| nn.ReLU(), | |
| nn.Linear(150, 1), | |
| nn.Sigmoid() | |
| ) | |
| def forward(self, x, masks): | |
| bert_out = self.bert(x, attention_mask=masks)[0][:, 0, :] | |
| out = self.linear(bert_out) | |
| return out |