📚 课程目录第3章 SMILES字符串的tokenization方法

3. SMILES 字符串的 tokenization 方法归纳总结

Tokenization 这件事，说白了就是把 SMILES 字符串拆成模型能吃的最小单位。我刚开始做这个课题时，觉得不就是按字符切一刀嘛，后来发现坑比想象中多得多。今天把我这几年的踩坑经验和你聊聊。

3.1 为什么 tokenization 如此关键？

注意力热力图的质量，很大程度上取决于 token 的颗粒度。你想想看，如果 token 切得太粗，比如把整个环结构当做一个 token，模型很难学到原子间的局部交互。如果切得太细，比如每个字符一个 token，序列长度暴增，训练效率直线下降。

我记得有一次项目，用简单字符级 tokenizer 做分子性质预测，模型收敛倒是快，但注意力热力图上到处都是噪声，根本看不出哪个原子对性质贡献最大。后来换成基于化学语义的 tokenizer，热力图一下子就清晰了。

  核心原则： token 边界应尽量对齐化学语义单元——原子、键、括号、数字、特殊符号。

3.2 常见 tokenization 方法对比

目前主流方法大概有三种，我按实践中遇到的频率排个序：

方法	实现思路	优点	缺点
字符级（Char-Level）	每个字符单独成 token，包括字母、数字、括号、等号等	实现简单，vocab 很小（约 30~50）	序列过长，丢失化学语义
原子级（Atom-Level）	将多字符原子如 Cl、Br、@ 等整体作为一个 token	序列缩短，保留元素完整性	对环数字和分支的处理需额外规则
BPE/WordPiece	从数据中学习最优子词切分	自适应 vocab，平衡长度和语义	需大量数据预训练，解释性略差

我个人习惯首选原子级 tokenization。为什么呢？因为 RDKit 本身就能帮我们做原子级拆分，代码量少，且产出结果和化学直觉高度一致。BPE 虽然时髦，但如果你只是想快速验证注意力机制的效果，原子级足够用了。

3.3 原子级 tokenization 的具体实现

这里直接给出一份我常用的 Python 代码。嗯，要注意细节——多字符原子、手性符号、同位素标记、环数字、分支括号，全都要照顾到。

import re

def atomwise_tokenize(smiles: str) -> list:
    """
    原子级 tokenizer：将 SMILES 拆分为化学语义单元。
    支持：
      - 多字符元素 (Cl, Br, Si, Se, Te, At, Rn, ...)
      - 手性 @ 和 @@
      - 同位素 [13C], [18O]
      - 环数字 (1~9 以及 % 两位数字)
      - 分支括号 ()
      - 键符号 - = # $ : /
    """
    # 多字符元素列表（非金属 + 常见半金属）
    multi_char = ['Cl', 'Br', 'Si', 'Se', 'Te', 'At', 'Rn', 'Ne', 'He', 'Li', 'Be', 'Na', 'Mg', 'Al', 'Ar', 'Ka', 'Ca', 'Sc', 'Ti', 'V', 'Cr', 'Mn', 'Fe', 'Co', 'Ni', 'Cu', 'Zn', 'Ga', 'Ge', 'As', 'Kr', 'Rb', 'Sr', 'Y', 'Zr', 'Nb', 'Mo', 'Tc', 'Ru', 'Rh', 'Pd', 'Ag', 'Cd', 'In', 'Sn', 'Sb', 'Te', 'Xe', 'Cs', 'Ba', 'La', 'Ce', 'Pr', 'Nd', 'Pm', 'Sm', 'Eu', 'Gd', 'Tb', 'Dy', 'Ho', 'Er', 'Tm', 'Yb', 'Lu', 'Hf', 'Ta', 'W', 'Re', 'Os', 'Ir', 'Pt', 'Au', 'Hg', 'Tl', 'Pb', 'Bi', 'Po', 'At', 'Rn', 'Fr', 'Ra', 'Ac', 'Th', 'Pa', 'U', 'Np', 'Pu', 'Am', 'Cm', 'Bk', 'Cf', 'Es', 'Fm', 'Md', 'No', 'Lr', 'Rf', 'Db', 'Sg', 'Bh', 'Hs', 'Mt', 'Ds', 'Rg', 'Cn', 'Nh', 'Fl', 'Mc', 'Lv', 'Ts', 'Og']
    multi_set = set(multi_char)
    # 构建正则：原子 + 同位素/电荷括号 + 手性 + 环数字 + 分支 + 键
    pattern = r'\[[^\]]*\]|Br|Cl|' + '|'.join(sorted(multi_char, key=len, reverse=True)) + r'|%[0-9]{2}|[0-9]|@@?|[-=#$:/\\(\\)\\[\\]\\+\\-]|.'
    tokens = re.findall(pattern, smiles)
    # 过滤空串和无关符号
    tokens = [t for t in tokens if t.strip()]
    return tokens

# 测试
test_smiles = "CCOC(=O)C1=C(C)C(=O)C2=C1C=CC=C2"
print(atomwise_tokenize(test_smiles))
# 输出: ['C', 'C', 'O', 'C', '(', '=', 'O', ')', 'C', '1', '=', 'C', '(', 'C', ')', 'C', '(', '=', 'O', ')', 'C', '2', '=', 'C', '1', 'C', '=', 'C', 'C', '=', 'C', '2']

避坑提醒： 多字符元素和单字符元素有重叠，比如 "Ca" 和 "C"+"a"。我的做法是先匹配多字符元素，再匹配单字符——重新排序一下模式即可。

3.4 处理特殊语法单元

有几个地方容易出问题，我曾经因为这个 debug 了一整个下午：

环数字 % 两位： 比如 %10 是两个字符，必须整体作为一个 token，否则拆成 % 和 10 会导致环索引错乱。
手性 @ 和 @@： 它们代表 tetrahedral 构型的朝向，拆开就丢失了立体化学信息。
方括号内同位素或电荷： 比如 [13CH3] 整体作为一个 token，不要拆成 [、13、C、H、3、]。
键符号 - = # $ : / \ ： 每个都是独立 token，但注意有些 SMILES 用 '/' 和 '\' 表示顺反异构。

重要： 如果你的模型不关心立体化学，可以把 @ 和 @@ 合并为同一个特殊 token，但千万别直接丢弃，否则训练时会出现大量未知 token。

3.5 从 token 序列构建注意力热力图的关键步骤

有了 token 序列后，下一步就是把它喂给 Transformer 模型。这里我画了一张图，帮你理清整个 pipeline：

从图中你可以看到，tokenization 是整个链条的基石。我经常跟团队讲，tokenization 做不好，后面的一切都是空中楼阁。

3.6 实用经验：避免 token 数量爆炸

有些分子特别长，比如多肽、大环化合物，SMILES 可能上百个字符。如果直接用字符级切分，序列长度可能突破 200。这会导致 Transformer 自注意力计算量呈 O(L²) 增长，训练时显存瞬间爆满。

我建议的优化策略：

尽量用原子级 tokenizer，它通常能将序列长度降低 20%~40%
对于常规 drug-like 分子，中位长度大约 30~60 个 token，很友好
如果遇到超长分子，可以加入 max_length 截断或使用滑动窗口
不要忘记 padding token 和 mask

3.7 总结一下本章要点

好的，我们来快速盘点一下：

Tokenization 的粒度直接影响注意力热力图的可解释性
原子级 tokenization 是当前最实用的方案，兼顾语义和效率
处理多字符元素、环数字、手性符号时要格外小心，正则表达式是最可靠的伙伴
将 token 序列映射回原子索引，是生成热力图前必须做的对齐工作