Add revision to hub

TParcollet · TParcollet · commit b78f6f86cfb4 · 2022-07-20T10:06:05.000+02:00
diff --git a/speechbrain/pretrained/fetching.py b/speechbrain/pretrained/fetching.py
@@ -30,6 +30,7 @@ def fetch(
     overwrite=False,
     save_filename=None,
     use_auth_token=False,
+    revision=None,
 ):
     """Ensures you have a local copy of the file, returns its path
 
@@ -65,6 +66,10 @@ def fetch(
     use_auth_token : bool (default: False)
         If true Hugginface's auth_token will be used to load private models from the HuggingFace Hub,
         default is False because majority of models are public.
+    revision : str
+        The model revision corresponding to the HuggingFace Hub model revision.
+        This is particularly useful if you wish to pin your code to a particular
+        version of a model hosted at HuggingFace.
     Returns
     -------
     pathlib.Path
@@ -113,7 +118,10 @@ def fetch(
         logger.info(MSG)
         try:
             fetched_file = huggingface_hub.hf_hub_download(
-                repo_id=source, filename=filename, use_auth_token=use_auth_token
+                repo_id=source,
+                filename=filename,
+                use_auth_token=use_auth_token,
+                revision=revision,
             )
         except HTTPError as e:
             if "404 Client Error" in str(e):
diff --git a/speechbrain/pretrained/interfaces.py b/speechbrain/pretrained/interfaces.py
@@ -289,6 +289,7 @@ def from_hparams(
         overrides={},
         savedir=None,
         use_auth_token=False,
+        revision=None,
         **kwargs,
     ):
         """Fetch and load based from outside source based on HyperPyYAML file
@@ -332,16 +333,20 @@ def from_hparams(
         use_auth_token : bool (default: False)
             If true Hugginface's auth_token will be used to load private models from the HuggingFace Hub,
             default is False because majority of models are public.
+        revision : str
+            The model revision corresponding to the HuggingFace Hub model revision.
+            This is particularly useful if you wish to pin your code to a particular
+            version of a model hosted at HuggingFace.
         """
         if savedir is None:
             clsname = cls.__name__
             savedir = f"./pretrained_models/{clsname}-{hashlib.md5(source.encode('UTF-8', errors='replace')).hexdigest()}"
         hparams_local_path = fetch(
-            hparams_file, source, savedir, use_auth_token
+            hparams_file, source, savedir, use_auth_token, revision
         )
         try:
             pymodule_local_path = fetch(
-                pymodule_file, source, savedir, use_auth_token
+                pymodule_file, source, savedir, use_auth_token, revision
             )
             sys.path.append(str(pymodule_local_path.parent))
         except ValueError:
