Update eval_logical_reasoning_all_epochs.py

llm_toolkit/eval_logical_reasoning_all_epochs.py

@@ -11,12 +11,14 @@ def evaluate_model_all_epochs_v2(
     load_in_4bit=True,
     num_of_entries=-1,
     result_file=None,
+    batch_size=2,
 ):
     new_env = os.environ.copy()
     new_env["MODEL_NAME"] = model_name
     model = model_name.split("/")[-1]
 
     new_env["LOAD_IN_4BIT"] = "true" if load_in_4bit else "false"
+    new_env["BATCH_SIZE"] = str(batch_size)
     if result_file is not None:
         new_env["LOGICAL_REASONING_RESULTS_PATH"] = result_file
 
@@ -80,6 +82,7 @@ if __name__ == "__main__":
     start_epoch = int(os.getenv("START_EPOCH", 0))
     load_in_4bit = os.getenv("LOAD_IN_4BIT", "true").lower() == "true"
     result_file = os.getenv("LOGICAL_REASONING_RESULTS_PATH", None)
+    batch_size = int(os.getenv("BATCH_SIZE", 2))
 
     num_of_entries = int(sys.argv[1]) if len(sys.argv) > 1 else -1
 
@@ -90,4 +93,5 @@ if __name__ == "__main__":
         load_in_4bit=load_in_4bit,
         num_of_entries=num_of_entries,
         result_file=result_file,
+        batch_size=batch_size,
     )