etna-team / etna

ETNA – Time-Series Library
https://docs.etna.ai
Apache License 2.0
142 stars 7 forks source link

Add possibility to load pretrained embedding models #452

Closed egoriyaa closed 2 months ago

egoriyaa commented 3 months ago

🚀 Task description

Add possibility to load pretrained embedding models

Plan

  1. For TS2VecEmbeddingModel and TSTCCEmbeddingModel classes modify load method in such way:

    def load(cls, path: Optional[pathlib.Path] = None, model_name: Optional[str] = None)
    • both parameters are specified - download the pretrained model and save it to path
    • path is specified, model_name is not - download the local model from path
    • path is not specified, model_name is specified - download the pretrained model to the etna folder in the home directory
    • both parameters are not specified - raise the exception
  2. For TS2VecEmbeddingModel and TSTCCEmbeddingModel classes add static method list_models

    • Return list of available models
  3. For TS2VecEmbeddingModel and TSTCCEmbeddingModel add parameter freezed in constructor. - If freezed is True, freeze model in constructor.

  4. Modify 210-embedding_models notebook to show how to use pretrained models

Test cases

  1. For load method:

    • 4 combinations of parameters
    • model_name not in list_models
  2. list_models returns right model names

Additional context

No response