Download pretrained model from AIStudio, MODELSCOPE or HUGGINGFACE automatically
(model_name_or_path, revision="master")
| 753 | |
| 754 | |
| 755 | def retrive_model_from_server(model_name_or_path, revision="master"): |
| 756 | """ |
| 757 | Download pretrained model from AIStudio, MODELSCOPE or HUGGINGFACE automatically |
| 758 | """ |
| 759 | if os.path.exists(model_name_or_path): |
| 760 | return model_name_or_path |
| 761 | model_source = envs.FD_MODEL_SOURCE |
| 762 | local_path = envs.FD_MODEL_CACHE |
| 763 | repo_id = model_name_or_path |
| 764 | if model_source == "AISTUDIO": |
| 765 | try: |
| 766 | if repo_id.lower().strip().startswith("baidu"): |
| 767 | repo_id = "PaddlePaddle" + repo_id.strip()[5:] |
| 768 | if local_path is None: |
| 769 | local_path = f'{os.getenv("HOME")}' |
| 770 | local_path = f"{local_path}/{repo_id}" |
| 771 | aistudio_download(repo_id=repo_id, revision=revision, local_dir=local_path) |
| 772 | model_name_or_path = local_path |
| 773 | except requests.exceptions.ConnectTimeout: |
| 774 | if os.path.exists(local_path): |
| 775 | llm_logger.error( |
| 776 | f"Failed to connect to aistudio, but detected that the model directory {local_path} exists. Attempting to start." |
| 777 | ) |
| 778 | return local_path |
| 779 | except Exception: |
| 780 | raise Exception( |
| 781 | f"The {revision} of {model_name_or_path} is not exist. Please check the model name or revision." |
| 782 | ) |
| 783 | elif model_source == "MODELSCOPE": |
| 784 | try: |
| 785 | from modelscope.hub.snapshot_download import ( |
| 786 | snapshot_download as modelscope_download, |
| 787 | ) |
| 788 | |
| 789 | if repo_id.lower().strip().startswith("baidu"): |
| 790 | repo_id = "PaddlePaddle" + repo_id.strip()[5:] |
| 791 | if local_path is None: |
| 792 | local_path = f'{os.getenv("HOME")}' |
| 793 | local_path = f"{local_path}/{repo_id}" |
| 794 | modelscope_download(repo_id=repo_id, revision=revision, local_dir=local_path) |
| 795 | model_name_or_path = local_path |
| 796 | except requests.exceptions.ConnectTimeout: |
| 797 | if os.path.exists(local_path): |
| 798 | llm_logger.error( |
| 799 | f"Failed to connect to modelscope, but detected that the model directory {local_path} exists. Attempting to start." |
| 800 | ) |
| 801 | return local_path |
| 802 | except Exception: |
| 803 | raise Exception( |
| 804 | f"The {revision} of {model_name_or_path} is not exist. Please check the model name or revision." |
| 805 | ) |
| 806 | elif model_source == "HUGGINGFACE": |
| 807 | try: |
| 808 | from huggingface_hub._snapshot_download import ( |
| 809 | snapshot_download as huggingface_download, |
| 810 | ) |
| 811 | |
| 812 | if revision == "master": |