feat: run any skops model from hf hub

Author: martibosch

detectree/classifier.py

@@ -264,6 +264,10 @@ def __init__(
         *,
         clf=None,
         clf_dict=None,
+        hf_hub_repo_id=None,
+        hf_hub_clf_filename=None,
+        hf_hub_download_kwargs=None,
+        skops_trusted=None,
         tree_val=None,
         nontree_val=None,
         refine=None,
@@ -285,6 +289,20 @@ def __init__(
         clf_dict : dictionary, optional
             Dictionary mapping a trained scikit-learn-like classifier to each
             first-level cluster label.
+        hf_hub_repo_id, hf_hub_clf_filename : str, optional
+            HuggingFace Hub repository id (string with the user or organization and
+            repository name separated by a `/`) and file name of the skops classifier
+            respectively. If no value is provided, the values set in
+            `settings.HF_HUB_REPO_ID` and `settings.HF_HUB_CLF_FILENAME` Ignored if
+            `clf` or `clf_dict` are provided.
+        hf_hub_download_kwargs : dict, optional
+            Additional keyword arguments (besides "repo_id", "filename", "library_name"
+            and "library_version") to pass to `huggingface_hub.hf_hub_download`.
+        skosp_trusted : list, optional
+            List of trusted object types to load the classifier from HuggingFace Hub,
+            passed to `skops.io.load`. If no value is provided, the value from
+            `settings.SKOPS_TRUSTED` is used. Ignored if `clf` or `clf_dict` are
+            provided.
         tree_val : int, optional
             Label used to denote tree pixels in the predicted images. If no value is
             provided, the value set in `settings.CLF_TREE_VAL` is used.
@@ -315,14 +333,32 @@ def __init__(
         elif clf is not None:
             self.clf = clf
         else:
+            if hf_hub_repo_id is None:
+                hf_hub_repo_id = settings.HF_HUB_REPO_ID
+            if hf_hub_clf_filename is None:
+                hf_hub_clf_filename = settings.HF_HUB_CLF_FILENAME
+            if hf_hub_download_kwargs is None:
+                _hf_hub_download_kwargs = {}
+            else:
+                _hf_hub_download_kwargs = hf_hub_download_kwargs.copy()
+                for key in [
+                    "repo_id",
+                    "filename",
+                    "library_name",
+                    "library_version",
+                ]:
+                    _ = _hf_hub_download_kwargs.pop(key, None)
+            if skops_trusted is None:
+                skops_trusted = settings.SKOPS_TRUSTED
             self.clf = io.load(
                 hf_hub.hf_hub_download(
-                    repo_id=settings.HF_HUB_REPO_ID,
-                    filename=settings.HF_HUB_FILENAME,
+                    repo_id=hf_hub_repo_id,
+                    filename=hf_hub_clf_filename,
                     library_name="skops",
                     library_version=skops.__version__,
+                    **_hf_hub_download_kwargs,
                 ),
-                trusted=settings.SKOPS_TRUSTED,
+                trusted=skops_trusted,
             )
 
         if tree_val is None:

detectree/settings.py

@@ -36,7 +36,7 @@
     "lightgbm.sklearn.LGBMClassifier",
 ]
 HF_HUB_REPO_ID = "martibosch/detectree"
-HF_HUB_FILENAME = "clf.skops"
+HF_HUB_CLF_FILENAME = "clf.skops"
 
 # LIDAR
 LIDAR_TREE_THRESHOLD = 15

tests/test_detectree.py

@@ -610,7 +610,14 @@ def test_classifier(self):
         # test that for the pre-trained classifier (no init `clf`/`clf_dict` arg) and
         # for the classifier initialized with the `clf` arg, the `clf` attribute is set
         # but not `clf_dict`
-        for c in [dtr.Classifier(), dtr.Classifier(clf=self.clf)]:
+        for c in [
+            dtr.Classifier(),
+            dtr.Classifier(clf=self.clf),
+            dtr.Classifier(
+                hf_hub_repo_id=settings.HF_HUB_REPO_ID,
+                hf_hub_clf_filename=settings.HF_CLF_FILENAME,
+            ),
+        ]:
             self.assertTrue(hasattr(c, "clf"))
             self.assertFalse(hasattr(c, "clf_dict"))
         # test that when initializing `clf_dict`, the `clf_dict` attribute is set but