Weights

We recommend installing huggingface_hub so that the required Uni-Mol models can be automatically downloaded at runtime! It can be installed by:

pip install huggingface_hub

huggingface_hub allows you to easily download and manage models from the Hugging Face Hub, which is key for using Uni-Mol models.

Models in Huggingface

The Uni-Mol pretrained models can be found at dptech/Uni-Mol-Models.

If the download is slow, you can use other mirrors, such as:

export HF_ENDPOINT=https://hf-mirror.com

By default unimol_tools first tries the official Hugging Face endpoint. If that fails and HF_ENDPOINT is not set, it automatically retries with https://hf-mirror.com. Set the variable yourself to choose a specific endpoint.

unimol_tools.weights.weight_hub.py control the logger.

unimol_tools.weights.weighthub.get_weight_dir()[source]

Return the directory where weights should be stored.

unimol_tools.weights.weighthub._snapshot_download_with_fallback(**kwargs)[source]

Try downloading with the current HF_ENDPOINT and fall back to the mirror.

The mirror is only tried when the user has not explicitly set HF_ENDPOINT and the first attempt fails.

unimol_tools.weights.weighthub.log_weights_dir()[source]

Logs the directory where the weights are stored.

unimol_tools.weights.weighthub.weight_download(pretrain, save_path, local_dir_use_symlinks=True)[source]

Downloads the specified pretrained model weights.

Parameters:
  • pretrain – (str), The name of the pretrained model to download.

  • save_path – (str), The directory where the weights should be saved.

  • local_dir_use_symlinks – (bool, optional), Whether to use symlinks for the local directory. Defaults to True.

unimol_tools.weights.weighthub.weight_download_v2(pretrain, save_path, local_dir_use_symlinks=True)[source]

Downloads the specified pretrained model weights.

Parameters:
  • pretrain – (str), The name of the pretrained model to download.

  • save_path – (str), The directory where the weights should be saved.

  • local_dir_use_symlinks – (bool, optional), Whether to use symlinks for the local directory. Defaults to True.

unimol_tools.weights.weighthub.download_all_weights(local_dir_use_symlinks=False)[source]

Downloads all available pretrained model weights to the WEIGHT_DIR.

Parameters:

local_dir_use_symlinks – (bool, optional), Whether to use symlinks for the local directory. Defaults to False.