can process !function

lm_eval/utils.py

@@ -7,6 +7,7 @@ import pathlib
 import functools
 import subprocess
 import collections
+import importlib.util
 
 from typing import List
 
@@ -256,6 +257,60 @@ def get_git_commit_hash():
     return git_hash
 
 
+def import_function(loader, node):
+
+    function_name = loader.construct_scalar(node)
+    yaml_path = os.path.dirname(loader.name)
+
+    module_name, function_name = function_name.split('.')
+    module_path = os.path.join(yaml_path, "{}.py".format(module_name))
+
+    spec = importlib.util.spec_from_file_location(module_name, module_path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+
+    function = getattr(module, function_name)
+    return function
+
+# Add the import_function constructor to the YAML loader
+yaml.add_constructor('!function', import_function)
+
+
+def load_yaml_config(yaml_path):
+    with open(yaml_path, 'rb') as file:
+        yaml_config = yaml.full_load(file)
+        yaml_dir = os.path.dirname(yaml_path)
+        
+        if 'include' in yaml_config:
+            include_path = yaml_config['include']
+            del yaml_config['include']
+
+            if type(include_path) == str:
+                include_path = [include_path]
+            
+            # Load from the last one first
+            include_path.reverse()
+            final_yaml_config = {}
+            for path in include_path:
+
+                # Assumes that path is a full path.
+                # If not found, assume the included yaml 
+                # is in the same dir as the original yaml
+                if not os.path.isfile(path):
+                    path = os.path.join(yaml_dir, path)
+
+                try:
+                    included_yaml_config = load_yaml_config(path)
+                    final_yaml_config.update(included_yaml_config)
+                except:
+                    # If failed to load, ignore
+                    pass
+
+            final_yaml_config.update(yaml_config)
+            return final_yaml_config
+        return yaml_config
+
+
 env = Environment(loader=BaseLoader, undefined=StrictUndefined)
 
 def apply_template(template, doc):