Make print_metrics.py recursive

scripts/print_results.py

@@ -1,21 +1,37 @@
 import os
 import json
-from contextlib import suppress
 
 import tyro
 
 
-def main(data_dir: str) -> None:
-    """Summarize the results collected for all models in the given directory."""
-    model_names = os.listdir(data_dir)
-    print(len(model_names), "models found")
+def main(data_dir: str, depth: int = 1) -> None:
+    """Summarize the results collected for all models in the given directory.
+    
+    Args:
+        data_dir: The directory containing the results.
+        depth: The depth of the directory tree to search. When it's 1, the
+            script expects to fine model directories directly under `data_dir`.
+            (Default: 1)
+    """
+    if depth < 1:
+        raise ValueError("depth must be >= 1")
 
-    for i, model_name in enumerate(model_names):
-        try:
-            benchmark = json.load(open(f"{data_dir}/{model_name}/benchmark.json"))
-            print(f"{i:2d} {len(benchmark):5d} results found for", model_name)
-        except json.JSONDecodeError:
-            print(f"{i:2d} [ERR] results found for {model_name}")
+    if depth == 1:
+        model_names = os.listdir(data_dir)
+        print(len(model_names), "models found in", data_dir)
+
+        for i, model_name in enumerate(model_names):
+            if not os.path.isdir(f"{data_dir}/{model_name}"):
+                continue
+            try:
+                benchmark = json.load(open(f"{data_dir}/{model_name}/benchmark.json"))
+                print(f"{i:2d} {len(benchmark):5d} results found for", model_name)
+            except json.JSONDecodeError:
+                print(f"{i:2d} [ERR] results found for {model_name}")
+
+    else:
+        for dir in os.listdir(data_dir):
+            main(f"{data_dir}/{dir}", depth - 1)
 
 
 if __name__ == "__main__":