构建高效ai化合分子自回归预训练数据集的工程实践

1次阅读

没有评论

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

在 AI 化合分子自回归预训练中，数据质量直接影响模型性能。我们常遇到以下核心问题：

分子数据异构性 ：不同来源的 SMILES 表达式存在异构体、不同质子化状态等变体，导致同一分子可能有多种表示形式。例如苯环可能被表示为 c1ccccc1 或 C1=CC=CC=C1
序列长度不匹配 ：自回归模型要求固定长度输入，但原始分子 SMILES 序列长度差异可达 10 倍（如甲烷 CH4 vs 复杂药物分子）
无效结构污染 ：约 5 -15% 的公开化合物数据包含非法价态或不可能的空间构型

我们设计的分级处理流程：

初级清洗
去除盐离子和溶剂分子
标准化原子杂化状态（sp3/sp2/sp）
统一芳香性表示（Kekule vs 非定域）
立体化学处理
手性中心规范化（R/ S 标记）
双键 E / Z 构型确认
使用 RDKit 的 SanitizeMol 进行价态验证
序列优化
通过深度优先遍历生成规范 SMILES
采用 Morgan 指纹去重（半径 =3，位数 =2048）

使用 Dask 实现三级并行：

任务级 ：将千万级分子库分块处理（每块 10 万分子）
分子级 ：每个 worker 节点处理独立分子批次
操作级 ：原子环境计算等细粒度操作并行化

关键配置参数：

from dask.distributed import Client
client = Client(
    n_workers=32,
    threads_per_worker=2,
    memory_limit='8GB'
)

import numpy as np
from rdkit import Chem, DataStructs
from rdkit.Chem import AllChem

def generate_morgan_fp(smiles: str, radius=2, n_bits=2048):
    """生成规范化的 Morgan 指纹"""
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        return None
    fp = AllChem.GetMorganFingerprintAsBitVect(mol, radius, nBits=n_bits)
    arr = np.zeros((1,), dtype=np.int8)
    DataStructs.ConvertToNumpyArray(fp, arr)
    return arr

import pandas as pd
import dask.dataframe as dd

def process_chunk(df_chunk):
    """处理单个数据块的 Dask UDF"""
    results = []
    for smi in df_chunk['smiles']:
        try:
            std_smi = standardize_smiles(smi)  # 标准化函数
            fp = generate_morgan_fp(std_smi)
            if fp is not None:
                results.append({
                    'smiles': std_smi,
                    'fingerprint': fp.tobytes()})
        except Exception as e:
            log_error(e)
    return pd.DataFrame(results)

# 主处理流程
ddf = dd.read_parquet('raw_data.parquet')
processed = ddf.map_partitions(process_chunk)
processed.to_parquet('processed_data')

格式	1 千万分子大小	读取速度	写入速度
Parquet	12.4GB	38s	2m12s
HDF5	15.7GB	1m53s	3m47s
CSV	23.1GB	4m21s	6m15s

Parquet 采用 Snappy 压缩时，体积可进一步减小 40%

绝对构型保留：使用 Chem.AssignAtomChiralTagsFromStructure
相对构型处理：对无法确定绝对构型的情况，统一采用 Chem.RemoveStereochemistry
内存泄漏预防：RDKit 对象需显式销毁

# 正确做法
def safe_processing(smi):
    mol = Chem.MolFromSmiles(smi)
    try:
        # 处理逻辑
        return result
    finally:
        if 'mol' in locals():
            del mol