Repository: huggingface/huggingface_hub
Pull request #11 (2021-01-19 23:58:46)

Adding mixin class for ease saving, uploading, downloading (as discussed in issue #9).
This PR will add feature as discussed in issue #9.

## Example
```python
from huggingface_hub import ModelHubMixin

class MyModel(nn.Module, ModelHubMixin):
   def __init__(self, **kwargs):
      super().__init__()
      self.config = kwargs.pop("config", None)
      self.layer = ...
   def forward(self, ...):
      return ...

model = MyModel()

# saving model to local directory & pushing to hub
model.save_pretrained("mymodel", upload_to_hub=True, config={"act": "gelu"})

# initiatizing model & loading it from trained-weights
model = MyModel.from_pretrained("username/mymodel@main")
```

## Note
 - In case user want to save some specific layers as per his needs, he can simple overwrite `._save_pretrained()` method.
 - Supporting all kind of `PyTorch` model.
 - `tf2` not supported currently


vasudevgupta7: Shall I integrate model uploading directly in `.save_pretrained(path, model_id="username/mymodel", upload_to_hub=True)` method instead of making separate method `.upload_to_hub(model_id="username/mymodel")`? I feel this looks better.

Please put your views on this.
julien-c: Hi @vasudevgupta7 thanks for the PR! 

I have no more bandwidth this week but plan on actively looking at this early next week!
julien-c: Ooops @vasudevgupta7 I wanted to reach out to let you know that we had merged the Repository PR. But I guess you saw that :)

Will take a look at your updated PR with @philschmid and others, soon. Thanks!
julien-c: > Shall I integrate model uploading directly in .save_pretrained(path, model_id="username/mymodel", upload_to_hub=True) method instead of making separate method .upload_to_hub(model_id="username/mymodel")? I feel this looks better.

Replying late to this but yeah I feel like it'd be nice to expose both.

julien-c: @vasudevgupta7 let us know when you feel like this is ready to squash/merge and I'll make a final review. Thanks again
vasudevgupta7: Hoping everything works now. Sorry for the delay.
julien-c: Looks good! 

Does anyone tagged here have any last comments? @LysandreJik and @philschmid maybe?

Otherwise I'll merge and push a release soon.

Patch:
diff --git a/src/huggingface_hub/__init__.py b/src/huggingface_hub/__init__.py
index f02353ed..04fe01a5 100644
--- a/src/huggingface_hub/__init__.py
+++ b/src/huggingface_hub/__init__.py
@@ -30,4 +30,5 @@ from .constants import (
 )
 from .file_download import cached_download, hf_hub_url
 from .hf_api import HfApi, HfFolder
+from .hub_mixin import ModelHubMixin
 from .repository import Repository
diff --git a/src/huggingface_hub/constants.py b/src/huggingface_hub/constants.py
index fb64c498..b1af0766 100644
--- a/src/huggingface_hub/constants.py
+++ b/src/huggingface_hub/constants.py
@@ -1,6 +1,11 @@
 import os
 
 
+# Possible values for env variables
+
+ENV_VARS_TRUE_VALUES = {"1", "ON", "YES", "TRUE"}
+ENV_VARS_TRUE_AND_AUTO_VALUES = ENV_VARS_TRUE_VALUES.union({"AUTO"})
+
 # Constants for file downloads
 
 PYTORCH_WEIGHTS_NAME = "pytorch_model.bin"
@@ -30,3 +35,9 @@ hf_cache_home = os.path.expanduser(
 default_cache_path = os.path.join(hf_cache_home, "hub")
 
 HUGGINGFACE_HUB_CACHE = os.getenv("HUGGINGFACE_HUB_CACHE", default_cache_path)
+
+HF_HUB_OFFLINE = os.environ.get("HF_HUB_OFFLINE", "AUTO").upper()
+if HF_HUB_OFFLINE in ENV_VARS_TRUE_VALUES:
+    HF_HUB_OFFLINE = True
+else:
+    HF_HUB_OFFLINE = False
diff --git a/src/huggingface_hub/file_download.py b/src/huggingface_hub/file_download.py
index 31d1602e..026ef3bb 100644
--- a/src/huggingface_hub/file_download.py
+++ b/src/huggingface_hub/file_download.py
@@ -6,6 +6,7 @@ import logging
 import os
 import sys
 import tempfile
+import time
 from contextlib import contextmanager
 from functools import partial
 from hashlib import sha256
@@ -16,6 +17,7 @@ from tqdm.auto import tqdm
 
 import requests
 from filelock import FileLock
+from huggingface_hub import constants
 
 from . import __version__
 from .constants import (
@@ -171,12 +173,73 @@ def http_user_agent(
     return ua
 
 
+class OfflineModeIsEnabled(ConnectionError):
+    pass
+
+
+def _raise_if_offline_mode_is_enabled(msg: Optional[str] = None):
+    """Raise a OfflineModeIsEnabled error (subclass of ConnectionError) if HF_HUB_OFFLINE is True."""
+    if constants.HF_HUB_OFFLINE:
+        raise OfflineModeIsEnabled(
+            "Offline mode is enabled."
+            if msg is None
+            else "Offline mode is enabled. " + str(msg)
+        )
+
+
+def _request_with_retry(
+    method: str,
+    url: str,
+    max_retries: int = 0,
+    base_wait_time: float = 0.5,
+    max_wait_time: float = 2,
+    timeout: float = 10.0,
+    **params,
+) -> requests.Response:
+    """Wrapper around requests to retry in case it fails with a ConnectTimeout, with exponential backoff.
+
+    Note that if the environment variable HF_HUB_OFFLINE is set to 1, then a OfflineModeIsEnabled error is raised.
+
+    Args:
+        method (str): HTTP method, such as 'GET' or 'HEAD'
+        url (str): The URL of the ressource to fetch
+        max_retries (int): Maximum number of retries, defaults to 0 (no retries)
+        base_wait_time (float): Duration (in seconds) to wait before retrying the first time. Wait time between
+            retries then grows exponentially, capped by max_wait_time.
+        max_wait_time (float): Maximum amount of time between two retries, in seconds
+        **params: Params to pass to `requests.request`
+    """
+    _raise_if_offline_mode_is_enabled(f"Tried to reach {url}")
+    tries, success = 0, False
+    while not success:
+        tries += 1
+        try:
+            response = requests.request(
+                method=method.upper(), url=url, timeout=timeout, **params
+            )
+            success = True
+        except requests.exceptions.ConnectTimeout as err:
+            if tries > max_retries:
+                raise err
+            else:
+                logger.info(
+                    f"{method} request to {url} timed out, retrying... [{tries/max_retries}]"
+                )
+                sleep_time = max(
+                    max_wait_time, base_wait_time * 2 ** (tries - 1)
+                )  # Exponential backoff
+                time.sleep(sleep_time)
+    return response
+
+
 def http_get(
     url: str,
     temp_file: BinaryIO,
     proxies=None,
     resume_size=0,
     headers: Optional[Dict[str, str]] = None,
+    timeout=10.0,
+    max_retries=0,
 ):
     """
     Donwload remote file. Do not gobble up errors.
@@ -184,7 +247,15 @@ def http_get(
     headers = copy.deepcopy(headers)
     if resume_size > 0:
         headers["Range"] = "bytes=%d-" % (resume_size,)
-    r = requests.get(url, stream=True, proxies=proxies, headers=headers)
+    r = _request_with_retry(
+        method="GET",
+        url=url,
+        stream=True,
+        proxies=proxies,
+        headers=headers,
+        timeout=timeout,
+        max_retries=max_retries,
+    )
     r.raise_for_status()
     content_length = r.headers.get("Content-Length")
     total = resume_size + int(content_length) if content_length is not None else None
@@ -254,8 +325,9 @@ def cached_download(
     etag = None
     if not local_files_only:
         try:
-            r = requests.head(
-                url,
+            r = _request_with_retry(
+                method="HEAD",
+                url=url,
                 headers=headers,
                 allow_redirects=False,
                 proxies=proxies,
@@ -276,15 +348,14 @@ def cached_download(
             # between the HEAD and the GET (unlikely, but hey).
             if 300 <= r.status_code <= 399:
                 url_to_download = r.headers["Location"]
+        except (requests.exceptions.SSLError, requests.exceptions.ProxyError):
+            # Actually raise for those subclasses of ConnectionError
+            raise
         except (
             requests.exceptions.ConnectionError,
             requests.exceptions.Timeout,
-        ) as exc:
-            # Actually raise for those subclasses of ConnectionError:
-            if isinstance(exc, requests.exceptions.SSLError) or isinstance(
-                exc, requests.exceptions.ProxyError
-            ):
-                raise exc
+            OfflineModeIsEnabled,
+        ):
             # Otherwise, our Internet connection is down.
             # etag is None
             pass
@@ -297,7 +368,7 @@ def cached_download(
     # etag is None == we don't have a connection or we passed local_files_only.
     # try to get the last downloaded one
     if etag is None:
-        if os.path.exists(cache_path):
+        if os.path.exists(cache_path) and not force_download:
             return cache_path
         else:
             matching_files = [
@@ -307,7 +378,7 @@ def cached_download(
                 )
                 if not file.endswith(".json") and not file.endswith(".lock")
             ]
-            if len(matching_files) > 0:
+            if len(matching_files) > 0 and not force_download:
                 return os.path.join(cache_dir, matching_files[-1])
             else:
                 # If files cannot be found and local_files_only=True,
diff --git a/src/huggingface_hub/hub_mixin.py b/src/huggingface_hub/hub_mixin.py
new file mode 100644
index 00000000..081be72b
--- /dev/null
+++ b/src/huggingface_hub/hub_mixin.py
@@ -0,0 +1,246 @@
+import json
+import logging
+import os
+from typing import Dict, Optional
+
+import requests
+
+from .constants import CONFIG_NAME, PYTORCH_WEIGHTS_NAME
+from .file_download import cached_download, hf_hub_url, is_torch_available
+from .hf_api import HfApi, HfFolder
+from .repository import Repository
+
+
+if is_torch_available():
+    import torch
+
+
+logger = logging.getLogger(__name__)
+
+
+class ModelHubMixin(object):
+    def __init__(self, *args, **kwargs):
+        """
+        Mix this class with your torch-model class for ease process of saving & loading from huggingface-hub
+
+        Example::
+
+            >>> from huggingface_hub import ModelHubMixin
+
+            >>> class MyModel(nn.Module, ModelHubMixin):
+            ...    def __init__(self, **kwargs):
+            ...        super().__init__()
+            ...        self.config = kwargs.pop("config", None)
+            ...        self.layer = ...
+            ...    def forward(self, ...)
+            ...        return ...
+
+            >>> model = MyModel()
+            >>> model.save_pretrained("mymodel", push_to_hub=False) # Saving model weights in the directory
+            >>> model.push_to_hub("mymodel", "model-1") # Pushing model-weights to hf-hub
+
+            >>> # Downloading weights from hf-hub & model will be initialized from those weights
+            >>> model = MyModel.from_pretrained("username/mymodel@main")
+        """
+
+    def save_pretrained(
+        self,
+        save_directory: str,
+        config: Optional[dict] = None,
+        push_to_hub: bool = False,
+        **kwargs,
+    ):
+        """
+        Saving weights in local directory.
+
+        Parameters:
+            save_directory (:obj:`str`):
+                Specify directory in which you want to save weights.
+            config (:obj:`dict`, `optional`):
+                specify config (must be dict) incase you want to save it.
+            push_to_hub (:obj:`bool`, `optional`, defaults to :obj:`False`):
+                Set it to `True` in case you want to push your weights to huggingface_hub
+            model_id (:obj:`str`, `optional`, defaults to :obj:`save_directory`):
+                Repo name in huggingface_hub. If not specified, repo name will be same as `save_directory`
+            kwargs (:obj:`Dict`, `optional`):
+                kwargs will be passed to `push_to_hub`
+        """
+
+        os.makedirs(save_directory, exist_ok=True)
+
+        # saving config
+        if isinstance(config, dict):
+            path = os.path.join(save_directory, CONFIG_NAME)
+            with open(path, "w") as f:
+                json.dump(config, f)
+
+        # saving model weights
+        path = os.path.join(save_directory, PYTORCH_WEIGHTS_NAME)
+        self._save_pretrained(path)
+
+        if push_to_hub:
+            return self.push_to_hub(save_directory, **kwargs)
+
+    def _save_pretrained(self, path):
+        """
+        Overwrite this method in case you don't want to save complete model, rather some specific layers
+        """
+        model_to_save = self.module if hasattr(self, "module") else self
+        torch.save(model_to_save.state_dict(), path)
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name_or_path: Optional[str],
+        strict: bool = True,
+        map_location: Optional[str] = "cpu",
+        force_download: bool = False,
+        resume_download: bool = False,
+        proxies: Dict = None,
+        use_auth_token: Optional[str] = None,
+        cache_dir: Optional[str] = None,
+        local_files_only: bool = False,
+        **model_kwargs,
+    ):
+        r"""
+        Instantiate a pretrained pytorch model from a pre-trained model configuration from huggingface-hub.
+        The model is set in evaluation mode by default using ``model.eval()`` (Dropout modules are deactivated). To
+        train the model, you should first set it back in training mode with ``model.train()``.
+
+        Parameters:
+            pretrained_model_name_or_path (:obj:`str` or :obj:`os.PathLike`, `optional`):
+                Can be either:
+                    - A string, the `model id` of a pretrained model hosted inside a model repo on huggingface.co.
+                      Valid model ids can be located at the root-level, like ``bert-base-uncased``, or namespaced under
+                      a user or organization name, like ``dbmdz/bert-base-german-cased``.
+                    - You can add `revision` by appending `@` at the end of model_id simply like this: ``dbmdz/bert-base-german-cased@main``
+                      Revision is the specific model version to use. It can be a branch name, a tag name, or a commit id,
+                      since we use a git-based system for storing models and other artifacts on huggingface.co, so ``revision`` can be any identifier allowed by git.
+                    - A path to a `directory` containing model weights saved using
+                      :func:`~transformers.PreTrainedModel.save_pretrained`, e.g., ``./my_model_directory/``.
+                    - :obj:`None` if you are both providing the configuration and state dictionary (resp. with keyword
+                      arguments ``config`` and ``state_dict``).
+            cache_dir (:obj:`Union[str, os.PathLike]`, `optional`):
+                Path to a directory in which a downloaded pretrained model configuration should be cached if the
+                standard cache should not be used.
+            force_download (:obj:`bool`, `optional`, defaults to :obj:`False`):
+                Whether or not to force the (re-)download of the model weights and configuration files, overriding the
+                cached versions if they exist.
+            resume_download (:obj:`bool`, `optional`, defaults to :obj:`False`):
+                Whether or not to delete incompletely received files. Will attempt to resume the download if such a
+                file exists.
+            proxies (:obj:`Dict[str, str], `optional`):
+                A dictionary of proxy servers to use by protocol or endpoint, e.g., :obj:`{'http': 'foo.bar:3128',
+                'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
+            local_files_only(:obj:`bool`, `optional`, defaults to :obj:`False`):
+                Whether or not to only look at local files (i.e., do not try to download the model).
+            use_auth_token (:obj:`str` or `bool`, `optional`):
+                The token to use as HTTP bearer authorization for remote files. If :obj:`True`, will use the token
+                generated when running :obj:`transformers-cli login` (stored in :obj:`~/.huggingface`).
+            model_kwargs (:obj:`Dict`, `optional`)::
+                model_kwargs will be passed to the model during initialization
+        .. note::
+            Passing :obj:`use_auth_token=True` is required when you want to use a private model.
+        """
+
+        model_id = pretrained_model_name_or_path
+        map_location = torch.device(map_location)
+
+        revision = None
+        if len(model_id.split("@")) == 2:
+            model_id, revision = model_id.split("@")
+
+        if model_id in os.listdir() and CONFIG_NAME in os.listdir(model_id):
+            config_file = os.path.join(model_id, CONFIG_NAME)
+        else:
+            try:
+                config_url = hf_hub_url(
+                    model_id, filename=CONFIG_NAME, revision=revision
+                )
+                config_file = cached_download(
+                    config_url,
+                    cache_dir=cache_dir,
+                    force_download=force_download,
+                    proxies=proxies,
+                    resume_download=resume_download,
+                    local_files_only=local_files_only,
+                    use_auth_token=use_auth_token,
+                )
+            except requests.exceptions.RequestException:
+                logger.warning("config.json NOT FOUND in HuggingFace Hub")
+                config_file = None
+
+        if model_id in os.listdir():
+            print("LOADING weights from local directory")
+            model_file = os.path.join(model_id, PYTORCH_WEIGHTS_NAME)
+        else:
+            model_url = hf_hub_url(
+                model_id, filename=PYTORCH_WEIGHTS_NAME, revision=revision
+            )
+            model_file = cached_download(
+                model_url,
+                cache_dir=cache_dir,
+                force_download=force_download,
+                proxies=proxies,
+                resume_download=resume_download,
+                local_files_only=local_files_only,
+                use_auth_token=use_auth_token,
+            )
+
+        if config_file is not None:
+            with open(config_file, "r", encoding="utf-8") as f:
+                config = json.load(f)
+            model_kwargs.update({"config": config})
+
+        model = cls(**model_kwargs)
+
+        state_dict = torch.load(model_file, map_location=map_location)
+        model.load_state_dict(state_dict, strict=strict)
+        model.eval()
+
+        return model
+
+    @staticmethod
+    def push_to_hub(
+        save_directory: Optional[str],
+        model_id: Optional[str] = None,
+        repo_url: Optional[str] = None,
+        commit_message: Optional[str] = "add model",
+        organization: Optional[str] = None,
+        private: bool = None,
+    ) -> str:
+        """
+        Parameters:
+            save_directory (:obj:`Union[str, os.PathLike]`):
+                Directory having model weights & config.
+            model_id (:obj:`str`, `optional`, defaults to :obj:`save_directory`):
+                Repo name in huggingface_hub. If not specified, repo name will be same as `save_directory`
+            repo_url (:obj:`str`, `optional`):
+                Specify this in case you want to push to existing repo in hub.
+            organization (:obj:`str`, `optional`):
+                Organization in which you want to push your model.
+            private (:obj:`bool`, `optional`):
+                private: Whether the model repo should be private (requires a paid huggingface.co account)
+            commit_message (:obj:`str`, `optional`, defaults to :obj:`add model`):
+                Message to commit while pushing
+
+        Returns:
+            url to commit on remote repo.
+        """
+        if model_id is None:
+            model_id = save_directory
+
+        token = HfFolder.get_token()
+        if repo_url is None:
+            repo_url = HfApi().create_repo(
+                token,
+                model_id,
+                organization=organization,
+                private=private,
+                repo_type=None,
+                exist_ok=True,
+            )
+
+        repo = Repository(save_directory, clone_from=repo_url, use_auth_token=token)
+
+        return repo.push_to_hub(commit_message=commit_message)
diff --git a/src/huggingface_hub/repository.py b/src/huggingface_hub/repository.py
index 2eb7a422..e1607d36 100644
--- a/src/huggingface_hub/repository.py
+++ b/src/huggingface_hub/repository.py
@@ -159,14 +159,26 @@ class Repository:
                     encoding="utf-8",
                     cwd=self.local_dir,
                 )
-                subprocess.run(
-                    ["git", "remote", "add", "origin", repo_url],
+
+                output = subprocess.run(
+                    "git remote -v".split(),
                     stderr=subprocess.PIPE,
                     stdout=subprocess.PIPE,
                     check=True,
                     encoding="utf-8",
                     cwd=self.local_dir,
                 )
+
+                if "origin" not in output.stdout.split():
+                    subprocess.run(
+                        ["git", "remote", "add", "origin", repo_url],
+                        stderr=subprocess.PIPE,
+                        stdout=subprocess.PIPE,
+                        check=True,
+                        encoding="utf-8",
+                        cwd=self.local_dir,
+                    )
+
                 subprocess.run(
                     "git fetch".split(),
                     stderr=subprocess.PIPE,
@@ -183,15 +195,27 @@ class Repository:
                     check=True,
                     cwd=self.local_dir,
                 )
-                # TODO(check if we really want the --force flag)
-                subprocess.run(
-                    "git checkout origin/main -ft".split(),
+
+                output = subprocess.run(
+                    "git branch".split(),
                     stderr=subprocess.PIPE,
                     stdout=subprocess.PIPE,
-                    encoding="utf-8",
                     check=True,
+                    encoding="utf-8",
                     cwd=self.local_dir,
                 )
+
+                if "main" not in output.stdout.split():
+                    # TODO(check if we really want the --force flag)
+                    subprocess.run(
+                        "git checkout origin/main -ft".split(),
+                        stderr=subprocess.PIPE,
+                        stdout=subprocess.PIPE,
+                        encoding="utf-8",
+                        check=True,
+                        cwd=self.local_dir,
+                    )
+
         except subprocess.CalledProcessError as exc:
             raise EnvironmentError(exc.stderr)
 


Test patch:
diff --git a/tests/test_file_download.py b/tests/test_file_download.py
index 9ad03479..f53b41b2 100644
--- a/tests/test_file_download.py
+++ b/tests/test_file_download.py
@@ -22,7 +22,12 @@ from huggingface_hub.constants import (
 )
 from huggingface_hub.file_download import cached_download, filename_to_url, hf_hub_url
 
-from .testing_utils import DUMMY_UNKWOWN_IDENTIFIER, SAMPLE_DATASET_IDENTIFIER
+from .testing_utils import (
+    DUMMY_UNKWOWN_IDENTIFIER,
+    SAMPLE_DATASET_IDENTIFIER,
+    OfflineSimulationMode,
+    offline,
+)
 
 
 MODEL_ID = DUMMY_UNKWOWN_IDENTIFIER
@@ -51,13 +56,26 @@ DATASET_SAMPLE_PY_FILE = "custom_squad.py"
 
 class CachedDownloadTests(unittest.TestCase):
     def test_bogus_url(self):
-        # This lets us simulate no connection
-        # as the error raised is the same
-        # `ConnectionError`
         url = "https://bogus"
         with self.assertRaisesRegex(ValueError, "Connection error"):
             _ = cached_download(url)
 
+    def test_no_connection(self):
+        invalid_url = hf_hub_url(
+            MODEL_ID, filename=CONFIG_NAME, revision=REVISION_ID_INVALID
+        )
+        valid_url = hf_hub_url(
+            MODEL_ID, filename=CONFIG_NAME, revision=REVISION_ID_DEFAULT
+        )
+        self.assertIsNotNone(cached_download(valid_url, force_download=True))
+        for offline_mode in OfflineSimulationMode:
+            with offline(mode=offline_mode):
+                with self.assertRaisesRegex(ValueError, "Connection error"):
+                    _ = cached_download(invalid_url)
+                with self.assertRaisesRegex(ValueError, "Connection error"):
+                    _ = cached_download(valid_url, force_download=True)
+                self.assertIsNotNone(cached_download(valid_url))
+
     def test_file_not_found(self):
         # Valid revision (None) but missing file.
         url = hf_hub_url(MODEL_ID, filename="missing.bin")
diff --git a/tests/test_hubmixin.py b/tests/test_hubmixin.py
new file mode 100644
index 00000000..23a3e5a2
--- /dev/null
+++ b/tests/test_hubmixin.py
@@ -0,0 +1,66 @@
+import unittest
+
+from huggingface_hub.file_download import is_torch_available
+from huggingface_hub.hub_mixin import ModelHubMixin
+
+
+if is_torch_available():
+    import torch.nn as nn
+
+
+HUGGINGFACE_ID = "vasudevgupta"
+DUMMY_REPO_NAME = "dummy"
+
+
+def require_torch(test_case):
+    """
+    Decorator marking a test that requires PyTorch.
+
+    These tests are skipped when PyTorch isn't installed.
+
+    """
+    if not is_torch_available():
+        return unittest.skip("test requires PyTorch")(test_case)
+    else:
+        return test_case
+
+
+@require_torch
+class DummyModel(ModelHubMixin):
+    def __init__(self, **kwargs):
+        super().__init__()
+        self.config = kwargs.pop("config", None)
+        self.l1 = nn.Linear(2, 2)
+
+    def forward(self, x):
+        return self.l1(x)
+
+
+@require_torch
+class DummyModelTest(unittest.TestCase):
+    def test_save_pretrained(self):
+        model = DummyModel()
+        model.save_pretrained(DUMMY_REPO_NAME)
+        model.save_pretrained(
+            DUMMY_REPO_NAME, config={"num": 12, "act": "gelu"}, push_to_hub=True
+        )
+        model.save_pretrained(
+            DUMMY_REPO_NAME, config={"num": 24, "act": "relu"}, push_to_hub=True
+        )
+        model.save_pretrained(
+            "dummy-wts", config=None, push_to_hub=True, model_id=DUMMY_REPO_NAME
+        )
+
+    def test_from_pretrained(self):
+        model = DummyModel()
+        model.save_pretrained(
+            DUMMY_REPO_NAME, config={"num": 7, "act": "gelu_fast"}, push_to_hub=True
+        )
+
+        model = DummyModel.from_pretrained(f"{HUGGINGFACE_ID}/{DUMMY_REPO_NAME}@main")
+        self.assertTrue(model.config == {"num": 7, "act": "gelu_fast"})
+
+    def test_push_to_hub(self):
+        model = DummyModel()
+        model.save_pretrained("dummy-wts", push_to_hub=False)
+        model.push_to_hub("dummy-wts", model_id=DUMMY_REPO_NAME)
diff --git a/tests/test_offline_utils.py b/tests/test_offline_utils.py
new file mode 100644
index 00000000..076bb29a
--- /dev/null
+++ b/tests/test_offline_utils.py
@@ -0,0 +1,36 @@
+from io import BytesIO
+
+import pytest
+
+import requests
+from huggingface_hub.file_download import http_get
+
+from .testing_utils import (
+    OfflineSimulationMode,
+    RequestWouldHangIndefinitelyError,
+    offline,
+)
+
+
+def test_offline_with_timeout():
+    with offline(OfflineSimulationMode.CONNECTION_TIMES_OUT):
+        with pytest.raises(RequestWouldHangIndefinitelyError):
+            requests.request("GET", "https://huggingface.co")
+        with pytest.raises(requests.exceptions.ConnectTimeout):
+            requests.request("GET", "https://huggingface.co", timeout=1.0)
+        with pytest.raises(requests.exceptions.ConnectTimeout):
+            http_get("https://huggingface.co", BytesIO())
+
+
+def test_offline_with_connection_error():
+    with offline(OfflineSimulationMode.CONNECTION_FAILS):
+        with pytest.raises(requests.exceptions.ConnectionError):
+            requests.request("GET", "https://huggingface.co")
+        with pytest.raises(requests.exceptions.ConnectionError):
+            http_get("https://huggingface.co", BytesIO())
+
+
+def test_offline_with_datasets_offline_mode_enabled():
+    with offline(OfflineSimulationMode.HF_HUB_OFFLINE_SET_TO_1):
+        with pytest.raises(ConnectionError):
+            http_get("https://huggingface.co", BytesIO())
diff --git a/tests/testing_utils.py b/tests/testing_utils.py
index 1c98239a..21668ef2 100644
--- a/tests/testing_utils.py
+++ b/tests/testing_utils.py
@@ -1,6 +1,9 @@
 import os
 import unittest
+from contextlib import contextmanager
 from distutils.util import strtobool
+from enum import Enum
+from unittest.mock import patch
 
 
 SMALL_MODEL_IDENTIFIER = "julien-c/bert-xsmall-dummy"
@@ -55,3 +58,70 @@ def require_git_lfs(test_case):
         return unittest.skip("test of git lfs workflow")(test_case)
     else:
         return test_case
+
+
+class RequestWouldHangIndefinitelyError(Exception):
+    pass
+
+
+class OfflineSimulationMode(Enum):
+    CONNECTION_FAILS = 0
+    CONNECTION_TIMES_OUT = 1
+    HF_HUB_OFFLINE_SET_TO_1 = 2
+
+
+@contextmanager
+def offline(mode=OfflineSimulationMode.CONNECTION_FAILS, timeout=1e-16):
+    """
+    Simulate offline mode.
+
+    There are three offline simulatiom modes:
+
+    CONNECTION_FAILS (default mode): a ConnectionError is raised for each network call.
+        Connection errors are created by mocking socket.socket
+    CONNECTION_TIMES_OUT: the connection hangs until it times out.
+        The default timeout value is low (1e-16) to speed up the tests.
+        Timeout errors are created by mocking requests.request
+    HF_HUB_OFFLINE_SET_TO_1: the HF_HUB_OFFLINE_SET_TO_1 environment variable is set to 1.
+        This makes the http/ftp calls of the library instantly fail and raise an OfflineModeEmabled error.
+    """
+    import socket
+
+    from requests import request as online_request
+
+    def timeout_request(method, url, **kwargs):
+        # Change the url to an invalid url so that the connection hangs
+        invalid_url = "https://10.255.255.1"
+        if kwargs.get("timeout") is None:
+            raise RequestWouldHangIndefinitelyError(
+                f"Tried a call to {url} in offline mode with no timeout set. Please set a timeout."
+            )
+        kwargs["timeout"] = timeout
+        try:
+            return online_request(method, invalid_url, **kwargs)
+        except Exception as e:
+            # The following changes in the error are just here to make the offline timeout error prettier
+            e.request.url = url
+            max_retry_error = e.args[0]
+            max_retry_error.args = (
+                max_retry_error.args[0].replace("10.255.255.1", f"OfflineMock[{url}]"),
+            )
+            e.args = (max_retry_error,)
+            raise
+
+    def offline_socket(*args, **kwargs):
+        raise socket.error("Offline mode is enabled.")
+
+    if mode is OfflineSimulationMode.CONNECTION_FAILS:
+        # inspired from https://stackoverflow.com/a/18601897
+        with patch("socket.socket", offline_socket):
+            yield
+    elif mode is OfflineSimulationMode.CONNECTION_TIMES_OUT:
+        # inspired from https://stackoverflow.com/a/904609
+        with patch("requests.request", timeout_request):
+            yield
+    elif mode is OfflineSimulationMode.HF_HUB_OFFLINE_SET_TO_1:
+        with patch("huggingface_hub.constants.HF_HUB_OFFLINE", True):
+            yield
+    else:
+        raise ValueError("Please use a value from the OfflineSimulationMode enum.")