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类是一个多模态生物分子词汇表系统，它能够：
- 统一处理DNA、RNA和蛋白质序列
- 将核苷酸字符(A/T/C/G/U)映射为数字(1-5)
- 支持BERT风格的特殊token([CLS]、[SEP]、[MASK]等)
- 为LucaOne模型提供序列到索引的转换能力
这使得LucaOne模型能够用统一的框架处理不同类型的生物序列数据。