added more registry

lm_eval/api/register.py

@@ -29,3 +29,77 @@ def register_group(name):
         return func
 
     return wrapper
+
+
+metric_registry = {}
+aggregation_registry = {}
+default_aggregation_registry = {}
+higher_is_better_registry = {}
+output_type_registry = {}
+metric2func_index = {}
+func2metric_index = {}
+aggregation2func_index = {}
+func2aggregation_index = {}
+
+
+def register_metric(name):
+    def wrapper(func):
+
+        metric_registry[name] = func
+        func2metric_index[func.__name__] = name
+        metric2func_index[name] = func.__name__
+        return func
+
+    return wrapper
+
+
+def register_aggregation(name):
+    def wrapper(func):
+
+        aggregation_registry[name] = func
+        func2aggregation_index[func.__name__] = name
+        aggregation2func_index[name] = func.__name__
+        return func
+
+    return wrapper
+
+
+def register_default_aggregation(aggregation):
+    def wrapper(func):
+
+        if aggregation in aggregation_registry:
+            metric_name = func2metric_index[func.__name__]
+            default_aggregation_registry[metric_name] = aggregation
+        else:
+            print("aggregation not registered")
+        return func
+
+    return wrapper
+
+
+def register_higher_is_better(higher_is_better):
+    def wrapper(func):
+
+        if func.__name__ in func2metric_index:
+            metric_name = func2metric_index[func.__name__]
+            higher_is_better_registry[metric_name] = higher_is_better
+        else:
+            pass
+
+        return func
+
+    return wrapper
+
+
+def register_output_type(output_type):
+    def wrapper(func):
+
+        if func.__name__ in func2metric_index:
+            metric_name = func2metric_index[func.__name__]
+            output_type_registry[metric_name] = output_type
+        else:
+            pass
+
+        return func
+
+    return wrapper