### 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()) ```