| import torch |
| import transformers |
| from torch.utils.data import Dataset, DataLoader |
| from transformers import RobertaModel, RobertaTokenizer, BertModel, BertTokenizer |
| import pandas as pd |
|
|
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
| MAX_LEN = 128 |
| BATCH_SIZE = 20 |
| text_col_name = 'sentence' |
|
|
| def scoring_data_prep(dataset): |
| out = [] |
| target = [] |
| mask = [] |
| |
| for i in range(len(dataset)): |
| rec = dataset[i] |
| out.append(rec['ids'].reshape(-1,MAX_LEN)) |
| mask.append(rec['mask'].reshape(-1,MAX_LEN)) |
|
|
| out_stack = torch.cat(out, dim = 0) |
| mask_stack = torch.cat(mask, dim =0 ) |
| out_stack = out_stack.to(device, dtype = torch.long) |
| mask_stack = mask_stack.to(device, dtype = torch.long) |
|
|
| return out_stack, mask_stack |
|
|
| class Triage(Dataset): |
| """ |
| This is a subclass of torch packages Dataset class. It processes input to create ids, masks and targets required for model training. |
| """ |
|
|
| def __init__(self, dataframe, tokenizer, max_len, text_col_name): |
| self.len = len(dataframe) |
| self.data = dataframe |
| self.tokenizer = tokenizer |
| self.max_len = max_len |
| self.text_col_name = text_col_name |
| |
|
|
| def __getitem__(self, index): |
| title = str(self.data[self.text_col_name][index]) |
| title = " ".join(title.split()) |
| inputs = self.tokenizer.encode_plus( |
| title, |
| None, |
| add_special_tokens=True, |
| max_length=self.max_len, |
| pad_to_max_length=True, |
| return_token_type_ids=True, |
| truncation=True, |
| ) |
| ids = inputs["input_ids"] |
| mask = inputs["attention_mask"] |
|
|
| return { |
| "ids": torch.tensor(ids, dtype=torch.long), |
| "mask": torch.tensor(mask, dtype=torch.long), |
| |
| } |
|
|
| def __len__(self): |
| return self.len |
|
|
| class BERTClass(torch.nn.Module): |
| def __init__(self, num_class, task): |
| super(BERTClass, self).__init__() |
| self.num_class = num_class |
| if task =="sustanability": |
| self.l1 = RobertaModel.from_pretrained("roberta-base") |
| else: |
| self.l1 = BertModel.from_pretrained("ProsusAI/finbert") |
| self.pre_classifier = torch.nn.Linear(768, 768) |
| self.dropout = torch.nn.Dropout(0.3) |
| self.classifier = torch.nn.Linear(768, self.num_class) |
| self.history = dict() |
|
|
| def forward(self, input_ids, attention_mask): |
| output_1 = self.l1(input_ids=input_ids, attention_mask=attention_mask) |
| hidden_state = output_1[0] |
| pooler = hidden_state[:, 0] |
| pooler = self.pre_classifier(pooler) |
| pooler = torch.nn.ReLU()(pooler) |
| pooler = self.dropout(pooler) |
| output = self.classifier(pooler) |
| return output |
|
|
| def do_predict(model, tokenizer, test_df): |
| test_set = Triage(test_df, tokenizer, MAX_LEN, text_col_name) |
| test_params = {'batch_size' : BATCH_SIZE, 'shuffle': False, 'num_workers':0} |
| test_loader = DataLoader(test_set, **test_params) |
| out_stack, mask_stack = scoring_data_prep(dataset = test_set) |
| n = 0 |
| combined_output = [] |
| model.eval() |
| with torch.no_grad(): |
| while n < test_df.shape[0]: |
| output = model(out_stack[n:n+BATCH_SIZE,:],mask_stack[n:n+BATCH_SIZE,:]) |
| n = n + BATCH_SIZE |
| combined_output.append(output) |
| combined_output = torch.cat(combined_output, dim = 0) |
| preds = torch.argsort(combined_output, axis = 1, descending = True) |
| preds = preds.to('cpu') |
| actual_predictions = [i[0] for i in preds.tolist()] |
| combined_output = combined_output.to('cpu') |
| prob_predictions= [i[1] for i in combined_output.tolist()] |
| return (actual_predictions, prob_predictions) |
| |
