Add possibility to load pretrained embedding models

etna/transforms/embeddings/models/ts2vec.py

@@ -1,8 +1,14 @@
+import os
 import pathlib
 import tempfile
+import warnings
 import zipfile
+from pathlib import Path
+from typing import List
 from typing import Literal
 from typing import Optional
+from urllib import request
+from urllib.error import HTTPError
 
 import numpy as np
 
@@ -12,6 +18,8 @@
 if SETTINGS.torch_required:
     from etna.libs.ts2vec import TS2Vec
 
+_DOWNLOAD_PATH = Path.home() / ".etna" / "embeddings" / "ts2vec"
+
 
 class TS2VecEmbeddingModel(BaseEmbeddingModel):
     """TS2Vec embedding model.
@@ -257,7 +265,7 @@
                 archive.write(model_save_path, "model.zip")
 
     @classmethod
-    def load(cls, path: pathlib.Path) -> "TS2VecEmbeddingModel":
+    def load(cls, path: Optional[pathlib.Path] = None, model_name: Optional[str] = None) -> "TS2VecEmbeddingModel":
         """Load an object.
 
         Model's weights are transferred to cpu during loading.
@@ -267,11 +275,50 @@
         path:
             Path to load object from.
 
+            - if `path` is not None and `model_name` is None, load the local model from `path`.
+            - if `path` is None and `model_name` is not None, save the external `model_name` model to the etna folder in the home directory and load it.
+            If `path` exists, external model will not be loaded.
+            - if `path` is not  None and `model_name` is not None, save the external `model_name` model to `path` and load it.
+            If `path` exists, external model will not be loaded.
+
+        model_name:
+            name of external model to load. To get list of available models use `list_models` method.
+
         Returns
         -------
         :
             Loaded object.
+
+        Raises
+        ------
+        ValueError:
+            If non of parameters path and model_name are set.
+        NotImplementedError:
+            If model_name is not from list of available model names
         """
+        if model_name is not None:
+            if path is None:
+                path = _DOWNLOAD_PATH / f"{model_name}.zip"
+            if os.path.exists(path):
+                warnings.warn(
+                    f"Path {path} already exists. Model {model_name} will not be downloaded. Loading existing local model."
+                )
+            else:
+                directory = os.path.dirname(path)
+                # If path not in current directory and it doesn't exist
+                if directory and not os.path.exists(directory):
+                    os.makedirs(directory)
+
+                try:
+                    url = f"http://etna-github-prod.cdn-tinkoff.ru/embeddings/ts2vec/{model_name}.zip"
+                    request.urlretrieve(url=url, filename=path)
+                except HTTPError:
+                    raise NotImplementedError(
+                        f"Model {model_name} is not available. To get list of available models use `list_models` method."
+                    )
+        elif path is None and model_name is None:
+            raise ValueError("Both path and model_name are not specified. Specify one parameter at least.")
+
         obj: TS2VecEmbeddingModel = super().load(path=path)
         obj.embedding_model = TS2Vec(
             input_dims=obj.input_dims,
@@ -292,3 +339,8 @@
                 obj.embedding_model.load(fn=str(model_path))
 
         return obj
+
+    @staticmethod
+    def list_models() -> List[str]:
+        """Return a list of available pretrained models."""
+        return ["ts2vec_tiny"]

etna/transforms/embeddings/models/tstcc.py

@@ -1,7 +1,14 @@
+import os
 import pathlib
 import tempfile
+import warnings
 import zipfile
+from pathlib import Path
+from typing import List
 from typing import Literal
+from typing import Optional
+from urllib import request
+from urllib.error import HTTPError
 
 import numpy as np
 
@@ -11,6 +18,8 @@
 if SETTINGS.torch_required:
     from etna.libs.tstcc import TSTCC
 
+_DOWNLOAD_PATH = Path.home() / ".etna" / "embeddings" / "tstcc"
+
 
 class TSTCCEmbeddingModel(BaseEmbeddingModel):
     """TSTCC embedding model.
@@ -252,7 +261,7 @@
                 archive.write(model_save_path, "model.zip")
 
     @classmethod
-    def load(cls, path: pathlib.Path) -> "TSTCCEmbeddingModel":
+    def load(cls, path: Optional[pathlib.Path] = None, model_name: Optional[str] = None) -> "TSTCCEmbeddingModel":
         """Load an object.
 
         Model's weights are transferred to cpu during loading.
@@ -262,11 +271,50 @@
         path:
             Path to load object from.
 
+            - if `path` is not None and `model_name` is None, load the local model from `path`.
+            - if `path` is None and `model_name` is not None, save the external `model_name` model to the etna folder in the home directory and load it.
+            If `path` exists, external model will not be loaded.
+            - if `path` is not  None and `model_name` is not None, save the external `model_name` model to `path` and load it.
+            If `path` exists, external model will not be loaded.
+
+        model_name:
+            name of external model to load. To get list of available models use `list_models` method.
+
         Returns
         -------
         :
             Loaded object.
+
+        Raises
+        ------
+        ValueError:
+            If non of parameters path and model_name are set.
+        NotImplementedError:
+            If model_name is not from list of available model names
         """
+        if model_name is not None:
+            if path is None:
+                path = _DOWNLOAD_PATH / f"{model_name}.zip"
+            if os.path.exists(path):
+                warnings.warn(
+                    f"Path {path} already exists. Model {model_name} will not be downloaded. Loading existing local model."
+                )
+            else:
+                directory = os.path.dirname(path)
+                # If path not in current directory and it doesn't exist
+                if directory and not os.path.exists(directory):
+                    os.makedirs(directory)
+
+                try:
+                    url = f"http://etna-github-prod.cdn-tinkoff.ru/embeddings/tstcc/{model_name}.zip"
+                    request.urlretrieve(url=url, filename=path)
+                except HTTPError:
+                    raise NotImplementedError(
+                        f"Model {model_name} is not available. To get list of available models use `list_models` method."
+                    )
+        elif path is None and model_name is None:
+            raise ValueError("Both path and model_name are not specified. Specify one parameter at least.")
+
         obj: TSTCCEmbeddingModel = super().load(path=path)
         obj.embedding_model = TSTCC(
             input_dims=obj.input_dims,
@@ -293,3 +341,8 @@
                 obj.embedding_model.load(fn=str(model_path))
 
         return obj
+
+    @staticmethod
+    def list_models() -> List[str]:
+        """Return a list of available pretrained models."""
+        return ["tstcc_medium"]