mambaphase

Mambaphase

用于蛋白质液–液相分离（LLPS）相关预测的推理代码。输入蛋白质 FASTA 序列，使用 ESM2 提取序列特征，再运行 LLPS、client、pH 和盐浓度分类头；独立的 scaffold 分类头可复用生成的特征。

该实现基于 Mambaphase，提供明确的命令行入口、模型加载检查和结构化输出。

推理模式

ESM 特征取自最后一层的首个 BOS/CLS token，每条序列对应一个 2560 维向量。base 是与已检查的分类头特征流程匹配的推荐起点；finetuned 使用不同的模型权重，两种模式的预测不能视为等价。

下载

发布仓库：NDWANG/mambaphase。

使用 Hugging Face CLI 下载仓库。建议把下载工具安装在独立环境，避免它的依赖影响下面固定版本的推理环境：

python3.10 -m venv .hf-cli
.hf-cli/bin/python -m pip install --upgrade huggingface_hub
.hf-cli/bin/hf download NDWANG/mambaphase --local-dir Mambaphase
cd Mambaphase

完整下载包含源码、分类头和两套 ESM 资产。若只使用基础 ESM 模式，可在下载命令中加上 --exclude "weights/esm_finetuned/*"；其余相对目录保持不变。下载后请核对下面的模型目录，推理程序本身不会联网补齐缺失文件。

安装

建议使用 Python 3.10 和独立虚拟环境（以下命令适用于 Linux/macOS shell）：

python3.10 -m venv .venv
source .venv/bin/activate
python -m pip install -r requirements-runtime.txt

请根据计算设备选择兼容的 PyTorch 构建。predict_base.py 要求 CUDA；predict_finetuned.py 支持 CUDA 或 CPU；predict_scaffold.py 使用 CPU。

finetuned 模式还需要提供 esm 模块的 fair-esm 包：

python -m pip install fair-esm

requirements-runtime.txt 记录了已检查的依赖版本组合，并非经过全新安装验证的锁文件。已验证的 fair-esm 安装自报版本为 2.0.1，该版本在检查的 PyPI 源中不可得；上面的安装命令用于取得公开发布的 fair-esm 包，不能保证重建相同环境。不同平台、CUDA 或 fair-esm 版本仍需检查兼容性。项目中的序列模块使用 PyTorch 与 einops 实现，无需安装 mamba-ssm。

模型资产

模型文件应保留以下相对目录。也可以在已有代码目录中使用 hf download NDWANG/mambaphase --local-dir . --include "weights/*" "model/*" 单独补充已发布的模型文件：

weights/
├── esm_base/
│   ├── config.json
│   ├── tokenizer_config.json
│   ├── vocab.txt
│   ├── pytorch_model.bin.index.json
│   ├── pytorch_model-00001-of-00002.bin
│   └── pytorch_model-00002-of-00002.bin
└── esm_finetuned/
    ├── config.json
    ├── vocab.txt
    └── model.pth

model/
├── pspweight/best_model.pth
├── clientweight/best_model.pth
├── phweight/best_model.pth
├── saltweight/best_model.pth
└── scaffoldweight/best_model.pth

• esm_base 对应 Hugging Face 格式的 ESM2 t36 3B，需包含配置、tokenizer、分片索引及全部权重分片。也支持带 model.safetensors.index.json 的分片格式；当前入口不接受只有单个未分片权重文件的目录。
• esm_finetuned/model.pth 必须是与原生 fair-esm ESM2 架构匹配的完整 state_dict，不能直接替换为 Hugging Face 格式权重。
• 分类头须使用与其网络结构匹配的权重。client 与 scaffold 是不同任务，不应交换它们的 checkpoint。

ESM 资产位置可通过 --base-dir，或 --esm-dir 与 --esm-weights 指定；scaffold 权重位置可通过 --weights 指定。

快速开始

准备 FASTA 文件，例如：

>protein_1
MKTAYIAKQRQISFVKSHFSRQ

基础 ESM 模式

python predict_base.py \
  --fasta input.fasta \
  --output-dir outputs/base \
  --device cuda

默认模式为 cuda_autocast_fp16：模型参数保持 FP32，ESM 前向计算启用 CUDA 混合精度。可用 --mode fp32 运行纯 FP32，或用 --mode both 分别生成两套结果。

微调 ESM 模式

python predict_finetuned.py \
  --fasta input.fasta \
  --output-dir outputs/finetuned \
  --device cuda

默认使用 FP32。--precision fp16 会显式转换模型精度，仅支持 CUDA；CPU 推理应使用 FP32。

Scaffold 分类

在特征提取完成后，可单独运行 scaffold 分类头：

python predict_scaffold.py \
  --features outputs/base/embeddings_cuda_autocast_fp16.npz \
  --output outputs/base/scaffold.json

该入口不重复加载 ESM。输入 NPZ 必须包含一一对应的 Unicode 字符串数组 ids、sequences，以及形状为 (N, 2560) 的浮点数组 embeddings。

查看完整参数：

python predict_base.py --help
python predict_finetuned.py --help
python predict_scaffold.py --help

输入要求

• FASTA 标题以 > 开始，第一个字段作为唯一序列 ID；序列不能为空，氨基酸字母必须大写。
• 两个 ESM 入口支持 ACDEFGHIKLMNPQRSTVWYU。scaffold 入口仅支持 20 种标准氨基酸，不接受 U。
• 不接受 X、B、Z、终止符 * 或缺口符；不会自动替换不支持的字符。
• 默认长度上限为 1024 aa，超限会报错，不会静默截断。--max-length 可调整校验上限，但不保证更长序列的资源需求或预测有效性。
• 每次运行请使用新的输出目录或文件，已有目标结果不会被覆盖。

输出格式

基础与微调模式的预测 JSON 在 predictions 中保存逐条记录，每条记录的 heads 包含 psp、client、ph、salt，各头提供类别概率 probabilities 与最大概率对应的 class_index。

client 的类别为 0 = nonLLPS、1 = client。独立 scaffold 输出的类别为 0 = client、1 = scaffold，其逐条概率直接位于预测记录中。PSP、pH 和盐浓度保留原始类别索引；这些索引不代表实际测量的 pH 数值或盐浓度。

ESM 模式的 metadata.json 记录参数、依赖版本、模型校验值、运行耗时和状态。成功结束时 status 为 end_to_end_passed；scaffold 输出的成功状态为 prediction_passed。模型加载或数值检查失败时，程序会报错，不会使用零向量代替失败的特征。

使用范围

本项目提供推理实现，不将程序运行成功视为准确率或生物学有效性的证明。不同 ESM 权重会改变特征分布，分类头与特征模型应配套使用；发布的分类头对原生微调特征的适配程度尚未独立验证。

PSP 发布权重对应的序列分支实际包含零个 Mamba 块，当前实现按该权重结构严格加载。模型概率未经独立校准，不应直接解释为实验成功率、最适 pH 或最适盐浓度。

许可证

代码采用 MIT License，保留上游版权声明。使用第三方模型资产时，还应遵守相应资产的许可证。