[evals] minor refactoring cli (openai#387)

evals/cli/oaieval.py

@@ -15,7 +15,7 @@
 import evals.base
 import evals.record
 from evals.base import ModelSpec, ModelSpecs
-from evals.registry import registry
+from evals.registry import Registry
 
 logger = logging.getLogger(__name__)
 
@@ -24,7 +24,7 @@ def _purple(str):
  return f"\033[1;35m{str}\033[0m"
 
 
-def parse_args(args=sys.argv[1:]) -> argparse.Namespace:
+def get_parser() -> argparse.ArgumentParser:
  parser = argparse.ArgumentParser(description="Run evals through the API")
  parser.add_argument("model", type=str, help="Name of a completion model.")
  parser.add_argument("eval", type=str, help="Name of an eval. See registry.")
@@ -44,7 +44,7 @@ def parse_args(args=sys.argv[1:]) -> argparse.Namespace:
  parser.add_argument("--local-run", action=argparse.BooleanOptionalAction, default=True)
  parser.add_argument("--dry-run", action=argparse.BooleanOptionalAction, default=False)
  parser.add_argument("--dry-run-logging", action=argparse.BooleanOptionalAction, default=True)
- return parser.parse_args(args)
+ return parser
 
 
 def n_ctx_from_model_name(model_name: str) -> Optional[int]:
@@ -122,7 +122,7 @@ def api_model_ids(self):
  return [m["id"] for m in openai.Model.list()["data"]]
 
 
-def run(args):
+def run(args, registry: Optional[Registry] = None):
  if args.debug:
  logging.getLogger().setLevel(logging.DEBUG)
 
@@ -131,6 +131,7 @@ def run(args):
  if args.max_samples is not None:
  evals.eval.set_max_samples(args.max_samples)
 
+ registry = registry or Registry()
  eval_spec = registry.get_eval(args.eval)
  assert (
  eval_spec is not None
@@ -224,7 +225,8 @@ def to_number(x):
 
 
 def main():
- args = parse_args()
+ parser = get_parser()
+ args = parser.parse_args(sys.argv[1:])
  logging.basicConfig(
  format="[%(asctime)s] [%(filename)s:%(lineno)d] %(message)s",
  level=logging.INFO,

evals/cli/oaievalset.py

@@ -5,6 +5,7 @@
 import json
 import subprocess
 from pathlib import Path
+from typing import Optional
 
 from evals.registry import Registry
 
@@ -41,7 +42,7 @@ def highlight(str: str) -> str:
  return f"\033[1;32m>>> {str}\033[0m"
 
 
-def main() -> None:
+def get_parser() -> argparse.ArgumentParser:
  parser = argparse.ArgumentParser(description="Run eval sets through the API")
  parser.add_argument("model", type=str, help="Name of a completion model.")
  parser.add_argument("eval_set", type=str, help="Name of eval set. See registry.")
@@ -57,9 +58,11 @@ def main() -> None:
  default=True,
  help="Exit if any oaieval command fails.",
  )
- args, unknown_args = parser.parse_known_args()
+ return parser
 
- registry = Registry()
+
+def run(args, unknown_args, registry: Optional[Registry] = None) -> None:
+ registry = registry or Registry()
  commands: list[Task] = []
  eval_set = registry.get_eval_set(args.eval_set)
  for eval in registry.get_evals(eval_set.evals):
@@ -92,5 +95,11 @@ def main() -> None:
  print(highlight("All done!"))
 
 
+def main() -> None:
+ parser = get_parser()
+ args, unknown_args = parser.parse_known_args()
+ run(args, unknown_args)
+
+
 if __name__ == "__main__":
  main()

evals/registry.py

@@ -21,10 +21,6 @@
 
 DEFAULT_PATHS = [Path(__file__).parents[0].resolve() / "registry", Path.home() / ".evals"]
 
-DEFAULT_SYSTEM_PATHS = [
- Path(__file__).parents[0].resolve() / "registry",
-]
-
 
 class Registry:
  def __init__(self, registry_paths: Sequence[Union[str, Path]] = DEFAULT_PATHS):