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) | |