Update README.md

README.md

@@ -95,7 +95,7 @@ class BERTClass(torch.nn.Module):
         output = self.classifier(pooler)
         return output
         
-def do_predict(model, tokenizer):
+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)
@@ -119,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(model_read, tokenizer_read)
+actual_predictions_read = do_predict(model_read, tokenizer_read, test_df)
 
 test_df['readability'] = ['readable' if i==1 else 'not_reabale' for i in actual_predictions_read]