initial commit

predict.py

@@ -13,8 +13,8 @@ from transformers import (
 from transformers.data.processors.squad import SquadResult, SquadV2Processor, SquadExample
 from transformers.data.metrics.squad_metrics import compute_predictions_logits
 
+
 def run_prediction(question_texts, context_text, model_path, n_best_size=1):
-    ### Setting hyperparameters
     max_seq_length = 512
     doc_stride = 256
     n_best_size = n_best_size
@@ -102,7 +102,7 @@ def run_prediction(question_texts, context_text, model_path, n_best_size=1):
     print(all_results)
 
     output_nbest_file = None
-    if n_best_size > 1:
+    if int(n_best_size) > 1:
         output_nbest_file = "nbest.json"
 
     timer = time.time()
@@ -123,4 +123,4 @@ def run_prediction(question_texts, context_text, model_path, n_best_size=1):
     )
     print(f'Logits converted to predictions in {time.time()-timer} seconds')
 
-    return final_predictions
+    return final_predictions