Repository: hiyouga/LLaMA-Factory
File: src/llamafactory/extras/misc.py (Python)

Code before:
def torch_gc() -> None:
    r"""
    Collects GPU or NPU memory.
    """
    gc.collect()
    if is_torch_xpu_available():
        torch.xpu.empty_cache()
    elif is_torch_npu_available():
        torch.npu.empty_cache()
    elif is_torch_mps_available():
        torch.mps.empty_cache()
    elif is_torch_cuda_available():
        torch.cuda.empty_cache()


def try_download_model_from_other_hub(model_args: "ModelArguments") -> str:
    if (not use_openmind() and not use_modelscope()) or os.path.exists(model_args.model_name_or_path):
        return model_args.model_name_or_path

    if use_openmind():
        try:
            import openmind
            from openmind.utils.hub import snapshot_download

            revision = "main" if model_args.model_revision == "main" else model_args.model_revision
            return snapshot_download(model_args.model_name_or_path, revision=revision, cache_dir=model_args.cache_dir)
        except ImportError:
            raise ImportError("Please install openmind and openmind_hub via `pip install openmind -U`")

    if use_modelscope():
        try:
            from modelscope import snapshot_download

            revision = "master" if model_args.model_revision == "main" else model_args.model_revision
            return snapshot_download(model_args.model_name_or_path, revision=revision, cache_dir=model_args.cache_dir)
        except ImportError:
            raise ImportError("Please install modelscope via `pip install modelscope -U`")

def use_openmind() -> bool:
    return os.environ.get("USE_OPENMIND_HUB", "0").lower() in ["true", "1"]

def use_modelscope() -> bool:
    return os.environ.get("USE_MODELSCOPE_HUB", "0").lower() in ["true", "1"]


Review comment (performance):
`return snapshot_download(model_args.model_name_or_path, revision=model_args.model_revision, cache_dir=model_args.cache_dir)`

Code after:
def torch_gc() -> None:
    r"""
    Collects GPU or NPU memory.
    """
    gc.collect()
    if is_torch_xpu_available():
        torch.xpu.empty_cache()
    elif is_torch_npu_available():
        torch.npu.empty_cache()
    elif is_torch_mps_available():
        torch.mps.empty_cache()
    elif is_torch_cuda_available():
        torch.cuda.empty_cache()


def try_download_model_from_other_hub(model_args: "ModelArguments") -> str:
    if (not use_openmind() and not use_modelscope()) or os.path.exists(model_args.model_name_or_path):
        return model_args.model_name_or_path

    if use_openmind():
        try:
            from openmind.utils.hub import snapshot_download

            return snapshot_download(model_args.model_name_or_path, revision=model_args.model_revision, cache_dir=model_args.cache_dir)
        except ImportError:
            raise ImportError("Please install openmind and openmind_hub via `pip install openmind -U`")

    if use_modelscope():
        try:
            from modelscope import snapshot_download

            revision = "master" if model_args.model_revision == "main" else model_args.model_revision
            return snapshot_download(model_args.model_name_or_path, revision=revision, cache_dir=model_args.cache_dir)
        except ImportError:
            raise ImportError("Please install modelscope via `pip install modelscope -U`")

def use_openmind() -> bool:
    return os.environ.get("USE_OPENMIND_HUB", "0").lower() in ["true", "1"]

def use_modelscope() -> bool:
    return os.environ.get("USE_MODELSCOPE_HUB", "0").lower() in ["true", "1"]