More verbose huggingface errors

doc_ufcn/models.py

@@ -5,6 +5,7 @@ import os
 
 import yaml
 from huggingface_hub import hf_hub_download
+from huggingface_hub.utils import RepositoryNotFoundError, RevisionNotFoundError
 
 HUGGING_FACE_REPO_PREFIX = "Teklia/doc-ufcn-"
 
@@ -19,20 +20,35 @@ def download_model(name, version=None):
 
     cache_dir = os.environ.get("XDG_CACHE_HOME", os.path.expanduser("~/.cache"))
     dir_path = os.path.join(cache_dir, "doc-ufcn", "models", name)
-    # Retrieve parameters.yml
-    parameters_path = hf_hub_download(
-        repo_id=HUGGING_FACE_REPO_PREFIX + name,
-        filename="parameters.yml",
-        cache_dir=dir_path,
-        revision=version,
-    )
-    # Retrieve parameters.yml
-    model_path = hf_hub_download(
-        repo_id=HUGGING_FACE_REPO_PREFIX + name,
-        filename="model.pth",
-        cache_dir=dir_path,
-        revision=version,
-    )
+
+    try:
+        # Retrieve parameters.yml
+        parameters_path = hf_hub_download(
+            repo_id=HUGGING_FACE_REPO_PREFIX + name,
+            filename="parameters.yml",
+            cache_dir=dir_path,
+            revision=version,
+        )
+        # Retrieve model.pth
+        model_path = hf_hub_download(
+            repo_id=HUGGING_FACE_REPO_PREFIX + name,
+            filename="model.pth",
+            cache_dir=dir_path,
+            revision=version,
+        )
+    except RepositoryNotFoundError as e:
+        logger.error(
+            f"Repository with name {name} was not found or you may not have access to it."
+        )
+        print(str(e))
+        raise
+    except RevisionNotFoundError as e:
+        logger.error(
+            f"Revision {version} was not found on the repository with name {name}."
+        )
+        print(str(e))
+        raise
+
     with open(parameters_path) as f:
         parameters = yaml.safe_load(f)
 

tests/test_models.py

@@ -3,12 +3,13 @@
 import os
 
 import pytest
+from huggingface_hub.utils import RepositoryNotFoundError, RevisionNotFoundError
 
 from doc_ufcn import models
 
 
 @pytest.mark.parametrize(
-    "name, version, expected_model_path, expected_parameters",
+    "name, version, expected_model_path, expected_parameters, correct_name",
     [
         # Correct name and version
         (
@@ -16,6 +17,7 @@ from doc_ufcn import models
             "main",
             "~/.cache/doc-ufcn/models/generic-page/models--Teklia--doc-ufcn-generic-page/snapshots/ef5519d77baa190bfe294ad3994146dbe7e72095",
             pytest.lazy_fixture("test_parameters"),
+            True,
         ),
         # Correct fullname and version
         (
@@ -23,30 +25,40 @@ from doc_ufcn import models
             "main",
             "~/.cache/doc-ufcn/models/generic-page/models--Teklia--doc-ufcn-generic-page/snapshots/ef5519d77baa190bfe294ad3994146dbe7e72095",
             pytest.lazy_fixture("test_parameters"),
+            True,
         ),
         # Correct name and incorrect version
-        ("generic-page", "version", None, None),
+        ("generic-page", "version", None, None, True),
         # Correct name and no version
         (
             "generic-page",
             None,
             "~/.cache/doc-ufcn/models/generic-page/models--Teklia--doc-ufcn-generic-page/snapshots/ef5519d77baa190bfe294ad3994146dbe7e72095",
             pytest.lazy_fixture("test_parameters"),
+            True,
         ),
         # Incorrect name and incorrect version
-        ("page_model", "version", None, None),
+        ("page_model", "version", None, None, False),
         # Incorrect name and no version
-        ("page_model", None, None, None),
+        ("page_model", None, None, None, False),
     ],
 )
-def test_download_model(name, version, expected_model_path, expected_parameters):
+def test_download_model(
+    name, version, expected_model_path, expected_parameters, correct_name
+):
     """
     Test of the download_model function.
     Check that the correct model is loaded.
     """
     if expected_model_path is None and expected_parameters is None:
-        with pytest.raises(Exception):
-            model_path, parameters = models.download_model(name, version)
+        if correct_name:
+            # Bad model name
+            with pytest.raises(RevisionNotFoundError):
+                model_path, parameters = models.download_model(name, version)
+        else:
+            # Bad revision
+            with pytest.raises(RepositoryNotFoundError):
+                model_path, parameters = models.download_model(name, version)
     else:
         model_path, parameters = models.download_model(name, version)
         assert model_path == os.path.join(