AI化合分子自回归预训练数据集构建指南：从数据清洗到模型训练

1次阅读

共计 2229 个字符，预计需要花费 6 分钟才能阅读完成。

在化学 AI 领域，分子结构的数字化表示是模型训练的基础。目前主流的表示方法有两种：

SMILES（Simplified Molecular Input Line Entry System）：用 ASCII 字符串描述分子结构，如乙醇表示为CCO
SELFIES（Self-referencing Embedded Strings）：SMILES 的改进版，天生具有 100% 语法有效性

自回归预训练（Autoregressive Pretraining）是指模型通过预测序列中的下一个元素来进行无监督学习。对于分子生成任务，这相当于让模型学习化学结构的生成规律。

PubChem 是最常用的公开化学数据库之一。通过其 REST API 可以高效获取分子数据：

import requests
import pandas as pd

# 示例：批量下载抗病毒化合物
url = 'https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/listkey/JSON'
payload = {
    'listkey': '123456',  # 实际使用需替换为有效 listkey
    'SMILES': 'property'
}

response = requests.post(url, data=payload)
compounds = response.json()['PropertyTable']['Properties']
df = pd.DataFrame(compounds)
print(f'获取到 {len(df)} 个分子')

使用 RDKit 进行分子验证和标准化是保证数据质量的关键步骤：

from rdkit import Chem
from rdkit.Chem import AllChem

def sanitize_mol(smiles):
    """分子标准化处理流程"""
    try:
        mol = Chem.MolFromSmiles(smiles)
        if mol is None:
            return None

        # 去盐处理
        mol = Chem.RemoveHs(mol)
        # 规范化化合价
        Chem.SanitizeMol(mol)
        # 标准化芳香键表示
        Chem.Kekulize(mol)

        return Chem.MolToSmiles(mol)
    except:
        return None

# 应用清洗函数
df['clean_smiles'] = df['SMILES'].apply(sanitize_mol)
df = df.dropna(subset=['clean_smiles'])

合理的验证集划分对评估模型泛化能力至关重要。建议采用：

骨架分割：基于分子骨架（Murcko scaffold）的分组
功能团分割：按特定官能团（如羧酸、氨基）划分
随机分割：基础验证方法

from rdkit.Chem.Scaffolds import MurckoScaffold

def generate_scaffold(smiles):
    mol = Chem.MolFromSmiles(smiles)
    scaffold = MurckoScaffold.GetScaffoldForMol(mol)
    return Chem.MolToSmiles(scaffold)

df['scaffold'] = df['clean_smiles'].apply(generate_scaffold)

处理百万级分子数据集时需要注意：

使用生成器避免内存爆炸
并行处理加速计算
分块存储中间结果

import multiprocessing as mp

def parallel_process(func, data, workers=4):
    """多进程处理函数"""
    with mp.Pool(workers) as pool:
        results = pool.map(func, data)
    return results

# 示例：并行清洗
chunks = [df[i:i+10000] for i in range(0, len(df), 10000)]
results = parallel_process(sanitize_chunk, chunks)

现象：Chem.MolFromSmiles 返回 None
解决：先用smiles_tokenizer 预处理，或转用 SELFIES 表示

现象：配位键处理异常
解决：设置sanitize=False 后手动处理

现象：同分异构体混淆
解决：保留@/@@ 标记或统一去除

from rdkit.Chem import Draw

sample = df.sample(3)
mols = [Chem.MolFromSmiles(s) for s in sample['clean_smiles']]
img = Draw.MolsToGridImage(mols, legends=list(sample.index))
img.save('molecule_samples.png')