modelscope_utils.py 1.3 KB

12345678910111213141516171819202122232425262728293031323334353637383940
  1. import os
  2. from modelscope.hub.snapshot_download import snapshot_download
  3. from pathlib import Path
  4. def check_model_dir(model_dir, model_name: str = "damo/speech_fsmn_vad_zh-cn-16k-common-pytorch"):
  5. model_dir = "/Users/zhifu/test_modelscope_pipeline/FSMN-VAD"
  6. cache_root = os.path.dirname(model_dir)
  7. dst_dir_root = os.path.join(cache_root, ".cache")
  8. dst = os.path.join(dst_dir_root, model_name)
  9. dst_dir = os.path.dirname(dst)
  10. os.makedirs(dst_dir, exist_ok=True)
  11. if not os.path.exists(dst):
  12. os.symlink(model_dir, dst)
  13. model_dir = snapshot_download(model_name, cache_dir=dst_dir_root)
  14. def get_default_cache_dir():
  15. """
  16. default base dir: '~/.cache/modelscope'
  17. """
  18. default_cache_dir = Path.home().joinpath('.cache', 'modelscope')
  19. return default_cache_dir
  20. def get_cache_dir(model_id):
  21. """cache dir precedence:
  22. function parameter > environment > ~/.cache/modelscope/hub
  23. Args:
  24. model_id (str, optional): The model id.
  25. Returns:
  26. str: the model_id dir if model_id not None, otherwise cache root dir.
  27. """
  28. default_cache_dir = get_default_cache_dir()
  29. base_path = os.getenv('MODELSCOPE_CACHE',
  30. os.path.join(default_cache_dir, 'hub'))
  31. return base_path if model_id is None else os.path.join(
  32. base_path, model_id + '/')