领域专用小型语言模型——生成蛋白质结构

数据智能老司机
0 阅读22分钟

本章涵盖：

  • 使用 LLM 生成蛋白质结构
  • 使用 LLM 生成晶体结构
  • 优化面向化学领域的 LLM 推理性能

第 7 章介绍了面向 Python 代码生成的领域专用小型语言模型，也就是 SLM 的实践示例，以及如何优化这些模型，使其能够在普通硬件甚至笔记本电脑上运行。本章将展示面向药物研发和化学领域的专门化 SLM 真实案例。

8.1 在化学中应用 Transformer

第 7 章中，你已经看到如何利用本书第一部分介绍的技巧和方法，克服推理方面的挑战。在后续章节进入更高级的推理优化和模型量化之前，我们会把这些技术应用到化学领域的专门化 SLM 上，说明它们并不局限于创建编码助手。如果你不熟悉化学，不用担心——我们会保持高层次讲解，只覆盖理解模型描述和推理优化策略所必需的内容。

化学是最受自然语言处理，也就是 NLP 进展影响的领域之一。研究人员已经把许多化学任务重新表述为文本序列问题。结合开放数据集，这激发了一波研究工作，最初从定义清晰的基础挑战开始。早期模型是单模态的，输入和输出都使用分子。较新的工作则是多模态的，加入了来自分析技术的光谱数据，以及人类语言，例如合成步骤，从而弥合化学建模与自然语言建模之间的差距。

一些化学术语

氨基酸——生物体用于构建蛋白质的有机分子。

抗体——免疫系统蛋白，帮助对抗细菌和病毒等病原体。

寡聚物——由少量重复单元组成的分子，这些重复单元称为单体。

蛋白质——由氨基酸构成的分子，支持生长、组织修复和其他身体功能。

反应物——在化学反应中发生变化或被消耗的物质。

试剂——被加入以触发反应，或用于检测某种反应或特定化学物质的物质。

SMILES——Simplified Molecular Input Line Entry System，一种用于以计算机可读形式表示分子结构的化学记法。

一个早期且值得注意的工作是 Molecular Transformer。有机合成是药物化学中的关键构件，而合成规划中一个必要但仍未解决的步骤是产物预测：给定反应物和试剂，最终产物会是什么？在这项工作中，反应预测被视为反应物、试剂和产物 SMILES 字符串之间的机器翻译问题。该模型优于早期方法，在一个常见 benchmark 数据集上实现了超过 90% 的 top-1 accuracy。Molecular Transformer 学习数据集中反应物、试剂和产物中化学基团存在或缺失之间的相关性，并且不需要手工规则。

药物发现是迄今为止采用语言模型获得最大收益的领域。正如第 1 章提到的，由于大型商业语言模型主要是在从 Web 抓取的数据上训练的，因此在需要专业知识的任务上表现较差。而这类知识通常需要领域专用、往往是专有的数据，这些数据并不为科技公司所拥有。在其他情况下，模型需要针对任务进行专门调优，或者数据结构过于复杂，使大型通用模型难以发挥作用。在所有这些场景中，大型商业模型的封闭性，以及微调它们的高成本，都使其无法达到药物发现所需的性能。

基于这些原因，学术界和工业界转向了更小的开源预训练语言模型。这些模型提供了完整控制权，并使发布面向多种药物发现任务的生产可用模型成为可能。每个月都会出现新的专门化模型，其中许多采用宽松开源许可证，通常不仅覆盖源代码和模型 checkpoint，还覆盖训练和测试数据集。我维护了一个 GitHub repo，用来跟踪这一领域中主要的完全开源项目。

本章将介绍面向蛋白质和晶体结构生成的领域专用 SLM 的重要案例，同时讨论 ML/AI 工程方面的考虑，以及提升推理性能的建议。

8.2 从自然语言到蛋白质结构

我们将探索的第一个 SLM 是 ProtGPT2。它是一个会“说”蛋白质语言的模型，可用于新型蛋白质设计与工程。ProtGPT2 生成的序列能够保持天然蛋白质的关键特征，例如氨基酸倾向性、二级结构内容和球状性，同时探索此前未被观察到的蛋白质空间区域。

ProtGPT2 是一个 decoder-only GPT-2 模型，拥有 7.38 亿参数。它以自监督方式预训练，用于预测序列中的下一个 token，也就是寡聚物。训练数据来自 UniRef50 数据库中的 87,162 条原始蛋白质序列，并且不包含相应注释。预处理后的训练集可在 Hugging Face Hub 上获得。给定自然语言 prompt 后，它可以生成蛋白质序列。预训练模型也可以在 Hugging Face Hub 上获得。

我们来看看这个模型如何工作，并用它在 zero-shot 设置下生成 de novo 蛋白质。本示例提供了一个配套 Colab notebook。由于该模型可在 Hugging Face Hub 上获取，我们会使用 Transformers 库下载它，并设置推理 pipeline：

from transformers import pipeline

protgpt2 = pipeline('text-generation', model="nferruz/ProtGPT2")

模型下载完成后，可以使用前面的 pipeline 生成蛋白质序列：

sequences = protgpt2("<|endoftext|>", max_length=100, 
                    do_sample=True, top_k=950,
                                   repetition_penalty=1.2, 
                    num_return_sequences=10,
                                   eos_token_id=0)

这与生成文本或代码没有区别。这里，我们要求模型生成 10 条蛋白质序列；结果看起来如下：

{'generated_text': 
'<|endoftext|>\nMGDIPAELIRNWNYYSYNFTLLHITFFLLGTLFHLIQKELKGKKPKSKGKKKQEESVAE
E\nKVKEEPPIMVQGKAKKEKTKTKKRDIQEAEPVPMVEPLPVDELEKLDKEMKKKKKKDKKK\nEKQEIEQKAV
VEQAPVVEQSAAVEQLPVAEQEKQEKKAKAKKSEPEEKKPEVKLKEKKRK\nEKEREETEEKPAEEPQAVEDEEH
EEVAKKSKSKKRKADVQEPEVEEVKVEAEAPKKKKKR\nKQSDDDSEQEKKGKKEKKSKK'}

{'generated_text': 
'<|endoftext|>\nMVQPQSASDPPASPPTTDHPLSPLGNPRVVVGRCLLLLALLFHTFTYTWAYAAGSLYCV
A\nVCMPTVLIYQCALLLCLTVSRRHHRVSRGGSRWPSRLVWLRSTKWAVVLLLIFFSLFGLA\nTQCFVIFCHP
TAYSQSTMQHHSPAANITTHCGPNAHNDSHHVALSQWVETAVAWAFVASV\nVAWFASVAVVIPRWTGYCRLCCR
DSASSSNDDASGMEDKWDEEEYEDEE'}

{'generated_text': 
'<|endoftext|>\nMPSRPRNPRTVAGAGPVPVAAHAPQHIETHGVRLHYVERGTGPPVVLVHGLGGQIDSMV
G\nALPALAPRHRVVALDQRGHGESDKPATGYGIAEIAADLDGFIAXXXCPAFGCHGRQGSGR\nGGPGRAAAAG
RSSEPRGRGLHARSAQPAQLDHGHRAGPRGGEVRRRGRRTRGCRPSLTPD\nLSTTPRGPA\n'}

{'generated_text': 
'<|endoftext|>\nMHIWGCGPSTIGLGINTKNPKTIFFILLITPDVLSISFAIPPKNKSLFSSIKYHVISKK
N\nKFLNISKIGIFLPSSSFCSCLYRCLIRGTTFTFNERPSTQTRSNNHQ\n'}

{'generated_text': 
'<|endoftext|>\nMPHPIALAAIAGAPLALCAVLWSHVAALEGQAAWSMAAMSAMLGTVTALRRTRLLGPAA
T\nAAFLAGFGWVALAELAAPPRAAHHLAALGPGAPVLVRVEREARADPVRTRYRLVTGDRAD\nGPGGRAWTAP
RPRLLVWFRAAGPPERWVPGDRAGARLRLMTPRGPTNPGALDAELAALRR\nGIAASGYLRGPAGVSASDPLSWR
RLRGRFERWLDRQAAAAGPRGAVAALATGDEAGIPAP\nDWAVFRDTGTAHLMAVSGLHIGLVAGLAGALAARLY
G'}

{'generated_text':
   '<|endoftext|>\nMKHAVWILLLFYFAETSAQLYSTHFHANPFNPLNMRPAHNSTELFFIPSHSFFGPVISA
Q\nNGFQQNPTAVNQYPTQTAQIYPTYSMYQPSSIPYYQQLTASTTPYNEYTSFPQNTYNTPS\nNYDAYYYPQT
QVYQYYQPSTTQTPMNYRSYNYTTHYPPIQQPTEYASNAYSMYQQQQTPP\nSQSPYNTQSTTTQPSTFSPQETM
LEVSPTPPASNNVPSPVAISPPPQVPPQSASSPISTA'}

{'generated_text':
'<|endoftext|>\nMSPIAWFFLASGSIALLGNVLPLWGPWRPRWLTPLAVVLFATWALQRWEAALFLLVLVV
A\nLIGLAGRAARVLRASLFARRLRGRAAAEVAVEPMDAEAAHAATAAAAAGDQP\n'}

{'generated_text': 
'<|endoftext|>\nMSGFVAVLFTATLFICCLIPCLAFIYILGNKDLGKDHPLIYLLSLASGTLCGDAFLHLI
P\nHLLENHETHGTESNGHSNMALWVLLGVLFFFMCEKLFHLITEYKQSHAHAMKISKKTMIV\nGDDAKLYDPV
TTRKINTSKHCCSSSGSPSYSEDTSGETICNNTNENGLLSVNNISNNDFN\nNANKNNHLNKNPNKHGTDNVETN
QNSSETDNKQNKKHSHNQKCNGDLEVSSKGNHRSTAS\nSS'}

{'generated_text': 
'<|endoftext|>\nMRALILFLLLCGCAVTTPATTVRTAQPTLEEPGRSERVWLWSTRAPERWLVVRDGTVRG
L\nDPASGRTIWQLAPGEERWGTEAVGDELLVTRAGTASVELAVRAKDLTGRAQLTVSTVPTS\nVPTPSPEPPL
WVVEGDAAPVDRARWGRAAVTLGDQQLVADPRERRVFAVGTPSGGWSLAR\nLDAATGAPLWQTPVAAGQALVAG
AAPGSVFVVGYDRLSAFDAASGAETWAGPVGEREGRG\nPAAAPTATVDGT'}

{'generated_text':
'<|endoftext|>\nMRLIWTFLIFFLTIGRAQDSLNNSRIIPLVTLKDFCNSGISIIPNGENEDELYKIKNKS
L\nANQFERIKYSLNRLDFEIDSIRCDSLNRFNDSLVLYDTSLFRIDSTFGNLEKSKNDIDNK\nISILEKNLKS
LLKSEHLKFSNVITFIVLIFFFIIAVLLWILKRSRNKTIINITINSKKIL\nELDRLEYEKKLTTMENELNQKIA
ENRLLMEQTRLKEIFETKTDSLDSQNQTLSSQNEKII\nAENIRFQEEKALLEQKL'}

模型选择生成以氨基酸 “M” 开头的序列。我们没有指定 prompt，所以这个选择由模型决定；不过你也可以指定任何其他氨基酸、寡聚物、片段或蛋白质作为起点。由于模型比较紧凑，在没有硬件加速的 Colab 免费层上，平均不到 2 分钟就可以生成一批 10 条序列。这是可以接受的，因为这些任务并不一定需要实时完成，而且模型可以运行在普通硬件上。对于这个模型来说，重点可能更应该放在生成蛋白质序列的质量上。下一节中，我们会评估一个与 ProtGPT2 不同、需要更多关注推理性能的模型。

评估 LLM 生成内容质量的一个常用指标是 perplexity score，第 6.2.1 节已经介绍过。在那一节中，我们也实现了一个基础函数，用于计算 LLM 生成文本的 perplexity。将同一个函数应用到这 10 条生成蛋白质序列上，会得到如下分数：

275.8756703543678
552.265406680504
222.6862045710289
739.1276632256441
202.7109209811726
449.55308995221174
419.3895659171086
385.7640621074327
404.39583964519915
573.9670349466911

与文本生成一样，perplexity 分数越低越好。

这个领域专用开源 SLM 表明，只要数据合适，你并不需要一个大模型，就可以解决真实世界中非平凡的任务。

8.3 使用 SLM 生成抗体

本章中的第二个 SLM 是 Antibody Generator，也被称为 AntibodyGPT。这是一个专门化蛋白质生成模型，用于创建可用于免疫学、疫苗开发和医疗治疗的治疗性抗体序列。抗体是保护我们身体的蛋白质。当异物，也就是抗原，进入身体时，免疫系统会产生抗体。抗体会与这些物质结合，帮助消除它们。Antibody Generator 基于 ProGen2。ProGen2 是 Salesforce 开发的高级语言模型，训练于超过 2.8 亿条蛋白质序列，并在论文《ProGen2: Exploring the Boundaries of Protein Language Models》中提出。

ProGen2 家族采用 decoder-only 架构，包含四种模型规模：small，1.51 亿参数；medium，7.65 亿参数；large，27 亿参数；extra-large，64 亿参数。Antibody Generator 模型以 ProGen2 的 small、medium 和 large 变体作为 baseline 构建，并在 5,000 个实验解析的抗体及其抗原晶体结构上进行微调，见图 8.1。三种 Antibody Generator 模型选项都可在 Hugging Face Hub 上获得。

image.png

图 8.1 —— ProGen2 和 Antibody Generator 模型家族的不同可用选项

我们从 small 模型开始，用 AntibodyGPT 运行推理。本示例提供配套 Colab notebook，并且需要 GPU 加速。

要从 Hugging Face Hub 下载该模型或另外两个版本的 checkpoint，并将其加载到内存中，你需要一个自定义 auto class：ProGenForCausalLM。它由 AntibodyGPT 的共同作者 Joseph Roberts 在一个公开 GitHub repo 中实现。首先克隆该 repo：

!git clone https://github.com/joethequant/docker_protein_generator.git
%cd ./docker_protein_generator/

现在可以下载模型和 tokenizer：

import torch
from models.progen.modeling_progen import ProGenForCausalLM
from tokenizers import Tokenizer

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

model_id = 'AntibodyGeneration/fine-tuned-progen2-small'
model = ProGenForCausalLM.from_pretrained(model_id)
model = model.to(device)
tokenizer = Tokenizer.from_pretrained(model_id)

模型加载到 GPU 内存之后，我们可以定义目标抗原序列：

target_sequence = 
'MQIPQAPWPVVWAVLQLGWRPGWFLDSPDRPWNPPTFSPALLVVTEGDNATFTCSFSNTSESFVLNWYRMSPSN
QTDKLAAFPEDRSQPGQDCRFRVTQLPNGRDFHMSVVRARRNDSGTYLCGAISLAPKAQIKESLRAELRVTERRA
EVPTAHPSPSPRPAGQFQTLVVGVVGGLLGSLVLLVWVLAVICSRAARGTIGARRTGQPLKEDPSAVPVFSVDYG
ELDFQWREKTPEPPVPCVPEQTEYATIVFPSGMGTSSPARRGSADGPRSAQPLRPEDGHCSWPL'

并且可以指定要生成多少条抗体序列：

number_of_sequences = 2

在把 prompt 传给模型并运行推理之前，需要对目标序列进行分词，并将其转换为 PyTorch tensor：

tokenized_sequence = tokenizer.encode(target_sequence)
input_tensor = torch.tensor([tokenized_sequence.ids]).to(device)

现在可以开始序列生成，也就是推理：

with torch.no_grad():
    output = model.generate(input_tensor, max_length=1024, 
    pad_token_id=tokenizer.encode('<|pad|>').ids[0], 
    do_sample=True, top_p=0.9, temperature=0.8, 
    num_return_sequences=number_of_sequences)

为了获取生成序列，需要按如下方式解码模型输出：

as_lists = lambda batch: [batch[i, ...].detach().cpu().numpy().tolist()
➥for i in range(batch.shape[0])]
sequences = tokenizer.decode_batch(as_lists(output))
if len(sequences) > 0:
    sequences = [x.replace('2', '') for x in sequences]

尽管这个模型比 ProtGPT 更小，但如果想快速生成仅两条目标序列，仍然需要 GPU。使用同一模型和设置，在没有硬件加速的 Colab 免费层上进行推理需要几个小时。在这些条件下，benchmark 显示平均延迟如下：

('P95 latency (ms) - 137426.0512944001; Average latency (ms) 
- 119778.24 +\- 27239.48;',
 137426.0512944001)

在 CUDA allocator 配置中，expandable-segment allocation 默认设置为 false。这可能会使 allocator 的内存段碎片化，阻止这些内存段在后续推理中扩展并复用。你可以测试改变 CUDA caching allocator 是否能改善延迟：

%env PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

应用前面的设置后重新运行 benchmark，可以获得略低的平均延迟：

('P95 latency (ms) - 128515.5656411501; Average latency (ms)
 - 103638.09 +\- 29278.26;',
 128515.5656411501)

在优化这个模型之前，我们先看看如何验证它生成的结果。对于这样的场景，perplexity 及类似统计指标本身意义不大。一个更可靠的选项是 ANARCI 工具，它由 Oxford Protein Informatics Group，也就是 OPIG，开发和维护，用于抗体编号和受体分类。ANARCI 会检查生成的抗体序列是否符合已知模式和结构，从而帮助我们评估其结构可行性，以及与既有框架的一致性。目标是让序列不仅新颖，还要具有生物学相关性，并可能具备功能。你可以把生成的抗体序列提交到 ANARCI 在线工具；图 8.2 展示了本示例中某条序列的结果。

image.png

图 8.2 —— 在线 ANARCI 工具对本示例中某条生成抗体序列的输出

在许多领域中，使用专用工具验证生成内容是一种常见做法。领域专家始终应该审查这些内容。你也可以使用 Python ANARCI API 以程序化方式执行这种验证。

接下来，我们看看这个模型是否能从量化中受益。使用 Bitsandbytes 库，我们可以将其量化为 8-bit。首先定义量化配置：

from transformers import BitsAndBytesConfig

bnb_config = BitsAndBytesConfig(
  load_in_8bit=True,
)

现在通过 ProGenForCausalLM auto class 使用它：

model_id="AntibodyGeneration/fine-tuned-progen2-small"
model_int8 = ProGenForCausalLM.from_pretrained(model_id,
                                            quantization_config=bnb_config,
                                             )
tokenizer = Tokenizer.from_pretrained(model_id)

符合预期，8-bit 量化模型比原始模型更小，156.73 MB 对比 617 MB。运行与原始模型相同的 benchmark，可以观察到如下性能：

('P95 latency (ms) - 80453.97859845006; Average latency (ms) 
- 61846.30 +\- 18507.29;',
 80453.97859845006)

在启用动态 expandable segments 之后，平均延迟通常比 vanilla 模型 benchmark 低约 40%。

8.4 从 CIF 文件到晶体结构

我们将探索的第三个 SLM 是 CrystaLLM，它来自论文《Crystal Structure Generation with Autoregressive Large Language Modeling》。它是一个基于 Transformer 的 decoder-only LLM，用于生成晶体结构。此前，晶体结构预测方法通常基于计算密集型的 ab initio 技术，并且通常依赖专用工具来优化候选结构。

CrystaLLM 是第一个训练于 crystallographic information files，也就是 CIF 的文本表示，而不是传统化学表示的模型。作者认为，符号序列可以表示复杂化学结构，并且 LLM 可以学习超越简单相关性的更深层因果关系。该模型在大量 CIF 数据上训练，可以生成物理上合理的无机晶体结构。作者还加入了 Monte Carlo tree search，也就是 MCTS，以及基于图的神经网络，以改善结构采样。

训练过程中，模型会接收来自 CIF 文件语料的 token 序列，并学习预测下一个 token。训练之后，它可以在给定起始 token 序列的条件下生成新的 CIF 文件。生成 CIF 文件的过程，是模型反复从 token 分布中采样，并以累计生成内容为条件继续生成，直到达到终止条件。

图 8.3 来自前面引用的论文，从高层展示了 CrystaLLM 的训练和推理流程。如图所示，在训练期间，CIF 文件会被分词成符号序列。模型处理该序列，并针对每个输入符号输出一个词汇表上的概率分布。这些预测分布会与目标分布进行比较，使用 cross-entropy loss 度量；目标分布把全部概率质量放在正确的下一个 token 上。目标是输入 token 向左移动一个位置，因为任务是根据前面的 token 预测下一个 token，就像许多 LLM 内容生成任务一样。token 被分组为 CIF tag、atom、数字和标点。输出 token，在训练中不采样，表示被分配最高概率的 token。带下划线的 token 表示正确下一个 token 被赋予相对较低概率的情况。

image.png

图 8.3 —— CrystaLLM 训练过程 a，以及使用 MCTS decoding 的生成过程 b

模型训练完成后，会根据 prompt 生成 CIF 文件。prompt 通过把符号数据与目标晶胞组成拼接而成，然后被分词并输入模型。在每一步中，模型会从预测分布中采样下一个 token，并把它追加到 CIF 内容中。这个过程会重复，直到达到预定义停止条件。模型基于 GPT-2；模型定义、训练和推理代码改编自 Andrej Karpathy 的 nanoGPT 仓库。

CrystaLLM 可以为范围广泛的无机化合物生成合理晶体结构，包括训练期间没有见过的化合物。它生成语法有效、物理合理的结构，并且可以在给定化学组成和对称性信息的情况下预测晶体结构。它也能够泛化到未见过的无机结构，在高达 40% 的情况下实现结构匹配，并可以将结构基序外推到新的组成上。这项工作很重要，因为它展示了通过生成候选晶体结构并供下游优化，从而加速材料发现的潜力。

我们会跳过剩余训练细节。本章中，我们使用预训练模型，并只聚焦推理过程。本节提供配套 Colab notebook。

为了设置环境，首先克隆 CrystaLLM GitHub 仓库：

!git clone https://github.com/lantunes/CrystaLLM.git
%cd CrystaLLM

接下来，将克隆下来的源码目录加入 Python path：

import sys
import os

sys.path.append('/content/CrystaLLM')
os.environ["PYTHONPATH"] += (":/content/CrystaLLM")

安装缺失依赖之后，下载预训练模型权重。CrystaLLM 有两个版本：small，2500 万参数；large，2 亿参数。这些权重不在 Hugging Face Hub 上，因此需要使用下面的 Python 脚本下载，这里请求的是 small 版本：

!python bin/download.py crystallm_v1_small.tar.gz
!tar xvf crystallm_v1_small.tar.gz

现在可以用它生成晶体结构。首先，需要创建一个 prompt，在这个场景中，它是一个预处理后 CIF 文件的部分内容。通常，prompt 包含晶胞组成：一个公式，用来指定单位晶胞中每种原子类型的总数。例如，我们可以指定 Na2Cl2，也就是二钠二氯化物，表示单位晶胞中有四个原子，并使用提供的 make_prompt_file.py 脚本准备并保存 prompt，供后续使用：

!python bin/make_prompt_file.py Na2Cl2 my_prompt.txt

为了用晶胞组成提示模型，格式应该使用 data_ 前缀，并在行末添加换行符：

data_<cell_composition>\n

make_prompt_file.py 脚本会处理这个格式约定，以及其他内容。在我们的示例中，模型接收到的 prompt 是：

data_Na2Cl2\n

现在可以生成 CIF 文件；我们会使用提供的 sample.py 脚本：

!python bin/sample.py \
out_dir=crystallm_v1_small \
start=FILE:sample_prompt.txt \
num_samples=2 \
top_k=10 \
max_new_tokens=3000 \
device=cpu \
target=file

这个脚本接收以下参数：

  • out_dir——包含已下载模型 checkpoint 的目录。
  • start——要么是所需格式的 prompt，要么是包含 prompt 的文件路径。
  • num_samples——要生成的样本数量。
  • top_k——保留的最可能 token 的最大数量。
  • max_new_tokens——每个样本生成的最大 token 数。
  • device——模型权重卸载到哪里：CPU 或 GPU。
  • target——生成内容应该发送到哪里。在本例中是 file。也可以是 console。

还有一些我们这里没有使用的参数：

  • temperature——含义与任何 Transformer 模型中相同。
  • seed——设置随机种子以保证可复现性；默认值为 1,337。
  • dtype——模型精度；默认值为 bfloat16。
  • compile——布尔标志，表示在 CIF 生成之前是否用 PyTorch 编译模型以加快执行。默认值为 False。

这个小模型不需要硬件加速。在所选配置下，也就是每条序列最多 3,000 个 token，在 Colab CPU VM 上生成两个 CIF 文件平均不到一分钟。下面是本次运行生成的一个 CIF 文件：

data_Na2Cl2
loop_
_atom_type_symbol
_atom_type_electronegativity
_atom_type_radius
_atom_type_ionic_radius
Na 0.9300 1.8000 1.1600
Cl 3.1600 1.0000 0.7800
_symmetry_space_group_name_H-M P4/nmm
_cell_length_a 4.9038
_cell_length_b 4.9038
_cell_length_c 3.9603
_cell_angle_alpha 90.0000
_cell_angle_beta 90.0000
_cell_angle_gamma 90.0000
_symmetry_Int_Tables_number 129
_chemical_formula_structural NaCl
_chemical_formula_sum 'Na2 Cl2'
_cell_volume 95.2964
_cell_formula_units_Z 2
loop_
_symmetry_equiv_pos_site_id
_symmetry_equiv_pos_as_xyz
1 'x, y, z'
loop_
_atom_site_type_symbol
_atom_site_label
_atom_site_symmetry_multiplicity
_atom_site_fract_x
_atom_site_fract_y
_atom_site_fract_z
_atom_site_occupancy
Na Na0 2 0.0000 0.0000 0.0000 1
Cl Cl1 2 0.0000 0.5000 0.9817 1

模型输出还不能直接使用：它生成的是原始 CIF 文件，需要后处理。项目提供了一个 Python 脚本；你需要按如下方式在生成文件上运行它：

!python bin/postprocess.py . colab_processed_cifs

postprocess.py 脚本接收两个参数：包含生成原始 CIF 文件的目录，以及后处理文件的输出目录。后处理之后，前面展示的原始 CIF 文件会变成这样：

data_Na2Cl2
_symmetry_space_group_name_H-M P4/nmm
_cell_length_a 4.9038
_cell_length_b 4.9038
_cell_length_c 3.9603
_cell_angle_alpha 90.0000
_cell_angle_beta 90.0000
_cell_angle_gamma 90.0000
_symmetry_Int_Tables_number 129
_chemical_formula_structural NaCl
_chemical_formula_sum 'Na2 Cl2'
_cell_volume 95.2964
_cell_formula_units_Z 2
loop_
 _symmetry_equiv_pos_site_id
 _symmetry_equiv_pos_as_xyz
  1  'x+1/2, y+1/2, -z'
  2  '-y, x, -z'
  3  '-x, y, z'
  4  'y+1/2, -x+1/2, z'
  5  'x+1/2, -y+1/2, -z'
  6  '-y+1/2, x+1/2, z'
  7  '-x, -y, z'
  8  'x, y, z'
  9  'y, x, -z'
  10  '-x+1/2, y+1/2, -z'
  11  '-y, -x, -z'
  12  '-x+1/2, -y+1/2, -z'
  13  'y, -x, -z'
  14  '-y+1/2, -x+1/2, z'
  15  'x, -y, z'
  16  'y+1/2, x+1/2, z'
loop_
_atom_site_type_symbol
_atom_site_label
_atom_site_symmetry_multiplicity
_atom_site_fract_x
_atom_site_fract_y
_atom_site_fract_z
_atom_site_occupancy
Na Na0 2 0.0000 0.0000 0.0000 1
Cl Cl1 2 0.0000 0.5000 0.9817 1

现在需要验证生成文件。项目提供了一个验证脚本，它要求输入 CIF 文件位于 .tar.gz 压缩包中：

!tar -czvf colab_processed_cifs.tar.gz ./colab_processed_cifs/

按如下方式触发验证：

!python bin/evaluate_cifs.py colab_processed_cifs.tar.gz -o 
➥colab_processed_cifs.csv

evaluate_cifs.py 脚本接收两个参数：包含待评估 CIF 文件的压缩包名称，以及保存结果的 CSV 文件名称。结果也会在最后打印到标准输出：

space group consistent: 2/2 (1.000)
atom site multiplicity consistent: 2/2 (1.000)
avg. bond length reasonableness score: 1.0000 ± 0.0000
bond lengths reasonable: 2/2 (1.000)
num valid: 2/2 (1.00)
longest valid generated length: 631
avg. valid generated length: 560.500 ± 70.500

在其他信息和统计结果之外，这些输出显示两个生成的 CIF 文件都是有效的。

如 8.4 节开头所述，见图 8.3，CrystaLLM 论文提出了第二种生成 CIF 文件的方法：在 MCTS 循环中运行 LLM。项目也提供了脚本。可以按如下方式运行基于 MCTS 的 decoding：

!python bin/mcts.py \
out_dir=crystallm_v1_small \
device=cpu \
dtype=bfloat16 \
start=FILE:sample_prompt.txt \
tree_width=5 \
max_depth=2000 \
selector=puct \
c=1.0 \
num_simulations=1000 \
reward_k=2.0 \
scorer=random \
top_child_weight_cutoff=0.9999 \
bypass_only_child=True \
mcts_out_dir=colab_mcts_cifs

mcts.py 使用与标准生成相同的模型、精度和样本 prompt 文件。它还接收 MCTS 专用参数：tree width、maximum depth、simulation 数量，以及每次 simulation 的 rollout 数量。

如果你想把 CrystaLLM 预训练模型，或者你自己的微调版本，上传到 Hugging Face Hub 仓库，并通过 Transformers 库使用它，可以按以下步骤操作。这里展示 small 模型；large 模型遵循相同步骤，但使用不同配置值。

首先，创建 Transformers auto class 所需的配置文件 config.json。你可以从 CrystaLLM GitHub 仓库中的 config/crystallm_v1_small.yaml 复制大多数值，large 模型也有类似文件。

crystall_small_config = """{
  "bias": true,
  "model_type": "gpt2",
  "block_size": 1024,
  "dropout": 0.1,
  "n_embd": 512,
  "n_head": 8,
  "n_layer": 8,
  "vocab_size": 50257
}"""

with open('/content/CrystaLLM/crystallm_v1_small/config.json', 'a') as f:
    f.write(crystall_small_config)

接下来，将模型 checkpoint 从 .pt 转换为 .bin，这是将其加载到 auto class 中所必需的：

import torch

checkpoint = torch.load("/content/CrystaLLM/crystallm_v1_small/ckpt.pt", 
                    map_location=torch.device('cpu'))
params = checkpoint["model"]
torch.save(params, 
➥"/content/CrystaLLM/crystallm_v1_small/pytorch_model.bin")

注意，这里不是访问 model_state_dict key，而是直接访问 model key，因为这个预训练模型在 checkpoint 字典中把权重保存在该 key 下。

现在可以把模型包装在一个 Model 类中，提供生成的配置，并让它可以通过 Transformers 库使用：

from transformers import AutoConfig, AutoModelForCausalLM

config = AutoConfig.from_pretrained(
            '/content/CrystaLLM/crystallm_v1_small/config.json')
transformer_model = AutoModelForCausalLM.from_pretrained(
            '/content/CrystaLLM/crystallm_v1_small', config=config)

然后把 Model 保存到磁盘，以生成剩余文件：

transformer_model.save_pretrained
➥("/content/CrystaLLM/crystallm_v1_small_hf")

如果需要，也可以将其上传到 Hugging Face Hub。首先，在 Colab 中创建一个名为 HF_TOKEN 的新 secret。从你的 Hugging Face Hub 个人设置中复制 access token key，并粘贴到 HF_TOKEN 的 Value 字段。然后切换 Colab UI 左侧的钥匙图标，让 notebook 可以访问该 secret。最后，通过程序加载 token：

import os
from google.colab import userdata

os.environ["HF_TOKEN"] = userdata.get("HF_TOKEN")

然后可以用它登录 Hugging Face Hub，这需要安装 huggingface_hub 包：

from huggingface_hub import login

login(os.environ["HF_TOKEN"])

现在可以把模型推送到 Hub：

transformer_model.push_to_hub("crystallm_v1_small_hf")

到这里，本书这一部分关于本地领域专用模型真实案例的内容就结束了。下一部分的章节将介绍更高级的优化技术，并最终讨论在不同硬件平台上部署，以及把小型专门化语言模型集成到更大系统中。

总结

  • 化学语言模型使用 NLP 技术，将分子结构表示为文本序列，用于药物发现。
  • ProtGPT2 使用一个 decoder-only GPT-2 模型生成蛋白质序列，该模型训练于 87,162 条 UniRef50 蛋白质序列。
  • 蛋白质生成使用 endoftext prompt，生成类似天然蛋白质的氨基酸序列。
  • Perplexity scoring 可以评估生成蛋白质序列，分数越低表示输出质量越高。
  • AntibodyGPT 在 5,000 个实验解析的晶体结构上微调 ProGen2，用于生成治疗性抗体。
  • ProGenForCausalLM 类从自定义仓库加载基于 ProGen2 的模型并运行推理。
  • CUDA 内存分配设置可以通过启用 expandable segments 改善内存管理，从而降低推理延迟。
  • 8-bit 量化将 AntibodyGPT 模型从 617 MB 缩小到 156 MB，并将平均延迟改善 40%。
  • ANARCI 工具用于验证生成抗体序列是否符合已知抗体结构模式。
  • CrystaLLM 根据 crystallographic information file 的文本表示生成晶体结构。
  • CIF 文件生成过程包括对晶体学数据进行分词，并使用自回归预测完成结构。
  • Monte Carlo tree search 可以提升晶体结构生成的采样质量。
  • 验证化学结构需要领域专用工具，而不是标准文本生成指标。
  • 与通用语言模型相比，领域专用模型可以用更少参数取得有效结果。
  • 专门化化学模型需要领域专家验证其生物学相关性和可行性。
avatar
文章
阅读
粉丝