small cleanups

Author: wenting-zhao

baselines/agents.py

@@ -1,3 +1,5 @@
+import sys
+import os
 from abc import ABC, abstractmethod
 from pathlib import Path
 
@@ -7,14 +9,18 @@
 
 
 class Agents(ABC):
+    def __init__(self, max_iteration: int):
+        self.max_iteration = max_iteration
+
     @abstractmethod
     def run(self) -> None:
         """Start agent"""
         raise NotImplementedError
 
 
 class AiderAgents(Agents):
-    def __init__(self, model_name: str):
+    def __init__(self, max_iteration: int, model_name: str):
+        super().__init__(max_iteration)
         self.model = Model(model_name)
 
     def run(
@@ -37,6 +43,10 @@ def run(
         log_dir.mkdir(parents=True, exist_ok=True)
         input_history_file = log_dir / ".aider.input.history"
         chat_history_file = log_dir / ".aider.chat.history.md"
+        print(
+            f"check {os.path.abspath(chat_history_file)} for prompts and lm generations",
+            file=sys.stderr,
+        )
         io = InputOutput(
             yes=True,
             input_history_file=input_history_file,
@@ -51,4 +61,6 @@ def run(
             test_cmd=test_cmd,
             io=io,
         )
+        coder.max_reflection = self.max_iteration
+        coder.stream = False
         coder.run(message)

baselines/class_types.py

@@ -26,3 +26,4 @@ class AgentConfig:
     max_lint_info_length: int
     pre_commit_config_path: str
     run_tests: bool
+    max_iteration: int

baselines/configs/base.yaml

@@ -25,6 +25,7 @@ agent_config:
   max_unit_tests_info_length: 10000
   max_reference_info_length: 10000
   max_lint_info_length: 10000
+  max_iteration: 3
 
 hydra:
   run:

baselines/run_agent.py

@@ -1,4 +1,5 @@
 import os
+import sys
 import hydra
 from datasets import load_dataset
 import traceback
@@ -57,7 +58,7 @@ def run_agent_for_repo(
     target_edit_files = get_target_edit_files(repo_path)
 
     if agent_config.agent_name == "aider":
-        agent = AiderAgents(agent_config.model_name)
+        agent = AiderAgents(agent_config.max_iteration, agent_config.model_name)
     else:
         raise NotImplementedError(
             f"{agent_config.agent_name} is not implemented; please add your implementations in baselines/agents.py."
@@ -127,6 +128,9 @@ def main() -> None:
     ]
     assert len(filtered_dataset) > 0, "No examples available"
 
+    if len(filtered_dataset) > 1:
+        sys.stdout = open(os.devnull, "w")
+
     with tqdm(
         total=len(filtered_dataset), smoothing=0, desc="Running Aider for repos"
     ) as pbar: