安装 moleculeSTM 踩坑日记

“学习 LLM ，在大模型时代为自己存张船票”。
相信很多人都有这样的想法。那么，在 AI for science 领域，哪些 LLM 模型值得一试呢？
笔者认为：

1. LLM 直接预测 SMILES 性质 or 直接生成 SMILES 的技术路线是行不通的。因为 SMILES 相比高精度 3D 数据是带有大量噪声的，从信息论的角度看，仅仅基于 SMILES 的模型是无法超过 3D 模型的。
2. 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 相关文件后，我们还需要准备另外两份：

1. 准备 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
2. 准备 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 时代的船票。