allow specifying an oe-eval branch

src/cookbook/cli/eval.py

@@ -233,6 +233,12 @@ def convert_checkpoint(
 )
 @click.option("-g", "--use-gantry", is_flag=True, help="Submit jobs with gantry directly.")
 @click.option("--beaker-retries", type=int, default=0, help="Number of retries for failed evals")
+@click.option(
+    "--oe-eval-branch",
+    type=str,
+    default=None,
+    help="Branch of the oe-eval toolkit to use; if not provided, use main",
+)
 @click.option(
     "--oe-eval-commit",
     type=str,
@@ -300,6 +306,7 @@ def convert_checkpoint(
     help="Suffix to add to the run name",
 )
 def evaluate_model(
+    oe_eval_branch: str,
     oe_eval_commit: str,
     checkpoint_path: str,
     aws_access_key_id: str,
@@ -353,6 +360,7 @@ def evaluate_model(
     tasks = [e for t in tasks for e in (ALL_EVAL_TASKS.get(t.lstrip("*"), [t]) if t.startswith("*") else [t])]
 
     evaluate_checkpoint(
+        oe_eval_branch=oe_eval_branch,
         oe_eval_commit=oe_eval_commit,
         checkpoint_path=checkpoint_path,
         aws_access_key_id=aws_access_key_id,

src/cookbook/cli/utils.py

@@ -173,6 +173,7 @@ def install_beaker_py(
 
 def install_oe_eval(
     commit_hash: Optional[str],
+    commit_branch: Optional[str],
     env: Optional[PythonEnv] = None,
     no_dependencies: bool = True,
     is_editable: bool = False,
@@ -182,7 +183,7 @@ def install_oe_eval(
     print("Installing beaker and gantry clients...")
     install_beaker_py(env)
 
-    oe_eval_dir = clone_repository(OE_EVAL_GIT_URL, commit_hash)
+    oe_eval_dir = clone_repository(OE_EVAL_GIT_URL, commit_hash, commit_branch)
 
     print(f"Installing OE-Eval from {oe_eval_dir}" + (" in editable mode" if is_editable else "") + "...")
     cmd = [
@@ -362,7 +363,7 @@ def make_eval_run_name(checkpoint_path: str, add_bos_token: bool, name_suffix: s
     )
 
 
-def clone_repository(git_url: str, commit_hash: Optional[str] = None) -> str:
+def clone_repository(git_url: str, commit_hash: Optional[str] = None, commit_branch: Optional[str] = None) -> str:
     # current directory
     current_dir = os.getcwd()
 
@@ -383,11 +384,17 @@ def clone_repository(git_url: str, commit_hash: Optional[str] = None) -> str:
         # Execute clone
         subprocess.run(cmd, check=True)
 
+        if commit_branch:
+            # Change directory to the cloned repo
+            os.chdir(tmp_dir)
+            subprocess.run(shlex.split(f"git fetch origin {commit_branch}:refs/remotes/origin/{commit_branch}"), check=True)
+            subprocess.run(shlex.split(f"git checkout -b {commit_branch} origin/{commit_branch}"), check=True)
+
         if commit_hash:
             # Change directory to the cloned repo
             os.chdir(tmp_dir)
             subprocess.run(shlex.split(f"git fetch origin '{commit_hash}'"), check=True)
-            subprocess.run(shlex.split(f"git checkout  '{commit_hash}'"), check=True)
+            subprocess.run(shlex.split(f"git checkout '{commit_hash}'"), check=True)
 
         return tmp_dir
 

src/cookbook/eval/evaluation.py

@@ -25,6 +25,7 @@
 
 def evaluate_checkpoint(
     oe_eval_commit: str,
+    oe_eval_branch: str,
     checkpoint_path: str,
     aws_access_key_id: str,
     aws_secret_access_key: str,
@@ -66,6 +67,7 @@ def evaluate_checkpoint(
     oe_eval_dir = install_oe_eval(
         env=env,
         commit_hash=oe_eval_commit,
+        commit_branch=oe_eval_branch,
         is_editable=use_gantry,
     )