diff --git a/magnet/backends/helm/models.py b/magnet/backends/helm/models.py new file mode 100644 index 0000000..b0f64e6 --- /dev/null +++ b/magnet/backends/helm/models.py @@ -0,0 +1,39 @@ +from helm.benchmark.model_deployment_registry import ALL_MODEL_DEPLOYMENTS +from helm.benchmark.config_registry import register_builtin_configs_from_helm_package +from helm.benchmark.model_metadata_registry import ModelMetadata, MODEL_NAME_TO_MODEL_METADATA, UNSUPPORTED_MODEL_TAG, DEPRECATED_MODEL_TAG + +class HELMModels: + """ + Simple collection of ModelDeployments in HELM with corresponding ModelMetadata. + + + Collects all available models, model metadata, and tokenizers from HELM package (helm/config/*.yaml). + + Example: + >>> from magnet.backends.helm.models import HELMModels + >>> models = HELMModels() + >>> len(models) # subject to change + 482 + """ + def __init__(self): + if not ALL_MODEL_DEPLOYMENTS: + # only run once to avoid duplicates + register_builtin_configs_from_helm_package() + + self.models = ALL_MODEL_DEPLOYMENTS + + self.model_metadata = { + model: metadata + for model, metadata in MODEL_NAME_TO_MODEL_METADATA.items() + if not (UNSUPPORTED_MODEL_TAG in metadata.tags or DEPRECATED_MODEL_TAG in metadata.tags) + } + + def get_metadata_from_model_name(self, model_name: str) -> ModelMetadata: + return self.model_metadata[model_name] + + def __len__(self): + """ + Retrieve number of unique supported models that have metadata + """ + # models.model_name includes 2 deprecated models + return len(self.model_metadata) \ No newline at end of file diff --git a/magnet/backends/helm/scenarios.py b/magnet/backends/helm/scenarios.py new file mode 100644 index 0000000..3e4b829 --- /dev/null +++ b/magnet/backends/helm/scenarios.py @@ -0,0 +1,64 @@ +import importlib +from inspect import getmembers, isclass +import pkgutil +from typing import Dict, Optional + +import helm.benchmark.scenarios as HELM_SCENARIOS_PKG +from helm.benchmark.scenarios.scenario import Scenario +from helm.benchmark.run_spec import discover_run_spec_functions, _REGISTERED_RUN_SPEC_FUNCTIONS, RunSpec + +HELM_IMPORTED_SCENARIOS: Dict[str, type[Scenario]] = {} + +class HELMScenarios: + """ + Simple collection of HELM scenarios for aggregate exploration + + + Collects available Scenario and RunSpec definitions in HELM package (helm/benchmark/scenarios). + + Example: + >>> from magnet.backends.helm.scenarios import HELMScenarios + >>> scenarios = HELMScenarios() + >>> len(scenarios) # subject to change; assumes crfm-helm[all] dependencies + 345 + """ + def __init__(self): + if not _REGISTERED_RUN_SPEC_FUNCTIONS: + discover_run_spec_functions() + + self.run_specs = _REGISTERED_RUN_SPEC_FUNCTIONS + + if not HELM_IMPORTED_SCENARIOS: + discover_scenarios() + + self.scenarios = HELM_IMPORTED_SCENARIOS + + def get_run_specs(self, scenario_name: str) -> Optional[RunSpec]: + """ + Return the run spec functions for provided scenario + + WIP: If RunSpecFunction has default args, then func().scenario_spec.class_name is path to Scenario + Only 166/302 values of _REGISTERED_RUN_SPEC_FUNCTIONS have no required args + """ + pass + + def __len__(self): + """ + Retrieve number of implemented scenarios + """ + return len(self.scenarios) + +# Adapted from discover_run_spec_functions() and helpers in helm/benchmark/run_spec.py +def discover_scenarios() -> None: + """ + Discover all scenarios under helm.benchmark.scenarios and store the classes by name + """ + for finder, name, ispkg in pkgutil.walk_packages(HELM_SCENARIOS_PKG.__path__, HELM_SCENARIOS_PKG.__name__ + "."): + try: + module = importlib.import_module(name) + for cls_name, cls in getmembers(module, isclass): + # remove "test_" per docs/scenarios.md + if issubclass(cls, Scenario) and cls is not Scenario and "test_" not in cls_name: + HELM_IMPORTED_SCENARIOS[cls_name] = cls + except ModuleNotFoundError as e: + print(f"Failed to import {name}; Missing {e}") \ No newline at end of file