Update README.md

README.md

@@ -43,13 +43,13 @@ 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, category_col):
+    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
-        self.category_col = category_col
+        
 
     def __getitem__(self, index):
         title = str(self.data[self.text_col_name][index])
@@ -69,9 +69,7 @@ class Triage(Dataset):
         return {
             "ids": torch.tensor(ids, dtype=torch.long),
             "mask": torch.tensor(mask, dtype=torch.long),
-            "targets": torch.tensor(
-                self.data[self.category_col][index], dtype=torch.long
-            ),
+            
         }
 
     def __len__(self):
@@ -97,7 +95,7 @@ class BERTClass(torch.nn.Module):
         output = self.classifier(pooler)
         return output
         
-def do_predict(tokenizer):
+def do_predict(model, tokenizer):
   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)
@@ -121,7 +119,7 @@ model_read.to(device)
 model_read.load_stat_dict(torch.load('pytorch_model.bin', map_location=device)['model_state_dict'])
 
 tokenizer_read = BertTokenizer.from_pretrained('ProsusAI/finbert')
-actual_predictions_read = do_predict(tokenizer_read)
+actual_predictions_read = do_predict(model_read, tokenizer_read)
 
 test_df['readability'] = ['readable' if i==1 else 'not_reabale' for i in actual_predictions_read]