nach0_flagos

## 模型介绍 nach0 由 Insilico Medicine 的研究团队提出，旨在解决化学分子语言与自然语言两类模态难以统一建模、单一领域任务需要各自训练专用模型的问题。它将分子 SMILES 与科学文献文本映射到同一序列空间，用一套编码器-解码器结构同时承担化学与语言两类任务。 **nach0: Multimodal Natural and Chemical Languages Foundation Model** <https://arxiv.org/abs/2311.12410> --- ## 模型描述 nach0 是一个面向化学与自然语言的多模态、多任务编码器-解码器模型，其核心结构是 **T5（Text-to-Text Transfer Transformer）**，通过科学文献、专利与分子 SMILES 的大规模预训练，再经指令微调（instruction tuning）获得跨领域能力；本仓库适配的 `nach0_base` 为 T5-base 规模，共 **222,935,808** 参数（12 层编码器 + 12 层解码器，`d_model=768`，`d_ff=2048`，12 个注意力头，词表 32170）。 模型的输入输出都是「文本序列」：化学结构先被拆成原子级 token 并包装为 `<sm_C>`、`<sm_1>` 这类特殊符号，与自然语言指令拼接后送入模型，生成的序列再还原为 SMILES。这种统一表示让反应产物预测、逆合成、分子描述生成等任务可以用同一套接口完成。 ```text 自然语言指令 + <sm_*> 分子序列 ──► T5 Encoder ──► T5 Decoder ──► <sm_*> 序列 ──► 还原为 SMILES ``` --- ## 适用场景 | 场景 | 说明 | | :--- | :--- | | **反应产物预测 / 逆合成等化学任务推理** | 给定反应物与试剂描述，生成可能的产物 SMILES；亦可完成分子描述生成、性质问答等跨模态任务 | | **本地快速验证** | 使用单个样例输入检查模型加载、分词、推理链路是否正常，确认适配无误 | | **ModelScope / OneCode 运行** | 作为独立模型包下载后直接安装依赖并运行脚本 | | **算子性能分析** | 使用 `prof.py` + `run_nach0.py` 采集单样本前向的算子调用与耗时分布，用于国产芯片算子覆盖与适配分析 | | **多卡训练** | 通过 `torchrun` 启动多进程训练（本仓库聚焦推理适配，未包含训练脚本） | --- ## 使用说明 ### 1. OneCode 使用 可通过 OneCode 在线环境体验智能化一键式 AI4S 编程： [点击体验智能化一键式 AI4S 编程](https://www.onescience.ai/) <!-- 待确认：OneCode 体验页的具体 URL，请按平台实际地址替换 --> ### 2. 手动安装使用 以下命令均假定在 nach0 项目根目录执行。 #### 硬件要求 - 推理可在 CPU 或 DCU 上运行；本仓库的适配验证在 DCU 容器的 CPU 模式下完成。训练建议使用 DCU 或 GPU。 - 多卡训练使用 NCCL 后端，请确保设备驱动、通信库和 PyTorch 版本匹配。 - DCU 用户需要预先安装 DTK，建议使用 DTK 25.04.2 以上版本或与当前集群匹配的 OneScience 推荐版本。 - 单样本前向（`1×202` token 输入、T5-base 规模）在 CPU 上约需 2 秒，内存占用约 2 GB。 #### 下载模型包 ```bash modelscope download --model FrontierScientificModels/nach0 --local_dir ./nach0 cd nach0 ``` 若使用 Hugging Face 权重，可直接下载 `insilicomedicine/nach0_base` 并放到 `nach0_base/` 目录下，所需文件如下： | 文件 | 大小（字节） | | :--- | ---: | | `pytorch_model.bin` | 990,666,933 | | `tokenizer.json` | 2,435,024 | | `tokenizer_config.json` | 2,349 | | `special_tokens_map.json` | 2,297 | | `config.json` | 1,540 | | `generation_config.json` | 142 | #### 安装运行环境 **DCU 环境** ```bash # 请先激活 DTK 及 CONDA conda create -n onescience311 python=3.11 -y conda activate onescience311 # 支持 uv 安装 pip install onescience[bio-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai ``` **GPU 环境** 暂未提供。药物研发领域（含生命科学）当前仅提供 DCU 版本，GPU 版本后续按需补充。 #### 运行前必须设置的环境变量 在 DCU 容器中，`torch` 自带的厂商设备后端扩展可能加载失败，且 CUDA 探测接口在无 GPU 环境下会触发段错误，因此运行任何脚本前需设置： ```bash export TORCH_DEVICE_BACKEND_AUTOLOAD=0 # 关闭厂商设备后端自动加载 export CUDA_VISIBLE_DEVICES="" # 屏蔽 CUDA 探测（CPU 推理 / 算子采集时必需） ``` #### 快速验证 使用单个样例检查模型加载、分词与推理链路： ```bash export TORCH_DEVICE_BACKEND_AUTOLOAD=0 python test.py --quick --device cpu --model-path ./nach0_base --max-length 128 ``` 该脚本依次检查运行环境、特殊符号预处理（`add_special_symbols` / `clean_output_sequence`）、依赖可用性与官方示例推理，并输出逐项 PASS / FAIL 汇总。 #### 推理 ```bash export TORCH_DEVICE_BACKEND_AUTOLOAD=0 # 独立实现，不依赖 transformers 的建模层 python nach0_standalone.py --model-path ./nach0_base --quick # 按论文的采样策略生成（do_sample + top_k + top_p） python nach0_standalone.py --model-path ./nach0_base --quick --do-sample --repetition-penalty 1.2 ``` 官方示例（反应产物预测）： ```text 输入： Given the following reactants and reagents, please provide a possible product. CCN(CC)CC.CCN=C=NCCCN(C)C.CN(C)C=O.Cl.NC1=CC=C(Cl)C=C1N.O.O=C(O)CCCCCNC(=O)C=C1C2=CC=CC=C2C2=CC=CC=C12.OC1=CC=CC2=C1N=NN2.[Cl-].[Na+] 参考输出： NC1=CC=C(Cl)C=C1NC(=O)CCCCCNC(=O)C=C1C2=CC=CC=C2C2=CC=CC=C12 ``` #### 算子性能分析 采集单样本前向（可选前向 + 反向）的算子调用与耗时分布： ```bash export TORCH_DEVICE_BACKEND_AUTOLOAD=0 # 推理：仅前向 CUDA_VISIBLE_DEVICES="" python prof.py run_nach0.py # 训练：前向 + 反向 CUDA_VISIBLE_DEVICES="" NACH0_PROF_MODE=train python prof.py run_nach0.py ``` 运行后在当前目录生成 `op_profiles/<run_id>/`，其中包含 `top_operators.md`（TOP-N 算子耗时表）与 `operator_profile.json`（完整统计）。采集范围是「第一个 `nn.Module` 前向到脚本结束」，因此每个 `run_*.py` 只执行一个样本的一次前向（或前向 + 反向）。 实测输入规格（用于归档命名）： ```text model_dtypes : ['torch.float32'] input_dtype : torch.int64 input_ids.shape : (1, 202) decoder_input.shape : (1, 60) hidden_states.shape : (1, 202, 768) ``` 对应产物目录名：`nach0_base_推理_1_float32_1x202x768`。 ```bash cd op_profiles tar -czf "nach0_base_推理_1_float32_1x202x768.tar.gz" "nach0_base_推理_1_float32_1x202x768" ``` #### 训练权重 本仓库聚焦推理适配，不包含训练脚本。推理所用权重来自 Hugging Face `insilicomedicine/nach0_base`，下载后置于 `nach0_base/` 目录即可。`nach0_large` 版本的推理请参考原论文与 NVIDIA NeMo 项目。 #### 本仓库脚本一览 | 文件 | 作用 | | :--- | :--- | | `test.py` | 完整自测：环境、预处理、依赖、官方示例、批量 / 确定性 / 性能 | | `nach0_standalone.py` | 独立 T5 实现（仅依赖 PyTorch），用于绕开环境中 transformers 建模层的兼容性问题 | | `run_nach0.py` | 算子采集用运行脚本：单样本、单次前向（可选反向） | | `prof.py` / `torch_op_profiler.py` | 算子采集入口与实现（官方提供） | | `check_env.py` / `diag_*.py` | 环境自检与故障诊断（可选） | --- ## OneScience 官方信息 | 平台 | OneScience 主仓库 | Skills 仓库 | | :--- | :--- | :--- | | Gitee | <https://gitee.com/onescience-ai/onescience> | <https://gitee.com/onescience-ai/oneskills> | | GitHub | <https://github.com/onescience-ai/OneScience> | <https://github.com/onescience-ai/oneskills> | --- ## 引用与许可证 **许可证** 本仓库为 **nach0** 原始论文的 OneScience 适配复现版本，参考 Hugging Face 上的官方权重与实现完成： - 源地址：<https://huggingface.co/insilicomedicine/nach0_base> - 论文：<https://arxiv.org/abs/2311.12410> 原始仓库的许可与使用条款保持不变，请务必遵守： - 模型权重采用 **CC BY-NC 4.0**（<https://creativecommons.org/licenses/by-nc/4.0/>），**仅限研究用途**。 - 配套数据集采用 **CC BY 4.0**，**仅限非商业使用**。 - 我们强烈建议使用者遵循最高的伦理标准，包括在研究过程中保持公平、透明与责任；任何可能导致危害或对社会造成损害的使用方式均被严格禁止。 本仓库代码与文档仅为上述工作的适配复现，其使用仍应以原项目中的许可及使用条款为准。 **引用** 如果您使用了本仓库的内容，请同时引用原始论文： ```bibtex @article{D4SC00966E, author = "Livne, Micha and Miftahutdinov, Zulfat and Tutubalina, Elena and Kuznetsov, Maksim and Polykovskiy, Daniil and Brundyn, Annika and Jhunjhunwala, Aastha and Costa, Anthony and Aliper, Alex and Aspuru-Guzik, Alán and Zhavoronkov, Alex", title = "nach0: multimodal natural and chemical languages foundation model", journal = "Chem. Sci.", year = "2024", volume = "15", issue = "22", pages = "8380-8389", publisher = "The Royal Society of Chemistry", doi = "10.1039/D4SC00966E", url = "http://dx.doi.org/10.1039/D4SC00966E", } ```