def download_model()

in picotron/utils.py [0:0]

• 13 lines of code
• 5 McCabe index (conditional complexity)

def download_model(model_name, hf_token):
    dst = os.path.join("hf_model", model_name)
    os.makedirs(dst, exist_ok=True)
    # check if model is already downloaded
    if os.path.exists(os.path.join(dst, "config.json")):
        print(f"Model {model_name} already exists at {dst}")
        return
    # Download HF model safetensors at the "dst" directory
    huggingface_hub.constants.HF_HUB_ENABLE_HF_TRANSFER = True
    print("Downloading SafeTensors files...")
    huggingface_hub.snapshot_download(model_name, repo_type="model", local_dir="hf_model_safetensors", token=hf_token,
                                      allow_patterns=["*.safetensors", "*.json"])
    # Check if the model has SafeTensors files
    if not glob.glob("hf_model_safetensors/*.safetensors"):
        raise ValueError(f"Model {model_name} does not have SafeTensors files.")
    print("SafeTensors files downloaded successfully! ✅")