| from transformers import BertPreTrainedModel, AutoModel, PretrainedConfig | |
| import sys | |
| sys.path.append("..") | |
| import torch.nn as nn | |
| from Classifier.pragformer_config import PragFormerConfig | |
| class BERT_Arch(BertPreTrainedModel): | |
| config_class = PragFormerConfig | |
| def __init__(self, config): | |
| super().__init__(config) | |
| self.bert = AutoModel.from_pretrained(config.bert['_name_or_path']) | |
| # dropout layer | |
| self.dropout = nn.Dropout(config.dropout) | |
| # relu activation function | |
| self.relu = nn.ReLU() | |
| # dense layer 1 | |
| self.fc1 = nn.Linear(self.config.bert['hidden_size'], config.fc1) | |
| # self.fc1 = nn.Linear(768, 512) | |
| # dense layer 2 (Output layer) | |
| self.fc2 = nn.Linear(config.fc1, config.fc2) | |
| # softmax activation function | |
| self.softmax = nn.LogSoftmax(dim = config.softmax_dim) | |
| # define the forward pass | |
| def forward(self, input_ids, attention_mask): | |
| # pass the inputs to the model | |
| _, cls_hs = self.bert(input_ids, attention_mask = attention_mask, return_dict=False) | |
| x = self.fc1(cls_hs) | |
| x = self.relu(x) | |
| x = self.dropout(x) | |
| # output layer | |
| x = self.fc2(x) | |
| # apply softmax activation | |
| x = self.softmax(x) | |
| return x | |
| # class BERT_Arch_new(BertPreTrainedModel): | |
| # def __init__(self, config): | |
| # super().__init__(config) | |
| # self.bert = AutoModel.from_pretrained('/home/talkad/Desktop/pragformer/PragFormer/DeepSCC-RoBERTa') | |
| # # dropout layer | |
| # self.dropout = nn.Dropout(0.2) | |
| # # relu activation function | |
| # self.relu = nn.ReLU() | |
| # # dense layer 1 | |
| # self.fc1 = nn.Linear(self.config.hidden_size, 512) | |
| # # self.fc1 = nn.Linear(768, 512) | |
| # # dense layer 2 (Output layer) | |
| # self.fc2 = nn.Linear(512, 2) | |
| # # softmax activation function | |
| # self.softmax = nn.LogSoftmax(dim = 1) | |
| # # define the forward pass | |
| # def forward(self, input_ids, attention_mask): | |
| # # pass the inputs to the model | |
| # _, cls_hs = self.bert(input_ids, attention_mask = attention_mask, return_dict=False) | |
| # x = self.fc1(cls_hs) | |
| # x = self.relu(x) | |
| # x = self.dropout(x) | |
| # # output layer | |
| # x = self.fc2(x) | |
| # # apply softmax activation | |
| # x = self.softmax(x) | |
| # return x | |
| # class BERT_Arch(nn.Module): | |
| # def __init__(self, bert): | |
| # super(BERT_Arch, self).__init__() | |
| # self.bert = bert | |
| # # dropout layer | |
| # self.dropout = nn.Dropout(0.2) | |
| # # relu activation function | |
| # self.relu = nn.ReLU() | |
| # # dense layer 1 | |
| # self.fc1 = nn.Linear(768, 512) | |
| # # dense layer 2 (Output layer) | |
| # self.fc2 = nn.Linear(512, 2) | |
| # # softmax activation function | |
| # self.softmax = nn.LogSoftmax(dim = 1) | |
| # # define the forward pass | |
| # def forward(self, input_ids, attention_mask): | |
| # # pass the inputs to the model | |
| # _, cls_hs = self.bert(input_ids, attention_mask = attention_mask, return_dict=False) | |
| # x = self.fc1(cls_hs) | |
| # x = self.relu(x) | |
| # x = self.dropout(x) | |
| # # output layer | |
| # x = self.fc2(x) | |
| # # apply softmax activation | |
| # x = self.softmax(x) | |
| # return x | |
| # def save_pretrained_model(self, path="", push=False, repo_name=""): | |
| # if not push: | |
| # self.bert.save_pretrained(path, repo_url=repo_name) | |
| # else: | |
| # self.bert.push_to_hub(repo_name) | |