Merge pull request #126 from lorenzomammana/feature/external-package-…

…integration

External package integration using plugins

lmms_eval/__main__.py

@@ -1,3 +1,4 @@
+import importlib
 import os
 import yaml
 import sys
@@ -236,6 +237,12 @@ def cli_evaluate_single(args: Union[argparse.Namespace, None] = None) -> None:
         eval_logger.info(f"Including path: {args.include_path}")
         include_path(args.include_path)
 
+    if os.environ.get("LMMS_EVAL_PLUGINS", None):
+        for plugin in os.environ["LMMS_EVAL_PLUGINS"].split(","):
+            package_tasks_location = importlib.util.find_spec(f"{plugin}.tasks").submodule_search_locations[0]
+            eval_logger.info(f"Including path: {args.include_path}")
+            include_path(package_tasks_location)
+
     if args.tasks is None:
         task_names = ALL_TASKS
     elif args.tasks == "list":

lmms_eval/models/__init__.py

@@ -1,3 +1,6 @@
+import importlib
+import os
+import hf_transfer
 from loguru import logger
 import sys
 
@@ -38,3 +41,17 @@
     except ImportError as e:
         # logger.warning(f"Failed to import {model_class} from {model_name}: {e}")
         pass
+
+if os.environ.get("LMMS_EVAL_PLUGINS", None):
+    # Allow specifying other packages to import models from
+    for plugin in os.environ["LMMS_EVAL_PLUGINS"].split(","):
+        m = importlib.import_module(f"{plugin}.models")
+        for model_name, model_class in getattr(m, "AVAILABLE_MODELS").items():
+            try:
+                exec(f"from {plugin}.models.{model_name} import {model_class}")
+            except ImportError:
+                pass
+
+import hf_transfer
+
+os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"