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_sustain.to(device)
 model_sustain.load_state_dict(torch.load('pytorch_model.bin', map_location=device)['model_state_dict'])
 
 tokenizer_sus = BertTokenizer.from_pretrained('roberta-base')
-actual_predictions_sus = do_predict(model_sustain, tokenizer_sus)
+actual_predictions_sus = do_predict(model_sustain, tokenizer_sus, test_df)
 
 test_df['sustainability'] = ['sustainable' if i==0 else 'unsustainable' for i in actual_predictions_read]
 ```