“学习 LLM ，在大模型时代为自己存张船票”。
相信很多人都有这样的想法。那么，在 AI for science 领域，哪些 LLM 模型值得一试呢？
笔者认为：
- LLM 直接预测 SMILES 性质 or 直接生成 SMILES 的技术路线是行不通的。因为 SMILES 相比高精度 3D 数据是带有大量噪声的，从信息论的角度看，仅仅基于 SMILES 的模型是无法超过 3D 模型的。
- LLM 的优势是可交互性强，能够通过自然语言和用户交流，即 open vocabulary。适合做公司前台，预处理用户复杂需求。
在清楚 LLM 定位后，笔者认为，当前做得比较好的一项作品是 2023_nature machine intelligence_moleculeSTM 。
这篇文章为大家展示了很科幻的应用场景：LLM+分子编辑任务
用户输入一个分子，并添加对其描述，模型吐出修改后的分子，使其更符合需求。如下图所示：
在了解 moleculeSTM 后，相信大家都跃跃欲试，想要体验一下。这篇文章就记录下本人在安装过程中踩的坑：
1. 安装 apex
纵观 moleculeSTM 的多项依赖，apex 是安装难度最大的一个。apex 本身是 nvidia 开发的，用于加速神经网络特定模块。在原文档里，我们需要 clone 作者的分支并进行安装：
git clone https://github.com/chao1224/apex.git
cd apex
pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./
其中难点在于，apex 要求环境中需要有 nvidia-cudakit，即
which nvcc
需要有响应。并且，pytorch 的版本需要跟 cuda 版本一致。例如，机器上的 cuda 型号是 12.1，这里只能安装在 cuda=12.1 下编译的 pytorch：
pip install -i http://mirrors.aliyun.com/pypi/simple/ --trusted-host mirrors.aliyun.com torch==2.0.0+cu121 -f https://download.pytorch.org/whl/cu121/torch_stable.html
在进行完这两项后，按照原文档可以将 apex 准备好。
2. 安装其他
基本上无需在意版本号，相对难度较低，按照原文档顺序安装即可。可能会缺一些东西，但基本上都很好解决，谷歌一下即可。一般来说，还需要再安装 nvidia-cutlass 库
3. 准备模型文件
首先，我们需要将 huggingface 上所有相关文件下载下来。一共是两个仓库，一个仓库是装 dataset 的，另一个仓库是装模型 checkpoint 的。
Dataset:
https://huggingface.co/datasets/chao1224/MoleculeSTM
Checkpoint:
https://huggingface.co/chao1224/MoleculeSTM/tree/main
如果网不好，可以淘宝找代下。
注意：
原作者在 issue 中指出，要想复现论文结果，我们需要载入特定模型：pretrained_MoleculeSTM/SciBERT-Graph-3e-5-1-1e-4-1-EBM_NCE-0.1-32-32
https://github.com/chao1224/MoleculeSTM/issues/16
这些模型可以在 hugging face 中找到。
准备好 moleculeSTM 相关文件后，我们还需要准备另外两份：
- 准备 pretrained_MegaMolBART，可以按照作者提供的脚本：https://github.com/chao1224/MoleculeSTM/blob/main/download_MegaMolBART.sh。但该脚本对我并不 work，随后是直接访问脚本中指向的谷歌链接，可以从那个地方下载：https://drive.usercontent.google.com/download?id=1O5u8b_n93HOrsjN1aezq6NhZojh-6dEe&export=download
- 准备 scibert。其实，如果开发机联网性能好的话，我们可以直接从 hugging face 中导入。但一般国内用户链接 hugging face 都比较困难。所以我这边还是选择直接下载好。原代码方法是：https://github.com/chao1224/MoleculeSTM/blob/9ba1584489369cd18bdc18d197e3bfdac6145d68/MoleculeSTM/downstream_molecule_edit_utils.py#L160。提前下载好，就是将此处的这个仓库下载到本地，并将这俩行中的地址进行替换，hugging face 地址在 https://huggingface.co/allenai/scibert_scivocab_uncased/tree/main
4. 开始跑
我们需要正确设定模型需要链接的位置。此处，我们以
https://github.com/chao1224/MoleculeSTM/blob/main/demos/demo_downstream_zero_shot_molecule_edit.ipynb
为例。我们需要修改所有改脚本中涉及文件路径的地方，将其改为，我们下载好的模型&数据目录。
最最最最后，可能还是跑不起来。
主要原因是，我们 load 的模型是在旧的 torch 框架下训练的。然而，由于第一条，我们开发机的 cuda 版本限制了 torch 版本。所以，我们只能将起冲突的地方，相关源码进行修改，使其适配新版本的 torch。以及，部分地方 load 的时候会报错：https://github.com/chao1224/MoleculeSTM/blob/9ba1584489369cd18bdc18d197e3bfdac6145d68/MoleculeSTM/downstream_molecule_edit_utils.py#L167
将其改为：
text_model.load_state_dict(state_dict, strict=False)
即可。
最终在经过长达一周的不断调试后，我在 cuda=12.1, torch=2.0.0 的开发机中成功跑通了 moleculeSTM，如下：
看到运行结果那一刻，我觉得所有的努力都是值得的，至少手里有了通往 LLM 时代的船票。
