Repository: huggingface/huggingface_hub
Pull request #73 (2021-06-04 06:45:58)

speechbrain library tag + code snippet

osanseviero: I did some changes for consistency with #94.

@TParcollet and team, do the code snippets seem appropriate for the respective models?
TParcollet: Yes! 

Patch:
diff --git a/interfaces/Libraries.ts b/interfaces/Libraries.ts
index bae1371b..28c234fe 100644
--- a/interfaces/Libraries.ts
+++ b/interfaces/Libraries.ts
@@ -12,6 +12,7 @@ export enum ModelLibrary {
 	'pyannote'               = 'Pyannote',
 	'sentence-transformers'  = 'Sentence Transformers',
 	'spacy'                  = 'spaCy',
+	'speechbrain'            = 'speechbrain',
 	'tensorflowtts'          = 'TensorFlowTTS',
 	'timm'                   = 'Timm',
 	'transformers'           = 'Transformers',
@@ -176,6 +177,37 @@ nlp = spacy.load("${nameWithoutNamespace(model.modelId)}")
 import ${nameWithoutNamespace(model.modelId)}
 nlp = ${nameWithoutNamespace(model.modelId)}.load()`;
 
+const speechbrainASR = (model: ModelData) =>
+`from speechbrain.pretrained import EncoderDecoderASR
+
+asr_model = EncoderDecoderASR.from_hparams(source="${model.modelId}")
+asr_model.transcribe_file("file.wav")`;
+
+const speechbrainEnhancement = (model: ModelData) =>
+`from speechbrain.pretrained import SpectralMaskEnhancement
+
+enhance_model = SpectralMaskEnhancement.from_hparams(source="${model.modelId}")
+enhanced = enhance_model.enhance_file("file.wav")`;
+
+const speechbrainSeparator = (model: ModelData) =>
+`from speechbrain.pretrained import SepformerSeparation
+
+separator_model = SepformerSeparation.from_hparams(source="${model.modelId}")
+est_sources = separator_model.separate_file("file.wav")`;
+
+const speechbrain = (model: ModelData) => {
+	if (model.tags?.includes("automatic-speech-recognition")){
+		return speechbrainASR(model);
+	} else if (model.tags?.includes("audio-to-audio")) {
+		if (model.tags?.includes("speech-enhancement")) {
+			return speechbrainEnhancement(model);
+		} else if (model.tags?.includes("audio-source-separation")) {
+			return speechbrainSeparator(model);
+		} 
+	}
+	return "# Unable to determine model type";
+};
+
 const transformers = (model: ModelData) =>
 `from transformers import AutoTokenizer, ${model.autoArchitecture}
   
@@ -236,6 +268,12 @@ export const MODEL_LIBRARIES_UI_ELEMENTS: { [key in keyof typeof ModelLibrary]:
 		repoUrl: "https://github.com/explosion/spaCy",
 		snippet: spacy,
 	},
+	speechbrain: {
+		btnLabel: "speechbrain",
+		repoName: "speechbrain",
+		repoUrl: "https://github.com/speechbrain/speechbrain",
+		snippet: speechbrain,
+	},
 	tensorflowtts : {
 		btnLabel: "TensorFlowTTS",
 		repoName: "TensorFlowTTS",
diff --git a/src/huggingface_hub/__init__.py b/src/huggingface_hub/__init__.py
index 29a628fa..2ad5972a 100644
--- a/src/huggingface_hub/__init__.py
+++ b/src/huggingface_hub/__init__.py
@@ -25,6 +25,7 @@ from .constants import (
     HUGGINGFACE_CO_URL_TEMPLATE,
     PYTORCH_WEIGHTS_NAME,
     REPO_TYPE_DATASET,
+    REPO_TYPE_SPACE,
     TF2_WEIGHTS_NAME,
     TF_WEIGHTS_NAME,
 )
diff --git a/src/huggingface_hub/commands/user.py b/src/huggingface_hub/commands/user.py
index 4153f3f5..1918b33e 100644
--- a/src/huggingface_hub/commands/user.py
+++ b/src/huggingface_hub/commands/user.py
@@ -18,11 +18,7 @@ from getpass import getpass
 from typing import List, Union
 
 from huggingface_hub.commands import BaseHuggingfaceCLICommand
-from huggingface_hub.constants import (
-    REPO_TYPE_DATASET,
-    REPO_TYPE_DATASET_URL_PREFIX,
-    REPO_TYPES,
-)
+from huggingface_hub.constants import REPO_TYPES, REPO_TYPES_URL_PREFIXES
 from huggingface_hub.hf_api import HfApi, HfFolder
 from requests.exceptions import HTTPError
 
@@ -67,7 +63,7 @@ class UserCommands(BaseHuggingfaceCLICommand):
         repo_create_parser.add_argument(
             "--type",
             type=str,
-            help='Optional: repo_type: set to "dataset" if creating a dataset, default is model.',
+            help='Optional: repo_type: set to "dataset" or "space" if creating a dataset or space, default is model.',
         )
         repo_create_parser.add_argument(
             "--organization", type=str, help="Optional: organization namespace."
@@ -239,8 +235,8 @@ class RepoCreateCommand(BaseUserCommand):
             print("Invalid repo --type")
             exit(1)
 
-        if self.args.type == REPO_TYPE_DATASET:
-            repo_id = REPO_TYPE_DATASET_URL_PREFIX + repo_id
+        if self.args.type in REPO_TYPES_URL_PREFIXES:
+            repo_id = REPO_TYPES_URL_PREFIXES[self.args.type] + repo_id
 
         print("You are about to create {}".format(ANSI.bold(repo_id)))
 
diff --git a/src/huggingface_hub/constants.py b/src/huggingface_hub/constants.py
index b1af0766..8155e23e 100644
--- a/src/huggingface_hub/constants.py
+++ b/src/huggingface_hub/constants.py
@@ -21,9 +21,13 @@ HUGGINGFACE_CO_URL_TEMPLATE = (
 )
 
 REPO_TYPE_DATASET = "dataset"
-REPO_TYPES = [None, REPO_TYPE_DATASET]
+REPO_TYPE_SPACE = "space"
+REPO_TYPES = [None, REPO_TYPE_DATASET, REPO_TYPE_SPACE]
 
-REPO_TYPE_DATASET_URL_PREFIX = "datasets/"
+REPO_TYPES_URL_PREFIXES = {
+    REPO_TYPE_DATASET: "datasets/",
+    REPO_TYPE_SPACE: "spaces/",
+}
 
 
 # default cache
diff --git a/src/huggingface_hub/file_download.py b/src/huggingface_hub/file_download.py
index d7b4994b..45c5ac00 100644
--- a/src/huggingface_hub/file_download.py
+++ b/src/huggingface_hub/file_download.py
@@ -23,9 +23,8 @@ from . import __version__
 from .constants import (
     HUGGINGFACE_CO_URL_TEMPLATE,
     HUGGINGFACE_HUB_CACHE,
-    REPO_TYPE_DATASET,
-    REPO_TYPE_DATASET_URL_PREFIX,
     REPO_TYPES,
+    REPO_TYPES_URL_PREFIXES,
 )
 from .hf_api import HfFolder
 
@@ -92,8 +91,8 @@ def hf_hub_url(
     if repo_type not in REPO_TYPES:
         raise ValueError("Invalid repo type")
 
-    if repo_type == REPO_TYPE_DATASET:
-        repo_id = REPO_TYPE_DATASET_URL_PREFIX + repo_id
+    if repo_type in REPO_TYPES_URL_PREFIXES:
+        repo_id = REPO_TYPES_URL_PREFIXES[repo_type] + repo_id
 
     if revision is None:
         revision = "main"
diff --git a/src/huggingface_hub/hf_api.py b/src/huggingface_hub/hf_api.py
index 385728d8..2d8d8ddf 100644
--- a/src/huggingface_hub/hf_api.py
+++ b/src/huggingface_hub/hf_api.py
@@ -22,17 +22,17 @@ from io import BufferedIOBase, RawIOBase
 from os.path import expanduser
 from typing import BinaryIO, Dict, Iterable, List, Optional, Tuple, Union
 
+import requests
+from requests.exceptions import HTTPError
+
+from .constants import REPO_TYPES, REPO_TYPES_URL_PREFIXES
+
 
 if sys.version_info >= (3, 8):
     from typing import Literal
 else:
     from typing_extensions import Literal
 
-import requests
-from requests.exceptions import HTTPError
-
-from .constants import REPO_TYPE_DATASET, REPO_TYPE_DATASET_URL_PREFIX, REPO_TYPES
-
 
 ENDPOINT = "https://huggingface.co"
 REMOTE_FILEPATH_REGEX = re.compile(r"^\w[\w\/]*(\.\w+)?$")
@@ -251,7 +251,7 @@ class HfApi:
         self, token: str, organization: Optional[str] = None
     ) -> List[RepoObj]:
         """
-        HuggingFace git-based system, used for models.
+        HuggingFace git-based system, used for models, datasets, and spaces.
 
         Call HF API to list all stored files for user (or one of their organizations).
         """
@@ -275,14 +275,14 @@ class HfApi:
         lfsmultipartthresh: Optional[int] = None,
     ) -> str:
         """
-        HuggingFace git-based system, used for models.
+        HuggingFace git-based system, used for models, datasets, and spaces.
 
         Call HF API to create a whole repo.
 
         Params:
             private: Whether the model repo should be private (requires a paid huggingface.co account)
 
-            repo_type: Set to "dataset" if creating a dataset, default is model
+            repo_type: Set to "dataset" or "space" if creating a dataset or space, default is model
 
             exist_ok: Do not raise an error if repo already exists
 
@@ -332,7 +332,7 @@ class HfApi:
         repo_type: Optional[str] = None,
     ):
         """
-        HuggingFace git-based system, used for models.
+        HuggingFace git-based system, used for models, datasets, and spaces.
 
         Call HF API to delete a whole repo.
 
@@ -374,13 +374,11 @@ class HfApi:
             namespace = organization
 
         path_prefix = "{}/api/".format(self.endpoint)
-        if repo_type == REPO_TYPE_DATASET:
-            path_prefix += REPO_TYPE_DATASET_URL_PREFIX
+        if repo_type in REPO_TYPES_URL_PREFIXES:
+            path_prefix += REPO_TYPES_URL_PREFIXES[repo_type]
 
         path = "{}{}/{}/settings".format(path_prefix, namespace, name)
         json = {"private": private}
-        if repo_type is not None:
-            json["type"] = repo_type
 
         r = requests.put(
             path,
@@ -417,7 +415,7 @@ class HfApi:
                 The repository to which the file will be uploaded, for example: :obj:`"username/custom_transformers"`
 
             repo_type (``str``, Optional):
-                Set to :obj:`"dataset"` if uploading to a dataset, :obj:`None` if uploading to a model. Default is :obj:`None`.
+                Set to :obj:`"dataset"` or :obj:`"space"` if uploading to a dataset or space, :obj:`None` if uploading to a model. Default is :obj:`None`.
 
             revision (``str``, Optional):
                 The git revision to commit from. Defaults to the :obj:`"main"` branch.
@@ -480,12 +478,8 @@ class HfApi:
                 )
             )
 
-        repo_id = "{prefix}{repo_id}".format(
-            prefix=REPO_TYPE_DATASET_URL_PREFIX
-            if (repo_type == REPO_TYPE_DATASET)
-            else "",
-            repo_id=repo_id,
-        )
+        if repo_type in REPO_TYPES_URL_PREFIXES:
+            repo_id = REPO_TYPES_URL_PREFIXES[repo_type] + repo_id
 
         revision = revision if revision is not None else "main"
 
diff --git a/src/huggingface_hub/hub_mixin.py b/src/huggingface_hub/hub_mixin.py
index a38473b1..bf28c09c 100644
--- a/src/huggingface_hub/hub_mixin.py
+++ b/src/huggingface_hub/hub_mixin.py
@@ -18,7 +18,7 @@ if is_torch_available():
 logger = logging.getLogger(__name__)
 
 
-class ModelHubMixin(object):
+class ModelHubMixin:
     def __init__(self, *args, **kwargs):
         """
         Mix this class with your torch-model class for ease process of saving & loading from huggingface-hub
diff --git a/src/huggingface_hub/repository.py b/src/huggingface_hub/repository.py
index e1607d36..5b1f4b22 100644
--- a/src/huggingface_hub/repository.py
+++ b/src/huggingface_hub/repository.py
@@ -206,9 +206,8 @@ class Repository:
                 )
 
                 if "main" not in output.stdout.split():
-                    # TODO(check if we really want the --force flag)
                     subprocess.run(
-                        "git checkout origin/main -ft".split(),
+                        "git checkout origin/main -t".split(),
                         stderr=subprocess.PIPE,
                         stdout=subprocess.PIPE,
                         encoding="utf-8",
@@ -415,7 +414,7 @@ class Repository:
 
     def push_to_hub(self, commit_message="commit files to HF hub") -> str:
         """
-        Helper to add, commit, and pushe file to remote repository on the HuggingFace Hub.
+        Helper to add, commit, and push files to remote repository on the HuggingFace Hub.
         Args:
             commit_message: commit message.
         """


Test patch:
diff --git a/tests/test_hf_api.py b/tests/test_hf_api.py
index 46e0fc49..0aa81d1b 100644
--- a/tests/test_hf_api.py
+++ b/tests/test_hf_api.py
@@ -21,7 +21,7 @@ import time
 import unittest
 from io import BytesIO
 
-from huggingface_hub.constants import REPO_TYPE_DATASET
+from huggingface_hub.constants import REPO_TYPE_DATASET, REPO_TYPE_SPACE
 from huggingface_hub.file_download import cached_download
 from huggingface_hub.hf_api import HfApi, HfFolder, ModelInfo, RepoObj
 from requests.exceptions import HTTPError
@@ -38,6 +38,7 @@ from .testing_utils import (
 REPO_NAME = "my-model-{}".format(int(time.time() * 10e3))
 REPO_NAME_LARGE_FILE = "my-model-largefiles-{}".format(int(time.time() * 10e3))
 DATASET_REPO_NAME = "my-dataset-{}".format(int(time.time() * 10e3))
+SPACE_REPO_NAME = "my-space-{}".format(int(time.time() * 10e3))
 WORKING_REPO_DIR = os.path.join(
     os.path.dirname(os.path.abspath(__file__)), "fixtures/working_repo"
 )
@@ -59,7 +60,7 @@ class HfApiLoginTest(HfApiCommonTest):
         self.assertIsInstance(token, str)
 
 
-class HfApiEndpointsTest(HfApiCommonTest):
+class HfApiCommonTestWithLogin(HfApiCommonTest):
     @classmethod
     def setUpClass(cls):
         """
@@ -67,6 +68,8 @@ class HfApiEndpointsTest(HfApiCommonTest):
         """
         cls._token = cls._api.login(username=USER, password=PASS)
 
+
+class HfApiEndpointsTest(HfApiCommonTestWithLogin):
     def test_whoami(self):
         user, orgs = self._api.whoami(token=self._token)
         self.assertEqual(user, USER)
@@ -93,25 +96,51 @@ class HfApiEndpointsTest(HfApiCommonTest):
 
     def test_create_update_and_delete_dataset_repo(self):
         self._api.create_repo(
-            token=self._token, name=REPO_NAME, repo_type=REPO_TYPE_DATASET
+            token=self._token, name=DATASET_REPO_NAME, repo_type=REPO_TYPE_DATASET
         )
         res = self._api.update_repo_visibility(
-            token=self._token, name=REPO_NAME, private=True, repo_type=REPO_TYPE_DATASET
+            token=self._token,
+            name=DATASET_REPO_NAME,
+            private=True,
+            repo_type=REPO_TYPE_DATASET,
         )
         self.assertTrue(res["private"])
         res = self._api.update_repo_visibility(
             token=self._token,
-            name=REPO_NAME,
+            name=DATASET_REPO_NAME,
             private=False,
             repo_type=REPO_TYPE_DATASET,
         )
         self.assertFalse(res["private"])
         self._api.delete_repo(
-            token=self._token, name=REPO_NAME, repo_type=REPO_TYPE_DATASET
+            token=self._token, name=DATASET_REPO_NAME, repo_type=REPO_TYPE_DATASET
+        )
+
+    @unittest.skip("skipped while spaces in beta")
+    def test_create_update_and_delete_space_repo(self):
+        self._api.create_repo(
+            token=self._token, name=SPACE_REPO_NAME, repo_type=REPO_TYPE_SPACE
+        )
+        res = self._api.update_repo_visibility(
+            token=self._token,
+            name=SPACE_REPO_NAME,
+            private=True,
+            repo_type=REPO_TYPE_SPACE,
+        )
+        self.assertTrue(res["private"])
+        res = self._api.update_repo_visibility(
+            token=self._token,
+            name=SPACE_REPO_NAME,
+            private=False,
+            repo_type=REPO_TYPE_SPACE,
+        )
+        self.assertFalse(res["private"])
+        self._api.delete_repo(
+            token=self._token, name=SPACE_REPO_NAME, repo_type=REPO_TYPE_SPACE
         )
 
 
-class HfApiUploadFileTest(HfApiEndpointsTest):
+class HfApiUploadFileTest(HfApiCommonTestWithLogin):
     def setUp(self) -> None:
         super().setUp()
         self.tmp_dir = tempfile.mkdtemp()