test: Add script to test model loading below n_parameters threshold

.github/workflows/model_loading.yml

@@ -0,0 +1,24 @@
+name: Model Loading
+
+on:
+  pull_request:
+    paths:
+      - 'mteb/models/**.py'
+
+jobs:
+  extract-and-run:
+    runs-on: ubuntu-latest
+
+    steps:
+    - name: Checkout repository
+      uses: actions/checkout@v3
+
+    - name: Set up Python
+      uses: actions/setup-python@v4
+      with:
+        python-version: '3.10'
+        cache: 'pip'
+
+    - name: Install dependencies and run tests
+      run: |
+        make model-load-test

.gitignore

@@ -145,3 +145,6 @@ tests/create_meta/model_card.md
 # removed results from mteb repo they are now available at: https://github.com/embeddings-benchmark/results
 results/
 uv.lock
+
+# model loading tests
+model_names.txt

Makefile

@@ -35,4 +35,11 @@ pr:
 build-docs:
 	@echo "--- 📚 Building documentation ---"
 	# since we do not have a documentation site, this just build tables for the .md files
-	python docs/create_tasks_table.py
+	python docs/create_tasks_table.py
+
+
+model-load-test:
+	@echo "--- 🚀 Running model load test ---"
+	pip install ".[dev, speedtask, pylate,gritlm,xformers,model2vec]"
+	python scripts/extract_model_names.py
+	python tests/test_models/test_model_loading.py

mteb/models/instruct_wrapper.py

@@ -24,7 +24,7 @@ def instruct_wrapper(
         from gritlm import GritLM
     except ImportError:
         raise ImportError(
-            f"Please install `pip install gritlm` to use {model_name_or_path}."
+            f"Please install `pip install mteb[gritlm]` to use {model_name_or_path}."
         )
 
     class InstructWrapper(GritLM, Wrapper):

pyproject.toml

@@ -65,6 +65,8 @@ openai = ["openai>=1.41.0", "tiktoken>=0.8.0"]
 model2vec = ["model2vec>=0.3.0"]
 pylate = ["pylate>=1.1.4"]
 bm25s = ["bm25s>=0.2.6", "PyStemmer>=2.2.0.3"]
+gritlm = ["gritlm>=1.0.2"]
+xformers = ["xformers>=0.0.29"]
 
 
 [tool.coverage.report]

scripts/extract_model_names.py

@@ -0,0 +1,63 @@
+from __future__ import annotations
+
+import ast
+import sys
+from pathlib import Path
+
+from git import Repo
+
+
+def get_changed_files(base_branch="main"):
+    repo_path = Path(__file__).parent.parent
+    repo = Repo(repo_path)
+    repo.remotes.origin.fetch(base_branch)
+
+    base_commit = repo.commit(f"origin/{base_branch}")
+    head_commit = repo.commit("HEAD")
+
+    diff = repo.git.diff("--name-only", base_commit, head_commit)
+
+    changed_files = diff.splitlines()
+    return [
+        f for f in changed_files if f.startswith("mteb/models/") and f.endswith(".py")
+    ]
+
+
+def extract_model_names(files: list[str]) -> list[str]:
+    model_names = []
+    for file in files:
+        with open(file) as f:
+            tree = ast.parse(f.read())
+            for node in ast.walk(tree):
+                if isinstance(node, ast.Assign):
+                    for target in node.targets:
+                        if (
+                            isinstance(target, ast.Name)
+                            and isinstance(node.value, ast.Call)
+                            and isinstance(node.value.func, ast.Name)
+                            and node.value.func.id == "ModelMeta"
+                        ):
+                            model_name = next(
+                                (
+                                    kw.value.value
+                                    for kw in node.value.keywords
+                                    if kw.arg == "name"
+                                ),
+                                None,
+                            )
+                            if model_name:
+                                model_names.append(model_name)
+    return model_names
+
+
+if __name__ == "__main__":
+    """
+    Can pass in base branch as an argument. Defaults to 'main'.
+    e.g. python extract_model_names.py mieb
+    """
+    base_branch = sys.argv[1] if len(sys.argv) > 1 else "main"
+    changed_files = get_changed_files(base_branch)
+    model_names = extract_model_names(changed_files)
+    output_file = Path(__file__).parent / "model_names.txt"
+    with output_file.open("w") as f:
+        f.write(" ".join(model_names))

scripts/model_loading.py

@@ -0,0 +1,116 @@
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+from pathlib import Path
+
+from huggingface_hub import scan_cache_dir
+
+from mteb import get_model, get_model_meta
+from mteb.models.overview import MODEL_REGISTRY
+
+logging.basicConfig(level=logging.INFO)
+
+
+def teardown_function():
+    hf_cache_info = scan_cache_dir()
+    all_revisions = []
+    for repo in list(hf_cache_info.repos):
+        for revision in list(repo.revisions):
+            all_revisions.append(revision.commit_hash)
+
+    delete_strategy = scan_cache_dir().delete_revisions(*all_revisions)
+    print("Will free " + delete_strategy.expected_freed_size_str)
+    delete_strategy.execute()
+
+
+def get_model_below_n_param_threshold(model_name: str) -> str:
+    """Test that we can get all models with a number of parameters below a threshold."""
+    model_meta = get_model_meta(model_name=model_name)
+    assert model_meta is not None
+    if model_meta.n_parameters is not None:
+        if model_meta.n_parameters >= 2e9:
+            return "Over threshold. Not tested."
+        elif "API" in model_meta.framework:
+            try:
+                m = get_model(model_name)
+                if m is not None:
+                    del m
+                    return "None"
+            except Exception as e:
+                logging.warning(f"Failed to load model {model_name} with error {e}")
+                return e.__str__()
+        try:
+            m = get_model(model_name)
+            if m is not None:
+                del m
+                return "None"
+        except Exception as e:
+            logging.warning(f"Failed to load model {model_name} with error {e}")
+            return e.__str__()
+        finally:
+            teardown_function()
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--omit_previous_success",
+        action="store_true",
+        default=False,
+        help="Omit models that have been successfully loaded in the past",
+    )
+    parser.add_argument(
+        "--run_missing",
+        action="store_true",
+        default=False,
+        help="Run the missing models in the registry that are missing from existing results.",
+    )
+    parser.add_argument(
+        "--model_name",
+        type=str,
+        nargs="+",
+        default=None,
+        help="Run the script for specific model names, e.g. model_1, model_2",
+    )
+
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    output_file = (
+        Path(__file__).parent.parent
+        / "tests"
+        / "test_models"
+        / "model_load_failures.json"
+    )
+
+    args = parse_args()
+
+    # Load existing results if the file exists
+    results = {}
+    if output_file.exists():
+        with output_file.open("r") as f:
+            results = json.load(f)
+
+    if args.model_name:
+        all_model_names = args.model_name
+    else:
+        omit_keys = []
+        if args.run_missing:
+            omit_keys = list(results.keys())
+        elif args.omit_previous_success:
+            omit_keys = [k for k, v in results.items() if v == "None"]
+
+        all_model_names = list(set(MODEL_REGISTRY.keys()) - set(omit_keys))
+
+    for model_name in all_model_names:
+        error_msg = get_model_below_n_param_threshold(model_name)
+        results[model_name] = error_msg
+
+        results = dict(sorted(results.items()))
+
+        # Write the results to the file after each iteration
+        with output_file.open("w") as f:
+            json.dump(results, f, indent=4)