### Install unimol_tools from source

Source: https://github.com/deepmodeling/unimol_tools/blob/main/docs/source/installation.md

Installs unimol_tools from source by installing dependencies, cloning the repository, and running the setup script.

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

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

## Install
python setup.py install
```

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

### Install unimol_tools from PyPI

Source: https://github.com/deepmodeling/unimol_tools/blob/main/docs/source/installation.md

Installs the unimol_tools package from PyPI. This is the recommended installation method.

```bash
pip install unimol_tools
```

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

### Install huggingface_hub

Source: https://github.com/deepmodeling/unimol_tools/blob/main/docs/source/weight.rst

Install the huggingface_hub package, which is required for automatic downloading of Uni-Mol models at runtime.

```bash
pip install huggingface_hub
```

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

### Classification example with MolTrain and MolPredict

Source: https://github.com/deepmodeling/unimol_tools/blob/main/docs/source/quickstart.md

Shows how to train a classification model and make predictions using unimol_tools. The default smiles column is 'SMILES' and the target column is 'target', but they can be customized with smiles_col and target_cols. The example loads the Ames dataset, drops the 'CAS_NO' column, renames 'Activity' to 'target', splits into train/test, trains with MolTrain, and predicts with MolPredict.

```Python
import pandas as pd
from unimol_tools import MolTrain, MolPredict

# Load the dataset
df = pd.read_csv('../datasets/Ames/Ames.csv')

df = df.drop(columns=['CAS_NO']).rename(columns={'Activity': 'target'})

# Divide the training set and test set
train_df = df.sample(frac=0.8, random_state=42)
test_df = df.drop(train_df.index)

train_df_dict = train_df.to_dict(orient='list')
test_df_dict = test_df.to_dict(orient='list')

clf = MolTrain(task='classification', 
                data_type='molecule', 
                epochs=20, 
                batch_size=16, 
                metrics='auc',
                smiles_col='Canonical_Smiles',
                )

clf.fit(train_df_dict)

predictor = MolPredict(load_model='./exp')

pred = predictor.predict(test_df_dict['Canonical_Smiles'])
```

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

### Train and predict with MolTrain and MolPredict

Source: https://github.com/deepmodeling/unimol_tools/blob/main/docs/source/quickstart.md

Example of training a classification model with MolTrain and predicting with MolPredict. Set model_name='unimolv2' for UniMol V2. Supports SMILES-based CSV/TXT files. The model is saved to './exp' and loaded for prediction.

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

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

### Molecule property prediction with MolTrain and MolPredict

Source: https://github.com/deepmodeling/unimol_tools/blob/main/README.md

Trains a classification model on molecular data (SMILES or SDF) using MolTrain, then predicts with MolPredict. The training data can be a CSV/TXT file with SMILES, an SDF file, or a custom dictionary with 'atoms' and 'coordinates'. Pretrained weights are downloaded automatically if paths are not provided. The example shows commented-out code for fitting with a dictionary and for specifying custom pretrained paths.

```python
from unimol_tools import MolTrain, MolPredict
clf = MolTrain(
    task='classification',
    data_type='molecule',
    epochs=10,
    batch_size=16,
    metrics='auc',
    # pretrained weights are downloaded automatically when left as ``None``
    # pretrained_model_path='/path/to/checkpoint.ckpt',
    # pretrained_dict_path='/path/to/dict.txt',
)
clf.fit(data = train_data)
# currently support data with smiles based csv/txt file, and sdf file with mol,
# and custom dict of {'atoms':[['C','C'],['C','H','O']], 'coordinates':[coordinates_1,coordinates_2]}

# The dict format can refer to the following format, or be obtained from sdf, 
# which can also be directly input into the model.
train_sdf = PandasTools.LoadSDF('exp/unimol_conformers_train.sdf')
train_dict = {
    'atoms': [list(atom.GetSymbol() for atom in mol.GetAtoms()) for mol in train_sdf['ROMol']],
    # atoms[0]: ['C', 'C', 'O', 'C', 'O', 'C', ...]
    'coordinates': [mol.GetConformers()[0].GetPositions() for mol in train_sdf['ROMol']],
    # coordinates[0]: array([[ 6.6462, -1.8268,  1.9275],
    #                        [ 6.1552, -1.9367,  0.4873],
    #                        [ 5.1832, -0.8757,  0.3007],
    #                        [ 5.4651, -0.0272, -0.7266],
    #                        [ 4.8586, -0.0844, -1.7917],
    #                        [ 6.5362,  0.9767, -0.3742],
    #                        ...,])
    'TARGET': train_sdf['TARGET'].tolist()
    # TARGET: [0, 1, 0, 0, 1, 0, ...]
}
# clf.fit(data = train_sdf)
# clf.fit(data = train_dict)


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

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

### Molecule representation with UniMolRepr

Source: https://github.com/deepmodeling/unimol_tools/blob/main/README.md

Generates UniMol representations for a list of SMILES strings using the UniMolRepr class. If no pretrained paths are provided, the pretrained model and dictionary are fetched from Hugging Face. The example prints the shape of the CLS token representation and the atomic-level representations, which align with RDKit atoms.

```python
import numpy as np
from unimol_tools import UniMolRepr
# single SMILES UniMol representation. If no paths are provided the
# pretrained model and dictionary are fetched from Hugging Face.
clf = UniMolRepr(
    data_type='molecule',
    remove_hs=False,
    # pretrained_model_path='/path/to/checkpoint.ckpt',
    # pretrained_dict_path='/path/to/dict.txt',
)
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)
```

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

### Configure MASTER_ADDR and MASTER_PORT for multi-node training

Source: https://github.com/deepmodeling/unimol_tools/blob/main/docs/source/quickstart.md

Shows how to set the MASTER_ADDR and MASTER_PORT environment variables for multi-node DDP training.

```bash
export MASTER_ADDR='localhost'
export MASTER_PORT='19198'
```

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

### Multi-node pretraining with torchrun

Source: https://github.com/deepmodeling/unimol_tools/blob/main/README.md

Launches multi-node pretraining with `torchrun`. Specify `--nnodes`, `--nproc_per_node`, `--node_rank`, `--master_addr`, and `--master_port`. Replace placeholders with actual values.

```bash
export TORCH_NCCL_ASYNC_ERROR_HANDLING=1
export HYDRA_FULL_ERROR=1
export OMP_NUM_THREADS=1

torchrun --nnodes=2 --nproc_per_node=8 --node_rank=0 \
    --master_addr=<master-ip> --master_port=<port> \
    -m unimol_tools.cli.run_pretrain ...
```

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

### Pretrain with CSV dataset using torchrun

Source: https://github.com/deepmodeling/unimol_tools/blob/main/docs/source/quickstart.md

Runs Uni-Mol pretraining with a CSV dataset. For multi-machine scaling, add torchrun arguments like --nnodes, --node_rank, --master_addr, --master_port. Checkpoints and dictionary are written to the output directory. Increase training.update_freq to accumulate gradients when GPU memory is limited.

```bash
torchrun --standalone --nproc_per_node=NUM_GPUS \
    -m unimol_tools.cli.run_pretrain \
    dataset.train_path=train.csv \
    dataset.valid_path=valid.csv \
    dataset.data_type=csv \
    dataset.smiles_column=smiles \
    training.total_steps=10000 \
    training.batch_size=16 \
    training.update_freq=1
```

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

### Get molecular representations with AutoModel

Source: https://github.com/deepmodeling/unimol_tools/blob/main/README.md

Loads the allh checkpoint, tokenizes SMILES strings, and extracts the CLS representation using model.get_cls_repr. The output shape is printed as torch.Size([2, 512]). Requires torch and the unimol_hf module.

```python
from importlib.resources import files

import torch
import unimol_hf
from transformers import AutoModel, AutoTokenizer
from unimol_hf import UnimolDataCollator

pretrained = files("unimol_hf").joinpath("pretrained/unimol-v1-allh")
smiles = ["CCO", "c1ccccc1"]

tokenizer = AutoTokenizer.from_pretrained(str(pretrained))
model = AutoModel.from_pretrained(str(pretrained)).eval()
collator = UnimolDataCollator(pad_token_id=tokenizer.pad_token_id)
batch = collator([tokenizer.encode(smi) for smi in smiles])

with torch.no_grad():
    cls_repr = model.get_cls_repr(**batch)

print(cls_repr.shape)  # torch.Size([2, 512])
```

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

### UniMolRepr molecule and pocket representation

Source: https://github.com/deepmodeling/unimol_tools/blob/main/docs/source/quickstart.md

Shows how to get CLS and atomic-level representations for a single SMILES string and for a pocket dictionary. For pockets, select atoms nearby and keep total atoms preferably not exceeding 256. The atomic_reprs align with rdkit mol.GetAtoms().

```python
import numpy as np
from unimol_tools import UniMolRepr
# single smiles unimol representation
clf = UniMolRepr(data_type='molecule', # avaliable: molecule, oled, pocket. Only for unimolv1.
                 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)

# For the pocket, please select and extract the atoms nearby; the total number of atoms should preferably not exceed 256.
clf = UniMolRepr(data_type='pocket',
                 remove_hs=False,
                 )
pocket_dict = {
    'atoms': atoms,
    'coordinates': coordinates,
    'residue': residue, # Optional
}
unimol_repr = clf.get_repr(pocket_dict, return_atomic_reprs=True)
```

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

### Run default unit tests

Source: https://github.com/deepmodeling/unimol_tools/blob/main/tests/README.md

Runs the default unit tests in the 'unit/' directory.

```bash
python -m pytest
```

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

### Pretrain Uni-Mol on CSV dataset

Source: https://github.com/deepmodeling/unimol_tools/blob/main/README.md

Launches single-node pretraining with `torchrun` on a CSV of SMILES strings. Overrides dataset and training options via Hydra. Replace `NUM_GPUS` with the number of GPUs.

```bash
export TORCH_NCCL_ASYNC_ERROR_HANDLING=1
export HYDRA_FULL_ERROR=1
export OMP_NUM_THREADS=1

torchrun --standalone --nproc_per_node=NUM_GPUS \
    -m unimol_tools.cli.run_pretrain \
    dataset.train_path=train.csv \
    dataset.valid_path=valid.csv \
    dataset.data_type=csv \
    dataset.smiles_column=smiles \
    training.total_steps=1000000 \
    training.batch_size=16 \
    training.update_freq=1
```

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

### Continue training with MolTrain using load_model_dir

Source: https://github.com/deepmodeling/unimol_tools/blob/main/docs/source/quickstart.md

Shows how to train a model and then continue training from a saved checkpoint by setting load_model_dir to the previous save_path. The first fit trains the model, the second fit continues training on new data.

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

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

### Pretrain Uni-Mol on LMDB dataset

Source: https://github.com/deepmodeling/unimol_tools/blob/main/README.md

Launches single-node pretraining with `torchrun` on an LMDB dataset. Sets environment variables for NCCL, Hydra, and OpenMP. `dataset.dict_path` is optional; effective batch size is `n_gpu * training.batch_size * training.update_freq`.

```bash
export TORCH_NCCL_ASYNC_ERROR_HANDLING=1
export HYDRA_FULL_ERROR=1
export OMP_NUM_THREADS=1

torchrun --standalone --nproc_per_node=NUM_GPUS \
    -m unimol_tools.cli.run_pretrain \
    dataset.train_path=train.lmdb \
    dataset.valid_path=valid.lmdb \
    dataset.data_type=lmdb \
    dataset.dict_path=dict.txt \
    training.total_steps=1000000 \
    training.batch_size=16 \
    training.update_freq=1
```

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

### Command-line prediction with run_predict

Source: https://github.com/deepmodeling/unimol_tools/blob/main/docs/source/quickstart.md

Launches prediction from the command line using a trained model. Loads the model from ./exp and predicts on test.csv.

```bash
python -m unimol_tools.cli.run_predict load_model=./exp data=test.csv
```

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

### Enable DDP training with MolTrain

Source: https://github.com/deepmodeling/unimol_tools/blob/main/docs/source/quickstart.md

Shows how to enable Distributed Data Parallel (DDP) training with MolTrain by setting use_ddp=True and use_gpu='all'. The code must be wrapped in if __name__ == '__main__': to avoid multiprocessing errors.

```python
from unimol_tools import MolTrain

# Initialize the training class with DDP enabled
if __name__ == '__main__':
    clf = MolTrain(
        task='regression',
        data_type='molecule',
        epochs=10,
        batch_size=16,
        save_path='./model_dir',
        remove_hs=False,
        target_cols='TARGET',
        use_ddp=True,
        use_gpu="all"
        )
    pred = clf.fit(data = train_data)
```

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

### Predict with a trained model using unimol_tools CLI

Source: https://github.com/deepmodeling/unimol_tools/blob/main/README.md

Runs prediction via `python -m unimol_tools.cli.run_predict`. Loads a trained model from `load_model` and applies it to `data_path`.

```bash
python -m unimol_tools.cli.run_predict load_model=./exp data_path=test.csv
```

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

### Load Uni-Mol pretrained checkpoint paths

Source: https://github.com/deepmodeling/unimol_tools/blob/main/README.md

Loads the paths to the two provided Uni-Mol checkpoints using importlib.resources.files. Use unimol-v1-allh to keep hydrogens (remove_hs=False) and unimol-v1-noh to remove hydrogens (remove_hs=True).

```python
from importlib.resources import files

all_h_pretrained = files("unimol_hf").joinpath("pretrained/unimol-v1-allh")
no_h_pretrained = files("unimol_hf").joinpath("pretrained/unimol-v1-noh")
```

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

### Regression with CSV input using MolTrain and MolPredict

Source: https://github.com/deepmodeling/unimol_tools/blob/main/docs/source/quickstart.md

Loads a CSV file with SMILES and target columns, trains a regression model with MolTrain, and predicts with MolPredict. Default column names are 'SMILES' and 'TARGET', customizable via smiles_col and target_cols. Uses metrics='mae'.

```python
from unimol_tools import MolTrain, MolPredict

# Load the dataset
train_data_path = '../datasets/ESol/train_data.csv'
test_data_path = '../datasets/ESol/test_data.csv'

reg = MolTrain(task='regression',
                data_type='molecule',
                epochs=20,
                batch_size=32,
                metrics='mae',
                smiles_col='smiles',
                target_cols=['ESOL predicted log solubility in mols per litre'],
                save_path='./exp_esol',
                )

reg.fit(train_data_path)

predictor = MolPredict(load_model='./exp_esol')
y_pred = predictor.predict(data=test_data_path)
```

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

### Multilabel Regression with CSV File Path

Source: https://github.com/deepmodeling/unimol_tools/blob/main/docs/source/quickstart.md

Trains a multilabel regression model by passing a CSV file path directly to MolTrain. Uses metrics='mae' and target columns 'expt,calc' with save_path='./exp_csv'.

```python
from unimol_tools import MolTrain, MolPredict

# Load the dataset
data_path = '../datasets/FreeSolv/SAMPL.csv'

mreg = MolTrain(task='multilabel_regression',
                data_type='molecule',
                epochs=20,
                batch_size=32,
                metrics='mae',
                smiles_col='smiles',
                target_cols='expt,calc',
                save_path='./exp_csv',
                )

mreg.fit(data_path)
```

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

### Train a regression model with unimol_tools CLI

Source: https://github.com/deepmodeling/unimol_tools/blob/main/README.md

Runs training via `python -m unimol_tools.cli.run_train`. Overrides options from YAML config files in `unimol_tools/config` using key-value pairs. Requires a CSV with SMILES and target columns; saves the model to `save_path`.

```bash
python -m unimol_tools.cli.run_train \
    train_path=train.csv \
    task=regression \
    save_path=./exp \
    smiles_col=smiles \
    target_cols=[target1] \
    epochs=10 \
    learning_rate=1e-4 \
    batch_size=16 \
    kfold=5
```

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

### Run both default CI test groups locally

Source: https://github.com/deepmodeling/unimol_tools/blob/main/tests/README.md

Runs both the unit and pipeline test groups locally.

```bash
python -m pytest tests/unit tests/pipeline
```

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

### Multilabel Regression with Manual Data Splitting

Source: https://github.com/deepmodeling/unimol_tools/blob/main/docs/source/quickstart.md

Loads a CSV dataset, manually splits into training/validation/test sets using numpy, and trains a multilabel regression model. Uses metrics='mae' and target columns 'expt' and 'calc' with save_path='./exp_dict'.

```python
from unimol_tools import MolTrain, MolPredict
import pandas as pd
import numpy as np

# Load the dataset
df = pd.read_csv('../datasets/FreeSolv/SAMPL.csv')

data_dict = {
    'SMILES': df['smiles'].tolist(),
    'expt': df['expt'].tolist(),
    'calc': df['calc'].tolist()
}

# Divide the training set and test set
num_molecules = len(data_dict['SMILES'])
np.random.seed(42)
indices = np.random.permutation(num_molecules)
train_end = int(0.8 * num_molecules)
train_val_idx = indices[:train_end]
test_idx  = indices[train_end:]

train_val_dict = {
    'SMILES': [data_dict['SMILES'][i] for i in train_val_idx],
    'expt': [data_dict['expt'][i] for i in train_val_idx],
    'calc': [data_dict['calc'][i] for i in train_val_idx],
}
test_dict = {
    'SMILES': [data_dict['SMILES'][i] for i in test_idx],
    'expt': [data_dict['expt'][i] for i in test_idx],
    'calc': [data_dict['calc'][i] for i in test_idx],
}

mreg = MolTrain(task='multilabel_regression',
                data_type='molecule',
                epochs=20,
                batch_size=32,
                metrics='mae',
                target_cols=['expt', 'calc'],
                save_path='./exp_dict',
                )

mreg.fit(train_val_dict)

predictor = MolPredict(load_model='./exp_dict')
y_pred = predictor.predict(data=test_dict, save_path='./pre_dict')
```

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

### Set custom weight directory

Source: https://github.com/deepmodeling/unimol_tools/blob/main/README.md

Sets the UNIMOL_WEIGHT_DIR environment variable to specify a directory for pre-trained weights when they are downloaded from another source.

```bash
export UNIMOL_WEIGHT_DIR=/path/to/your/weights/dir/
```

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

### Generate representations with unimol_tools CLI

Source: https://github.com/deepmodeling/unimol_tools/blob/main/README.md

Runs representation extraction via `python -m unimol_tools.cli.run_repr`. Requires `data_path` and `smiles_col` to specify the input data.

```bash
python -m unimol_tools.cli.run_repr data_path=test.csv smiles_col=smiles
```

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

### Multilabel Classification with SDF File Input

Source: https://github.com/deepmodeling/unimol_tools/blob/main/docs/source/quickstart.md

Loads an SDF file using RDKit's PandasTools, fills missing values for specific columns, and trains a multilabel classification model. Uses save_path='./exp_sdf' and target columns 'SR-HSE' and 'NR-AR'.

```python
from unimol_tools import MolTrain, MolPredict
from rdkit.Chem import PandasTools

# Load the dataset
data_path = '../datasets/tox21.sdf'

data = PandasTools.LoadSDF(data_path)

# Fill missing values ​​with 0
data['SR-HSE'] = data['SR-HSE'].fillna(0)
data['NR-AR'] = data['NR-AR'].fillna(0)

mlclf = MolTrain(task='multilabel_classification',
                    data_type='molecule',
                    epochs=20,
                    batch_size=32,
                    metrics='auc',
                    target_cols=['SR-HSE', 'NR-AR'],
                    save_path='./exp_sdf',
                    )
mlclf.fit(data)
```

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

### Regression with Trainer using unimol_hf

Source: https://github.com/deepmodeling/unimol_tools/blob/main/README.md

Trains a regression model with Trainer using the allh checkpoint and problem_type='regression'. Expects train.csv with SMILES and TARGET columns. Saves the model and tokenizer to './hf_reg_exp'.

```python
from importlib.resources import files

import pandas as pd
import unimol_hf
from transformers import AutoModelForSequenceClassification, AutoTokenizer, Trainer, TrainingArguments
from unimol_hf import UnimolConfig, UnimolDataCollator, UnimolSmilesDataset

pretrained = files("unimol_hf").joinpath("pretrained/unimol-v1-allh")
train_data = pd.read_csv("train.csv")  # columns: SMILES, TARGET

tokenizer = AutoTokenizer.from_pretrained(str(pretrained))
config = UnimolConfig.from_pretrained(
    str(pretrained),
    num_labels=1,
    problem_type="regression",
)
model = AutoModelForSequenceClassification.from_pretrained(str(pretrained), config=config)

train_dataset = UnimolSmilesDataset(
    train_data,
    tokenizer,
    smiles_col="SMILES",
    target_col="TARGET",
    problem_type="regression",
)
collator = UnimolDataCollator(
    pad_token_id=tokenizer.pad_token_id,
    problem_type="regression",
)
args = TrainingArguments(
    output_dir="./hf_reg_exp",
    per_device_train_batch_size=16,
    num_train_epochs=10,
    learning_rate=1e-4,
    remove_unused_columns=False,
    report_to=[],
)

trainer = Trainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
    data_collator=collator,
)
trainer.train()
model.save_pretrained("./hf_reg_exp")
tokenizer.save_pretrained("./hf_reg_exp")
```

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

### Classification with Trainer using unimol_hf

Source: https://github.com/deepmodeling/unimol_tools/blob/main/README.md

Trains a sequence classification model with Hugging Face Trainer. Uses the allh checkpoint, sets problem_type to 'single_label_classification' for CrossEntropyLoss, and saves the model and tokenizer to './hf_cls_exp'. Requires train.csv with SMILES and TARGET columns.

```python
from importlib.resources import files

import pandas as pd
import unimol_hf  # register Uni-Mol Auto classes
from transformers import AutoModelForSequenceClassification, AutoTokenizer, Trainer, TrainingArguments
from unimol_hf import UnimolConfig, UnimolDataCollator, UnimolSmilesDataset

pretrained = files("unimol_hf").joinpath("pretrained/unimol-v1-allh")
train_data = pd.read_csv("train.csv")  # columns: SMILES, TARGET

tokenizer = AutoTokenizer.from_pretrained(str(pretrained))
config = UnimolConfig.from_pretrained(
    str(pretrained),
    num_labels=2,
    problem_type="single_label_classification",
)
model = AutoModelForSequenceClassification.from_pretrained(str(pretrained), config=config)

train_dataset = UnimolSmilesDataset(
    train_data,
    tokenizer,
    smiles_col="SMILES",
    target_col="TARGET",
    problem_type="single_label_classification",
)
collator = UnimolDataCollator(
    pad_token_id=tokenizer.pad_token_id,
    problem_type="single_label_classification",
)
args = TrainingArguments(
    output_dir="./hf_cls_exp",
    per_device_train_batch_size=16,
    num_train_epochs=10,
    learning_rate=1e-4,
    remove_unused_columns=False,
    report_to=[],
)

trainer = Trainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
    data_collator=collator,
)
trainer.train()
model.save_pretrained("./hf_cls_exp")
tokenizer.save_pretrained("./hf_cls_exp")
```

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

### Enable DDP for molecular representation with UniMolRepr

Source: https://github.com/deepmodeling/unimol_tools/blob/main/docs/source/quickstart.md

Shows how to enable DDP for molecular representation using UniMolRepr with use_ddp=True and use_gpu='0,1'. The code must be wrapped in if __name__ == '__main__': to avoid multiprocessing errors.

```python
from unimol_tools import UniMolRepr

# Initialize the UniMolRepr class with DDP enabled
if __name__ == '__main__':
    repr_model = UniMolRepr(
        data_type='molecule',
        batch_size=32,
        remove_hs=False,
        model_name='unimolv2',
        model_size='84m',
        use_ddp=True,  # Enable Distributed Data Parallel
        use_gpu='0,1'  # Use GPU 0 and 1
    )

    unimol_repr = repr_model.get_repr(smiles_list, return_atomic_reprs=True)

    # CLS token representation
    print(unimol_repr['cls_repr'])
    # Atomic-level representation
    print(unimol_repr['atomic_reprs'])
```

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

### Multiclass classification with pandas DataFrame using MolTrain and MolPredict

Source: https://github.com/deepmodeling/unimol_tools/blob/main/docs/source/quickstart.md

Loads a CSV into a pandas DataFrame, filters targets to <=4, splits into train/test, and trains a multiclass classifier with metrics='acc'. Uses numpy for random splitting.

```python
import pandas as pd
from unimol_tools import MolTrain, MolPredict
import numpy as np

# Load the dataset
df = pd.read_csv('../datasets/ESOL/ESOL.csv')

data_dict = {
    'SMILES': df['smiles'].tolist(),
    'target': df['Number of H-Bond Donors'].tolist()
}

data_dict['SMILES'] = [smiles for i, smiles in enumerate(data_dict['SMILES']) if data_dict['target'][i] <= 4]
data_dict['target'] = [target for target in data_dict['target'] if target <= 4]

# Divide the training set and test set
num_molecules = len(data_dict['SMILES'])
np.random.seed(42)
indices = np.random.permutation(num_molecules)
train_end = int(0.8 * num_molecules)
train_val_idx = indices[:train_end]
test_idx  = indices[train_end:]

train_val_dict = {
    'SMILES': [data_dict['SMILES'][i] for i in train_val_idx],
    'target': [data_dict['target'][i] for i in train_val_idx],
}
test_dict = {
    'SMILES': [data_dict['SMILES'][i] for i in test_idx],
    'target': [data_dict['target'][i] for i in test_idx],
}

mclf = MolTrain(task='multiclass',
                data_type='molecule',
                epochs=20,
                batch_size=32,
                metrics='acc',
                save_path='./exp',
                )

mclf.fit(train_val_dict)

predictor = MolPredict(load_model='./exp')
y_pred = predictor.predict(data=test_dict)
```

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

### Run offline pipeline smoke tests

Source: https://github.com/deepmodeling/unimol_tools/blob/main/tests/README.md

Runs the offline pipeline smoke tests in the 'pipeline/' directory.

```bash
python -m pytest tests/pipeline
```