LucaOne模型的词汇表系统

https://github1s.com/lucaone/LucaOne/blob/master/src/models/alphabet.py#L1-L205

通过定义词汇表，就能输入蛋白质或者DNA。

这段代码定义了 LucaOne模型的词汇表系统(Alphabet类)，是模型处理不同生物分子序列的核心组件。让我详细解释：

主要功能

1. 定义三种预定义词汇表

# 基因专用词汇表（DNA/RNA）
gene_standard_toks = ['1', '2', '3', '4', '5', '.', '-', '*']  
# 蛋白质专用词汇表（氨基酸）
prot_standard_toks = ['L', 'A', 'G', 'V', 'S', 'E', 'R', 'T', 'I', 'D', 'P', 'K', 'Q', 'N', 'F', 'Y', 'M', 'H', 'W', 'C', 'X', 'B', 'U', 'Z', 'O', 'J', '.', '-', '*']
# 基因+蛋白质混合词汇表
gene_prot_standard_toks = ['1', '2', '3', '4', '5', 'L', 'A', 'G', 'V', 'S', 'E', 'R', 'T', 'I', 'D', 'P', 'K', 'Q', 'N', 'F', 'Y', 'M', 'H', 'W', 'C', 'X', 'B', 'U', 'Z', 'O', 'J', '.', '-', '*']

2. 特殊Token定义

gene_prot_prepend_toks = ['[PAD]', '[UNK]']  # 前置特殊token
gene_prot_append_toks = ['[CLS]', '[SEP]', '[MASK]']  # 后置特殊token

3. 关键特性解析

DNA/RNA序列处理

# 检测是否为核酸序列
ATCGU = {"A", "T", "C", "G", "U"}
if seq_type in ["gene", "dna", "rna", "nucleic_acid", "nucleotide"]:if len(ATCGU & set(list(seq.upper()))) > 0:  # 检查是否包含ATCGU字符seq = gene_seq_replace(seq)  # 将ATCGU转换为数字编码

数字编码映射

基因序列使用数字编码：

• A/T/C/G/U → 1/2/3/4/5
• 这样做是为了统一处理不同生物分子类型

4. Alphabet类的主要方法

编码序列

def encode(self, seq_type, seq):# 如果是DNA/RNA，先进行数字编码转换# 然后tokenize并转换为索引

掩码语言模型编码

def encode_for_eval_mask(self, seq_type, seq):# 用于评估掩码任务，将'-'替换为[MASK] token

实际应用示例

# 创建混合词汇表
alphabet = Alphabet.from_predefined("gene_prot")# 编码DNA序列
dna_seq = "ATCGATTAGCT"
encoded = alphabet.encode("gene", dna_seq)  # 输出: [数字索引列表]# 编码蛋白质序列  
protein_seq = "MALWMRLLPLL"
encoded = alphabet.encode("prot", protein_seq)  # 输出: [氨基酸索引列表]

总结

这个Alphabet类是一个多模态生物分子词汇表系统，它能够：

1. 统一处理DNA、RNA和蛋白质序列
2. 将核苷酸字符(A/T/C/G/U)映射为数字(1-5)
3. 支持BERT风格的特殊token([CLS]、[SEP]、[MASK]等)
4. 为LucaOne模型提供序列到索引的转换能力

这使得LucaOne模型能够用统一的框架处理不同类型的生物序列数据。