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