update docs (#243)

Author: emotionor

docs/source/features.md

@@ -1,3 +1,5 @@
 # New Features
 
-## Coming Soon...
+## 2024-06-25
+
+Unimol_tools has been publish to pypi! Huggingface has been used to manage the pretrain models.

docs/source/index.rst

@@ -24,6 +24,7 @@ Uni-Mol tools is a easy-use wrappers for property prediction,representation and
    models
    task
    utils
+   weight
    features
 
 Indices and tables

docs/source/weight.rst

@@ -0,0 +1,30 @@
+.. _weights:
+
+Weights
+=======
+
+We recommend installing ``huggingface_hub`` so that the required UniMol models can be automatically downloaded at runtime! It can be installed by:
+
+.. code-block:: bash
+
+    pip install huggingface_hub
+
+``huggingface_hub`` allows you to easily download and manage models from the Hugging Face Hub, which is key for using UniMol models.
+
+Models in Huggingface
+---------------------
+
+The UniMol pretrained models can be found at `dptech/Uni-Mol-Models <https://huggingface.co/dptech/Uni-Mol-Models/tree/main>`_.
+
+If the download is slow, you can use other mirrors, such as:
+
+.. code-block:: bash
+
+    export HF_ENDPOINT=https://hf-mirror.com
+
+Setting the ``HF_ENDPOINT`` environment variable specifies the mirror address for the Hugging Face Hub to use when downloading models.
+
+`unimol_tools.weights.weight_hub.py <https://github.com/dptech-corp/Uni-Mol/blob/docs/unimol_tools/unimol_tools/weights/weighthub.py>`_ control the logger.
+
+.. automodule:: unimol_tools.weights.weighthub
+   :members:

unimol_tools/unimol_tools/weights/weighthub.py

@@ -16,11 +16,10 @@ def snapshot_download(*args, **kwargs):
 def weight_download(pretrain, save_path, local_dir_use_symlinks=True):
     """
     Downloads the specified pretrained model weights.
-    
-    Args:
-        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.
+
+    :param pretrain: (str), The name of the pretrained model to download.
+    :param save_path: (str), The directory where the weights should be saved.
+    :param local_dir_use_symlinks: (bool, optional), Whether to use symlinks for the local directory. Defaults to True.
     """
     if os.path.exists(os.path.join(save_path, pretrain)):
         logger.info(f'{pretrain} exists in {save_path}')
@@ -40,8 +39,7 @@ def download_all_weights(local_dir_use_symlinks=False):
     """
     Downloads all available pretrained model weights to the WEIGHT_DIR.
     
-    Args:
-        local_dir_use_symlinks (bool, optional): Whether to use symlinks for the local directory. Defaults to False.
+    :param local_dir_use_symlinks: (bool, optional), Whether to use symlinks for the local directory. Defaults to False.
     """
     logger.info(f'Downloading all weights to {WEIGHT_DIR}')
     snapshot_download(