Uni-Mol

https://unimol.readthedocs.io/en/latest/
Admin
Uni-Mol is a universal large-scale three-dimensional Molecular Representation Learning (MRL)...
Tokens:65,145
Snippets:222
Trust Score:8.3
Update:1 year ago
Show doc for...
Tokens:
Raw
### Install Uni-Mol from Source

Source: https://unimol.readthedocs.io/en/latest/installation

Provides commands to install Uni-Mol from its source code repository. It includes installing dependencies from requirements.txt, cloning the repository, navigating to the tools directory, and running the setup script.

```bash
## Dependencies installation
pip install -r requirements.txt

## Clone repository
git clone https://github.com/deepmodeling/Uni-Mol.git
cd Uni-Mol/unimol_tools

## Install
python setup.py install
```

--------------------------------

### Install Uni-Mol from Source

Source: https://unimol.readthedocs.io/en/latest/_sources/installation

Installs Uni-Mol by cloning the repository and building from source. This method involves installing dependencies from `requirements.txt`, cloning the git repository, and then running the setup script.

```bash
pip install -r requirements.txt

git clone https://github.com/deepmodeling/Uni-Mol.git
cd Uni-Mol/unimol_tools

python setup.py install
```

--------------------------------

### Install huggingface_hub for Model Downloading

Source: https://unimol.readthedocs.io/en/latest/_sources/installation

Installs the `huggingface_hub` library, which facilitates the automatic download of Uni-Mol models from the Hugging Face Hub at runtime. This is an optional but recommended step.

```bash
pip install huggingface_hub
```

--------------------------------

### Install Uni-Mol from PyPI

Source: https://unimol.readthedocs.io/en/latest/_sources/installation

Installs the Uni-Mol tools package using pip. This is the recommended installation method. It requires PyTorch and optionally installs `huggingface_hub` for automatic model downloading.

```bash
pip install unimol_tools
```

--------------------------------

### Continue Training (Re-train) Uni-Mol Models

Source: https://unimol.readthedocs.io/en/latest/_sources/quickstart

This example illustrates how to continue training an existing Uni-Mol model. It involves specifying the task, data type, and training parameters, along with the path to the previously saved model directory using `load_model_dir`. This allows for iterative model improvement.

```python
clf = MolTrain(task='regression',
                data_type='molecule',
                epochs=10,
                batch_size=16,
                save_path='./model_dir',
                remove_hs=False,
                target_cols='TARGET',
                )
pred = clf.fit(data = train_data)
# After train a model, set load_model_dir='./model_dir' to continue training

clf2 = MolTrain(task='regression',
                data_type='molecule',
                epochs=10,
                batch_size=16,
                save_path='./retrain_model_dir',
                remove_hs=False,
                target_cols='TARGET',
                load_model_dir='./model_dir',
                )

pred2 = clf.fit(data = retrain_data)
```

--------------------------------

### Configure Hugging Face Hub Mirror

Source: https://unimol.readthedocs.io/en/latest/_sources/installation

Sets an environment variable to specify an alternative mirror for downloading models from the Hugging Face Hub. This can be useful if the default download is slow.

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

--------------------------------

### Install huggingface_hub for Uni-Mol Model Downloads

Source: https://unimol.readthedocs.io/en/latest/weight

Installs the huggingface_hub library, which is recommended for automatically downloading Uni-Mol pretrained models. This library simplifies model management from the Hugging Face Hub.

```bash
pip install huggingface_hub

```

--------------------------------

### Train and Predict Molecule Properties with UniMol

Source: https://unimol.readthedocs.io/en/latest/_sources/quickstart

This snippet shows how to train a classification model using SMILES strings and corresponding target values. It also demonstrates how to load a pre-trained model to make predictions on test data. Ensure data is provided in a CSV or TXT file format with SMILES as input.

```python
from unimol_tools import MolTrain, MolPredict
clf = MolTrain(task='classification', 
                data_type='molecule', 
                epochs=10, 
                batch_size=16, 
                metrics='auc',
                model_name='unimolv1', # avaliable: unimolv1, unimolv2
                model_size='84m', # work when model_name is unimolv2. avaliable: 84m, 164m, 310m, 570m, 1.1B.
                )
pred = clf.fit(data = train_data)
# currently support data with smiles based csv/txt file

clf = MolPredict(load_model='../exp')
res = clf.predict(data = test_data)
```

--------------------------------

### Generate Conformers (Multi-Process)

Source: https://unimol.readthedocs.io/en/latest/_modules/unimol_tools/data/conformer

Initiates the conformer generation process for a list of SMILES strings, utilizing multiprocessing for efficiency. If `multi_process` is enabled, it uses a `Pool` of workers to process SMILES strings in parallel. Otherwise, it processes them sequentially. Logs the start of the process.

```python
def transform(self, smiles_list):
        logger.info('Start generating conformers...')
        if self.multi_process:
            pool = Pool(processes=min(8, os.cpu_count()))
            inputs = [item for item in tqdm(pool.imap(self.single_process, smiles_list))]
            pool.close()
        else:

```

--------------------------------

### Generate Uni-Mol Molecule and Atomic Representations

Source: https://unimol.readthedocs.io/en/latest/_sources/quickstart

This code demonstrates how to obtain Uni-Mol representations for molecules and atoms. It allows for specifying whether to remove hydrogens and choosing between UniMol V1 and V2 models with different sizes. The output includes CLS token representations and atomic-level representations.

```python
import numpy as np
from unimol_tools import UniMolRepr
# single smiles unimol representation
clf = UniMolRepr(data_type='molecule', 
                 remove_hs=False,
                 model_name='unimolv1', # avaliable: unimolv1, unimolv2
                 model_size='84m', # work when model_name is unimolv2. avaliable: 84m, 164m, 310m, 570m, 1.1B.
                 )
smiles = 'c1ccc(cc1)C2=NCC(=O)Nc3c2cc(cc3)[N+](=O)[O]'
smiles_list = [smiles]
unimol_repr = clf.get_repr(smiles_list, return_atomic_reprs=True)
# CLS token repr
print(np.array(unimol_repr['cls_repr']).shape)
# atomic level repr, align with rdkit mol.GetAtoms()
print(np.array(unimol_repr['atomic_reprs']).shape)
```

--------------------------------

### MolTrain Initialization and Configuration - Python

Source: https://unimol.readthedocs.io/en/latest/_modules/unimol_tools/train

Illustrates the initialization of the MolTrain class and how configuration parameters are read and applied. It shows the loading of configuration from a file or default settings, and the subsequent update of these configurations based on provided arguments.

```python
import os
import logging
from unimol_tools.util import YamlHandler

logger = logging.getLogger(__name__)

# Assume class structure and YamlHandler are defined elsewhere
class MolTrain:
    def __init__(self, task, data_type, epochs, learning_rate, batch_size, early_stopping, metrics, split, split_group_col, kfold, remove_hs, smiles_col, target_cols, target_col_prefix, target_anomaly_check, smiles_check, target_normalize, max_norm, use_cuda, use_amp, freeze_layers, freeze_layers_reversed, load_model_dir=None, model_name='unimolv1', model_size='84m', save_path=None):
        if load_model_dir is not None:
            config_path = os.path.join(load_model_dir, 'config.yaml')
            logger.info('Load config file from {}'.format(config_path))
        else:
            config_path = os.path.join(os.path.dirname(__file__), 'config/default.yaml')
        self.yamlhandler = YamlHandler(config_path)
        config = self.yamlhandler.read_yaml()
        config.task = task
        config.data_type = data_type
        config.epochs = epochs
        config.learning_rate = learning_rate
        config.batch_size = batch_size
        config.patience = early_stopping
        config.metrics = metrics
        config.split = split
        config.split_group_col = split_group_col
        config.kfold = kfold
        config.remove_hs = remove_hs
        config.smiles_col = smiles_col
        config.target_cols = target_cols
        config.target_col_prefix = target_col_prefix
        config.anomaly_clean = target_anomaly_check in ['filter']
        config.smi_strict = smiles_check in ['filter']
        config.target_normalize = target_normalize
        config.max_norm = max_norm
        config.use_cuda = use_cuda
        config.use_amp = use_amp
        config.freeze_layers = freeze_layers
        config.freeze_layers_reversed = freeze_layers_reversed
        config.load_model_dir = load_model_dir
        config.model_name = model_name
        config.model_size = model_size
        self.save_path = save_path
        self.config = config

```

--------------------------------

### Get Activation Function Utility

Source: https://unimol.readthedocs.io/en/latest/models

A utility function to retrieve the appropriate activation function based on a string identifier. This simplifies the process of configuring different activation functions within neural network architectures.

```python
def get_activation_fn(_activation_):
    """
    Returns the activation function corresponding to activation 
    """
    pass
```

--------------------------------

### Safe Indexing in a List (Python)

Source: https://unimol.readthedocs.io/en/latest/_modules/unimol_tools/data/conformer

Provides a safe way to get the index of an element in a list. If the element is not found, it returns the index of the last element in the list, preventing 'ValueError' exceptions.

```python
def safe_index(l, e):
    """
    Return index of element e in list l. If e is not present, return the last index
    """
    try:
        return l.index(e)
    except:
        return len(l) - 1
```

--------------------------------

### Download All Uni-Mol Pretrained Model Weights

Source: https://unimol.readthedocs.io/en/latest/weight

Downloads all available pretrained model weights to the WEIGHT_DIR. It includes an option to use symlinks for the local directory, defaulting to False.

```python
unimol_tools.weights.weighthub.download_all_weights(_local_dir_use_symlinks =False_)

```

--------------------------------

### Download All UniMol Model Weights (Python)

Source: https://unimol.readthedocs.io/en/latest/_modules/unimol_tools/weights/weighthub

Downloads all available pretrained model weights from the 'dptech/Uni-Mol-Models' repository to the `WEIGHT_DIR`. It uses `snapshot_download` with `allow_patterns='*'` to fetch all files.

```Python
import os

from ..utils import logger

try:
    from huggingface_hub import snapshot_download
except:
    huggingface_hub_installed = False
    def snapshot_download(*args, **kwargs):
        raise ImportError('huggingface_hub is not installed. If weights are not avaliable, please install it by running: pip install huggingface_hub. Otherwise, please download the weights manually from https://huggingface.co/dptech/Uni-Mol-Models')

WEIGHT_DIR = os.environ.get('UNIMOL_WEIGHT_DIR', os.path.dirname(os.path.abspath(__file__))) 

if 'UNIMOL_WEIGHT_DIR' in os.environ:
    logger.warning(f'Using custom weight directory from UNIMOL_WEIGHT_DIR: {WEIGHT_DIR}')
else:
    logger.info(f'Weights will be downloaded to default directory: {WEIGHT_DIR}')

os.environ["HF_ENDPOINT"] = "https://hf-mirror.com" # use mirror to download weights


def download_all_weights(local_dir_use_symlinks=False):
    """
    Downloads all available pretrained model weights to the WEIGHT_DIR.
    
    :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(
        repo_id="dptech/Uni-Mol-Models",
        local_dir=WEIGHT_DIR,
        allow_patterns='*',
        local_dir_use_symlinks=local_dir_use_symlinks,
        #max_workers=8
    )


if '__main__' == __name__:
    download_all_weights()
```

--------------------------------

### Get Uni-Mol Representations

Source: https://unimol.readthedocs.io/en/latest/train

Illustrates how to obtain molecular representations using the UniMolRepr class. This class supports different data types like molecules and OLEDs, and allows for custom model loading and GPU usage.

```python
from unimol_tools.predictor import UniMolRepr

# Example with SMILES string
smiles_data = 'CCO'
repr_generator = UniMolRepr(model_name='unimolv1', use_gpu=True)
mol_repr = repr_generator.get_repr(data=smiles_data)

# Example with custom conformers
custom_conformer_data = {
    'atoms': ['C', 'O'],
    'coordinates': [[0, 0, 0], [0.5, 0, 0]]
}
mol_repr_conformer = repr_generator.get_repr(data=custom_conformer_data, return_atomic_reprs=True)
```

--------------------------------

### Initialize DataHub with ML Task Configuration

Source: https://unimol.readthedocs.io/en/latest/_modules/unimol_tools/data/datahub

Initializes the DataHub for machine learning tasks, setting up data handling and preprocessing based on provided parameters. It configures target scaling and data initialization according to the specified task type.

```python
class DataHub(object):
    """
    The DataHub class is responsible for storing and preprocessing data for machine learning tasks.
    It initializes with configuration options to handle different types of tasks such as regression, 
    classification, and others. It also supports data scaling and handling molecular data.
    """


    def __init__(self, data=None, is_train=True, save_path=None, **params):
        """
        Initializes the DataHub instance with data and configuration for the ML task.

        :param data: Initial dataset to be processed.
        :param is_train: (bool) Indicates if the DataHub is being used for training.
        :param save_path: (str) Path to save any necessary files, like scalers.
        :param params: Additional parameters for data preprocessing and model configuration.
        """
        self.data = data
        self.is_train = is_train
        self.save_path = save_path
        self.task = params.get('task', None)
        self.target_cols = params.get('target_cols', None)
        self.multiclass_cnt = params.get('multiclass_cnt', None)
        self.ss_method = params.get('target_normalize', 'none')
        self._init_data(**params)
        self._init_split(**params)
```

--------------------------------

### Unimol Predict Module Source Code

Source: https://unimol.readthedocs.io/en/latest/_modules/unimol_tools/predict

This snippet displays the source code for the unimol_tools.predict module. It includes the copyright notice and the MIT license information. No specific function definitions or usage examples are provided in this excerpt.

```python
# Copyright (c) DP Technology.
# This source code is licensed under the MIT license found in the

```

--------------------------------

### Get Activation Function Utility (PyTorch)

Source: https://unimol.readthedocs.io/en/latest/_modules/unimol_tools/models/transformers

A utility function that returns the corresponding PyTorch functional activation function based on a string identifier. It supports 'relu', 'gelu', 'tanh', and 'linear' activations, raising an error for unsupported types.

```python
import torch
import torch.nn.functional as F



[docs]
def get_activation_fn(activation):
    """ Returns the activation function corresponding to `activation` """

    if activation == "relu":
        return F.relu
    elif activation == "gelu":
        return F.gelu
    elif activation == "tanh":
        return torch.tanh
    elif activation == "linear":
        return lambda x: x
    else:
        raise RuntimeError("--activation-fn {} not supported".format(activation))
```

--------------------------------

### Weight Download Functions

Source: https://unimol.readthedocs.io/en/latest/weight

Functions for downloading Uni-Mol pretrained model weights.

```APIDOC
## unimol_tools.weights.weighthub.weight_download

### Description
Downloads the specified pretrained model weights.

### Method
Python Function Call

### Endpoint
N/A

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
*   **pretrain** (str) - Required - The name of the pretrained model to download.
*   **save_path** (str) - Required - 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.

### Request Example
```python
weight_download(pretrain='dptech/Uni-Mol-Models', save_path='./models')
```

### Response
#### Success Response (200)
N/A (This is a Python function)

#### Response Example
N/A
```

```APIDOC
## unimol_tools.weights.weighthub.weight_download_v2

### Description
Downloads the specified pretrained model weights.

### Method
Python Function Call

### Endpoint
N/A

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
*   **pretrain** (str) - Required - The name of the pretrained model to download.
*   **save_path** (str) - Required - 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.

### Request Example
```python
weight_download_v2(pretrain='dptech/Uni-Mol-Models', save_path='./models')
```

### Response
#### Success Response (200)
N/A (This is a Python function)

#### Response Example
N/A
```

```APIDOC
## unimol_tools.weights.weighthub.download_all_weights

### Description
Downloads all available pretrained model weights to the WEIGHT_DIR.

### Method
Python Function Call

### Endpoint
N/A

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
*   **local_dir_use_symlinks** (bool) - Optional - Whether to use symlinks for the local directory. Defaults to False.

### Request Example
```python
download_all_weights()
```

### Response
#### Success Response (200)
N/A (This is a Python function)

#### Response Example
N/A
```

--------------------------------

### Initialize MolPredict Class

Source: https://unimol.readthedocs.io/en/latest/_modules/unimol_tools/predict

Initializes the MolPredict class by loading a pre-trained model and its configuration. It requires the path to the model directory, which should contain a 'config.yaml' file. The configuration is parsed to set up prediction parameters such as target columns and the prediction task.

```python
class MolPredict(object):
    """A :class:`MolPredict` class is responsible for interface of predicting process of molecular data."""


    def __init__(self, load_model=None):
        """ 
        Initialize a :class:`MolPredict` class.

        :param load_model: str, default=None, path of model to load.
        """
        if not load_model:
            raise ValueError("load_model is empty")
        self.load_model = load_model
        config_path = os.path.join(load_model, 'config.yaml')
        self.config = YamlHandler(config_path).read_yaml()
        self.config.target_cols = self.config.target_cols.split(',')
        self.task = self.config.task
        self.target_cols = self.config.target_cols

```

--------------------------------

### UniMolModel Initialization (Python)

Source: https://unimol.readthedocs.io/en/latest/_modules/unimol_tools/models/unimol

Initializes the UniMolModel, dynamically configuring its architecture based on the specified data type (e.g., 'molecule', 'protein'). It handles loading pretrained weights and dictionary files, and sets up embedding and encoder layers. Dependencies include PyTorch, custom utilities for padding, and configuration/data modules.

```python
class UniMolModel(nn.Module):
    """
    UniMolModel is a specialized model for molecular, protein, crystal, or MOF (Metal-Organic Frameworks) data. 
    It dynamically configures its architecture based on the type of data it is intended to work with. The model
    supports multiple data types and incorporates various architecture configurations and pretrained weights.

    Attributes:
        - output_dim: The dimension of the output layer.
        - data_type: The type of data the model is designed to handle.
        - remove_hs: Flag to indicate whether hydrogen atoms are removed in molecular data.
        - pretrain_path: Path to the pretrained model weights.
        - dictionary: The dictionary object used for tokenization and encoding.
        - mask_idx: Index of the mask token in the dictionary.
        - padding_idx: Index of the padding token in the dictionary.
        - embed_tokens: Embedding layer for token embeddings.
        - encoder: Transformer encoder backbone of the model.
        - gbf_proj, gbf: Layers for Gaussian basis functions or numerical embeddings.
        - classification_head: The final classification head of the model.
    """


    def __init__(self, output_dim=2, data_type='molecule', **params):
        """
        Initializes the UniMolModel with specified parameters and data type.

        :param output_dim: (int) The number of output dimensions (classes).
        :param data_type: (str) The type of data (e.g., 'molecule', 'protein').
        :param params: Additional parameters for model configuration.
        """
        super().__init__()
        if data_type == 'molecule':
            self.args = molecule_architecture()
        elif data_type == 'oled':
            self.args = oled_architecture()
        elif data_type == 'protein':
            self.args = protein_architecture()
        elif data_type == 'crystal':
            self.args = crystal_architecture()
        else:
            raise ValueError('Current not support data type: {}'.format(data_type))
        self.output_dim = output_dim
        self.data_type = data_type
        self.remove_hs = params.get('remove_hs', False)
        if data_type == 'molecule':
            name = "no_h" if self.remove_hs else "all_h" 
            name = data_type + '_' + name
        else:
            name = data_type
        if not os.path.exists(os.path.join(WEIGHT_DIR, MODEL_CONFIG['weight'][name])):
            weight_download(MODEL_CONFIG['weight'][name], WEIGHT_DIR)
        if not os.path.exists(os.path.join(WEIGHT_DIR, MODEL_CONFIG['dict'][name])):
            weight_download(MODEL_CONFIG['dict'][name], WEIGHT_DIR)
        self.pretrain_path = os.path.join(WEIGHT_DIR, MODEL_CONFIG['weight'][name])
        self.dictionary = Dictionary.load(os.path.join(WEIGHT_DIR, MODEL_CONFIG['dict'][name]))
        self.mask_idx = self.dictionary.add_symbol("[MASK]", is_special=True)
        self.padding_idx = self.dictionary.pad()
        self.embed_tokens = nn.Embedding(
            len(self.dictionary), self.args.encoder_embed_dim, self.padding_idx
        )
        self.encoder = BACKBONE[self.args.backbone](
            encoder_layers=self.args.encoder_layers,
            embed_dim=self.args.encoder_embed_dim,
            ffn_embed_dim=self.args.encoder_ffn_embed_dim,
            attention_heads=self.args.encoder_attention_heads,
            emb_dropout=self.args.emb_dropout,
            dropout=self.args.dropout,
            attention_dropout=self.args.attention_dropout,
            activation_dropout=self.args.activation_dropout,
            max_seq_len=self.args.max_seq_len,
            activation_fn=self.args.activation_fn,
            no_final_head_layer_norm=self.args.delta_pair_repr_loss < 0,
        )
        K = 128
        n_edge_type = len(self.dictionary) * len(self.dictionary)
        self.gbf_proj = NonLinearHead(
            K, self.args.encoder_attention_heads, self.args.activation_fn
        )
        if self.args.kernel == 'gaussian':
            self.gbf = GaussianLayer(K, n_edge_type)
        else:
            self.gbf = NumericalEmbed(K, n_edge_type)
        self.classification_head = ClassificationHead(
            input_dim=self.args.encoder_embed_dim,
            inner_dim=self.args.encoder_embed_dim,
            num_classes=self.output_dim,
            activation_fn=self.args.pooler_activation_fn,
            pooler_dropout=self.args.pooler_dropout,
        )
```

--------------------------------

### Initialize and Apply K-Fold Splitter in Python

Source: https://unimol.readthedocs.io/en/latest/_modules/unimol_tools/data/datahub

Initializes a data splitter with K-Fold cross-validation and applies it to the dataset. It handles cases where kfold is 1 (all data for training) and logs the splitting process. The output is a list of tuples, where each tuple contains training and testing indices for a fold.

```python
self.kfold = params.get('kfold', kfold)
        self.method = params.get('split', method)
        self.split_seed = params.get('split_seed', 42)
        self.data['kfold'] = self.kfold
        if not self.is_train:
            return
        self.splitter = Splitter(self.method, self.kfold, seed=self.split_seed)
        split_nfolds = self.splitter.split(**self.data)
        if self.kfold == 1:
            logger.info(f"Kfold is 1, all data is used for training.")
        else:
            logger.info(f"Split method: {self.method}, fold: {self.kfold}")
        nfolds = np.zeros(len(split_nfolds[0][0])+len(split_nfolds[0][1]), dtype=int)
        for enu, (tr_idx, te_idx) in enumerate(split_nfolds):
            nfolds[te_idx] = enu
        self.data['split_nfolds'] = split_nfolds
        return split_nfolds
```

--------------------------------

### Focal Loss Calculation (PyTorch)

Source: https://unimol.readthedocs.io/en/latest/_modules/unimol_tools/models/loss

Implements the Focal Loss function, designed to address class imbalance by down-weighting easy examples and focusing training on hard negatives. It takes predicted probabilities and ground truth labels, along with optional alpha and gamma parameters for balancing and focusing, respectively. The function returns the computed focal loss averaged over the batch.

```python
def FocalLoss(y_pred, y_true, alpha=0.25, gamma=2):
    """
    Calculates the Focal Loss, used to address class imbalance by focusing on hard examples.

    :param y_pred: (torch.Tensor) Predicted probabilities.
    :param y_true: (torch.Tensor) Ground truth labels.
    :param alpha: (float) Weighting factor for balancing positive and negative examples. Defaults to 0.25.
    :param gamma: (float) Focusing parameter to scale the loss. Defaults to 2.

    :return: (torch.Tensor) Computed focal loss.
    """
    if y_pred.shape != y_true.shape:
        y_true = y_true.flatten()
    y_true = y_true.long()
    y_pred = y_pred.float()
    y_true = y_true.float()
    y_true = y_true.unsqueeze(1)
    y_pred = y_pred.unsqueeze(1)
    y_true = torch.cat((1-y_true, y_true), dim=1)
    y_pred = torch.cat((1-y_pred, y_pred), dim=1)
    y_pred = y_pred.clamp(1e-5, 1.0)
    loss = -alpha * y_true * torch.pow((1 - y_pred), gamma) * torch.log(y_pred)
    return torch.mean(torch.sum(loss, dim=1))
```

--------------------------------

### Config API

Source: https://unimol.readthedocs.io/en/latest/utils

Handles reading and writing YAML configuration files, with utilities for converting between dict and addict objects.

```APIDOC
## unimol_tools.utils.config_handler.YamlHandler

### Description
A class to read and write YAML files.

### Parameters
- **file_path** (str) - The path to the YAML configuration file.

### Methods
#### read_yaml

### Description
Reads a YAML file and converts its content to an addict dictionary.

### Parameters
- **encoding** (str) - The encoding method for reading the file, defaults to 'utf-8'.

### Returns
- Dict (addict) - The content of the YAML file as an addict dictionary.

#### write_yaml

### Description
Writes data to a YAML file.

### Parameters
- **data** (dict or Dict(addict)) - The data to write to the YAML file.
- **out_file_path** (str) - The path to write the YAML file to.
- **encoding** (str) - The encoding method for writing the file, defaults to 'utf-8'.

## unimol_tools.utils.config_handler.addict2dict

### Description
Converts an addict object to a standard Python dictionary.

### Parameters
- **addict_obj** (Dict(addict)) - The addict object to convert.

### Returns
- Dict - The converted dictionary.
```

--------------------------------

### Apply Custom Loss Functions (GHM, Focal, MAE)

Source: https://context7.com/context7/unimol_readthedocs_io_en/llms.txt

Applies specialized loss functions including Gradient Harmonizing Mechanism (GHM) loss for imbalanced datasets, Focal loss for hard example mining, and MAE with NaN support. Uses classes from unimol_tools.models.loss. Requires PyTorch. Inputs are model predictions and targets; outputs are calculated loss values.

```python
from unimol_tools.models.loss import GHMC_Loss, GHMR_Loss, FocalLossWithLogits, MAEwithNan
import torch

# GHM loss for classification (handles imbalanced data)
ghm_loss = GHMC_Loss(bins=10, alpha=0.5)

logits = torch.randn(32, 1)  # (batch_size, 1)
targets = torch.randint(0, 2, (32, 1)).float()
loss = ghm_loss(logits, targets)
print(f"GHM Classification loss: {loss.item()}")

# GHM loss for regression
ghmr_loss = GHMR_Loss(bins=10, alpha=0.5, mu=0.02)

pred = torch.randn(32, 1)
target = torch.randn(32, 1)
loss = ghmr_loss(pred, target)
print(f"GHM Regression loss: {loss.item()}")

# Focal loss for classification
focal_loss = FocalLossWithLogits

logits = torch.randn(32, 1)
targets = torch.randint(0, 2, (32, 1)).float()
loss = focal_loss(logits, targets, alpha=0.25, gamma=2.0)
print(f"Focal loss: {loss.item()}")

# MAE with NaN support
pred = torch.tensor([[1.5], [2.3], [float('nan')], [4.5]])
target = torch.tensor([[1.4], [2.5], [3.0], [4.3]])
loss = MAEwithNan(pred, target)
print(f"MAE (ignoring NaN): {loss.item()}")
```

--------------------------------

### Build UniMolModel Instance

Source: https://unimol.readthedocs.io/en/latest/_modules/unimol_tools/models/unimol

A class method to construct a new UniMolModel. It takes model configuration arguments and returns an initialized instance of the UniMolModel class.

```python
@classmethod
def build_model(cls, args):
        """
        Class method to build a new instance of the UniMolModel.

        :param args: Arguments for model configuration.
        :return: An instance of UniMolModel.
        """
        return cls(args)
```

--------------------------------

### Initialize Trainer Parameters

Source: https://unimol.readthedocs.io/en/latest/_modules/unimol_tools/tasks/trainer

Initializes common and Neural Network-specific parameters for the trainer, including data splitting, seeding, logging level, learning rate, batch size, epochs, warmup ratio, patience, max norm, CUDA usage, automatic mixed precision (AMP), and device selection. It also initializes a gradient scaler if AMP and CUDA are enabled.

```python
def _init_trainer(self, **params):
    """
    Initializing the trainer class to train model.

    :param params: Containing training arguments.
    """
    ### init common params ###
    self.split_method = params.get('split_method', '5fold_random')
    self.split_seed = params.get('split_seed', 42)
    self.seed = params.get('seed', 42)
    self.set_seed(self.seed)
    self.logger_level = int(params.get('logger_level', 1))
    ### init NN trainer params ###
    self.learning_rate = float(params.get('learning_rate', 1e-4))
    self.batch_size = params.get('batch_size', 32)
    self.max_epochs = params.get('epochs', 50)
    self.warmup_ratio = params.get('warmup_ratio', 0.1)
    self.patience = params.get('patience', 10)
    self.max_norm = params.get('max_norm', 1.0)
    self.cuda = params.get('cuda', False)
    self.amp = params.get('amp', False)
    self.device = torch.device(
        "cuda:0" if torch.cuda.is_available() and self.cuda else "cpu")
    self.scaler = torch.cuda.amp.GradScaler(
    ) if self.device.type == 'cuda' and self.amp == True else None
```

--------------------------------

### Dataset Handling API

Source: https://unimol.readthedocs.io/en/latest/models

Documentation for NNDataset and TorchDataset for handling data in Uni-Mol.

```APIDOC
## `NNDataset()`

### Description
A dataset class for neural network operations.
```

```APIDOC
## `TorchDataset` Class

### Description
Dataset implementation specifically for PyTorch.

### Methods

- **`TorchDataset.__init__(...)`**
  - Initializes the TorchDataset.
```

--------------------------------

### get_linear_schedule_with_warmup - Learning Rate Scheduler

Source: https://unimol.readthedocs.io/en/latest/task

Creates a schedule with a learning rate that decreases linearly after a warmup period.

```APIDOC
## POST /get_linear_schedule_with_warmup

### Description
Create a schedule with a learning rate that decreases linearly from the initial lr set in the optimizer to 0, after a warmup period during which it increases linearly from 0 to the initial lr set in the optimizer.

### Method
POST

### Endpoint
/get_linear_schedule_with_warmup

### Parameters
#### Request Body
- **optimizer** (object) - The optimizer for which to schedule the learning rate.
- **num_warmup_steps** (int) - The number of steps for the warmup phase.
- **num_training_steps** (int) - The total number of training steps.
- **last_epoch** (int) - Optional. The index of the last epoch when resuming training. Defaults to -1.

### Request Example
```json
{
  "optimizer": "your_optimizer_object",
  "num_warmup_steps": 100,
  "num_training_steps": 1000
}
```

### Response
#### Success Response (200)
- **scheduler** (object) - torch.optim.lr_scheduler.LambdaLR with the appropriate schedule.

#### Response Example
```json
{
  "scheduler": "your_lr_scheduler_object"
}
```
```

--------------------------------

### Transformer Module Utilities

Source: https://unimol.readthedocs.io/en/latest/models

Details on helper functions for activation and dropout within the transformer module.

```APIDOC
## unimol_tools.models.transformers.softmax_dropout

### Description
Applies softmax dropout to an input tensor, with optional mask and bias.

### Method
`softmax_dropout(_input_, _dropout_prob_, _is_training=True_, _mask=None_, _bias=None_, _inplace=True_)`

### Parameters
#### Arguments
- **input** (torch.Tensor) - The input tensor.
- **dropout_prob** (float) - The dropout probability.
- **is_training** (bool, optional) - Whether the module is in training mode. Defaults to True.
- **mask** (torch.Tensor, optional) - An optional mask tensor to be added to the input.
- **bias** (torch.Tensor, optional) - An optional bias tensor to be added to the input.
- **inplace** (bool, optional) - If True, performs the operation in-place. Defaults to True.

### Returns
- **torch.Tensor**: The result after applying softmax dropout.

## unimol_tools.models.transformers.get_activation_fn

### Description
Returns the appropriate activation function based on the provided name.

### Method
`get_activation_fn(_activation_)`

### Parameters
#### Arguments
- **activation** (str) - The name of the activation function (e.g., 'gelu').

### Returns
- A callable activation function.
```

--------------------------------

### Download Specific Uni-Mol Model Weights

Source: https://unimol.readthedocs.io/en/latest/weight

Downloads specified pretrained model weights using the weight_download function. It takes the model name, save path, and an option for using local directory symlinks as parameters.

```python
unimol_tools.weights.weighthub.weight_download(_pretrain_ , _save_path_ , _local_dir_use_symlinks =True_)

```

--------------------------------

### Initialize MolTrain for Molecular Data Training

Source: https://unimol.readthedocs.io/en/latest/train

Instantiates the MolTrain class for training molecular datasets. This class handles various training configurations such as task type, data preprocessing, optimization parameters, and model saving. It supports different data types and metrics.

```python
from unimol_tools import MolTrain
import numpy as np

# Example initialization for classification task
train = MolTrain(
    task='classification',
    data_type='molecule',
    epochs=10,
    learning_rate=0.0001,
    batch_size=16,
    early_stopping=5,
    metrics='auc',
    split='random',
    save_path='./exp',
    remove_hs=False,
    smiles_col='SMILES',
    target_cols=['TARGET_A', 'TARGET_B'],
    use_cuda=True,
    use_amp=True,
    model_name='unimolv1',
    model_size='84m'
)
```

--------------------------------

### Import Core Libraries for PyTorch Training

Source: https://unimol.readthedocs.io/en/latest/_modules/unimol_tools/tasks/trainer

Imports essential libraries for building and training PyTorch models. This includes numerical computation with NumPy, tensor operations with PyTorch, data loading utilities, optimizers like Adam, learning rate scheduling, and gradient clipping.

```python
# Copyright (c) DP Technology.
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

from __future__ import absolute_import, division, print_function

import os
import numpy as np
import torch
from torch.utils.data import DataLoader as TorchDataLoader
from torch.optim import Adam
from torch.optim.lr_scheduler import LambdaLR
from functools import partial
from torch.nn.utils import clip_grad_norm_
```

--------------------------------

### Metrics Class Initialization

Source: https://unimol.readthedocs.io/en/latest/_modules/unimol_tools/utils/metrics

Initializes the Metrics class with a specified task type and optionally a comma-separated string of metric names. It sets up the task, defines thresholds for calculations, and initializes the metric dictionary based on provided or default metrics.

```python
class Metrics(object):
    """
    Class for calculating metrics for different tasks.

    :param task: The task type. Supported tasks are 'regression', 'multilabel_regression',
                 'classification', 'multilabel_classification', and 'multiclass'.
    :param metrics_str: Comma-separated string of metric names. If provided, only the specified metrics will be calculated. If not provided or an empty string, default metrics for the task will be used.
    """



    def __init__(self, task=None, metrics_str=None, **params):
        self.task = task
        self.threshold = np.arange(0, 1., 0.1)
        self.metric_dict = self._init_metrics(self.task, metrics_str, **params)
        self.METRICS_REGISTER = METRICS_REGISTER[task]
```

--------------------------------

### Download Specific UniMol Model Weights (Python)

Source: https://unimol.readthedocs.io/en/latest/_modules/unimol_tools/weights/weighthub

Downloads weights for a specific UniMol pretrained model. It checks if the model already exists locally before initiating the download using `snapshot_download` from `huggingface_hub`. The weights are saved to a specified path.

```Python
import os

from ..utils import logger

try:
    from huggingface_hub import snapshot_download
except:
    huggingface_hub_installed = False
    def snapshot_download(*args, **kwargs):
        raise ImportError('huggingface_hub is not installed. If weights are not avaliable, please install it by running: pip install huggingface_hub. Otherwise, please download the weights manually from https://huggingface.co/dptech/Uni-Mol-Models')

WEIGHT_DIR = os.environ.get('UNIMOL_WEIGHT_DIR', os.path.dirname(os.path.abspath(__file__))) 

if 'UNIMOL_WEIGHT_DIR' in os.environ:
    logger.warning(f'Using custom weight directory from UNIMOL_WEIGHT_DIR: {WEIGHT_DIR}')
else:
    logger.info(f'Weights will be downloaded to default directory: {WEIGHT_DIR}')

os.environ["HF_ENDPOINT"] = "https://hf-mirror.com" # use mirror to download weights


def weight_download(pretrain, save_path, local_dir_use_symlinks=True):
    """
    Downloads the specified pretrained model weights.

    :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}')
        return
    
    logger.info(f'Downloading {pretrain}')
    snapshot_download(
        repo_id="dptech/Uni-Mol-Models",
        local_dir=save_path,
        allow_patterns=pretrain,
        local_dir_use_symlinks=local_dir_use_symlinks,
        #max_workers=8
    )
```

--------------------------------

### Transformer Utilities API

Source: https://unimol.readthedocs.io/en/latest/models

Documentation for utility functions and layers related to Transformers.

```APIDOC
## `softmax_dropout()` Function

### Description
Applies softmax followed by dropout.
```

```APIDOC
## `get_activation_fn()` Function

### Description
Retrieves an activation function based on its name.
```

```APIDOC
## `SelfMultiheadAttention` Class

### Description
Implements self multi-head attention mechanism.

### Methods

- **`SelfMultiheadAttention.__init__(...)`**
  - Initializes the SelfMultiheadAttention layer.
```

```APIDOC
## `TransformerEncoderLayer` Class

### Description
Represents a single layer in a Transformer encoder.

### Methods

- **`TransformerEncoderLayer.__init__(...)`**
  - Initializes the TransformerEncoderLayer.

- **`TransformerEncoderLayer.forward(...)`**
  - Performs the forward pass of the encoder layer.
```

```APIDOC
## `TransformerEncoderWithPair` Class

### Description
Transformer encoder layer that considers pair-wise interactions.

### Methods

- **`TransformerEncoderWithPair.__init__(...)`**
  - Initializes the TransformerEncoderWithPair.

- **`TransformerEncoderWithPair.forward(...)`**
  - Performs the forward pass of the encoder considering pair-wise information.
```

--------------------------------

### UniMolModel API

Source: https://unimol.readthedocs.io/en/latest/models

Details the UniMolModel class and its methods for building and loading models.

```APIDOC
## `UniMolModel` Class

### Description
Represents the main Uni-Mol model, handling initialization, loading weights, and forward passes.

### Methods

- **`UniMolModel.__init__(...)`**
  - Initializes the UniMolModel.

- **`UniMolModel.load_pretrained_weights(...)`**
  - Loads pre-trained weights into the model.

- **`UniMolModel.build_model(...)`**
  - Constructs the model architecture.

- **`UniMolModel.forward(...)`**
  - Performs the forward pass of the model.

- **`UniMolModel.batch_collate_fn(...)`**
  - Collates batches for model input.
```

--------------------------------

### Set Hugging Face Hub Mirror Endpoint

Source: https://unimol.readthedocs.io/en/latest/weight

Configures the HF_ENDPOINT environment variable to use a specified mirror for downloading models from the Hugging Face Hub. This is useful if the default download is slow.

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

```

--------------------------------

### Reading Molecular Data with MolDataReader (Python)

Source: https://unimol.readthedocs.io/en/latest/_modules/unimol_tools/data/datareader

This snippet demonstrates how to initialize and use the MolDataReader class to read molecular data. It covers loading data from different input types (file path, dictionary, list/numpy array) and setting up parameters for data processing, such as target columns and SMILES column names.

```python
from unimol.data.loader import MolDataReader

# Example 1: Reading from a file path
data_reader = MolDataReader()
data_dict_from_file = data_reader.read_data(data='/path/to/your/data.csv', smiles_col='SMILES', target_cols=['target1', 'target2'])

# Example 2: Reading from a dictionary
my_data_dict = {'SMILES': ['CCO', 'CCC'], 'target1': [0.1, 0.2], 'target2': [1, 0]}
data_dict_from_dict = data_reader.read_data(data=my_data_dict)

# Example 3: Reading from a list of SMILES strings
smiles_list = ['CCO', 'CCC']
data_dict_from_list = data_reader.read_data(data=smiles_list)

# Example 4: Reading for prediction (target placeholder)
data_dict_predict = data_reader.read_data(data=smiles_list, is_train=False)

print(data_dict_from_file.keys())
print(data_dict_from_dict.keys())
print(data_dict_from_list.keys())
print(data_dict_predict.keys())
```
Tokens:
Raw