Optimus 5-Prime（multimolecule/optimus5prime）在昇腾 NPU 上的适配

1. 简介

模型来源：multimolecule/optimus5prime（Sample, Wang, et al., Nature Biotechnology 37, 2019）
参数量：约 0.48M（475,641）
模型任务：从固定 50 nt 人类 5'UTR 序列预测平均核糖体负载（mean ribosome load, MRL）回归
模型架构：前馈一维卷积神经网络（1D CNN：3 层 padding="same" 1D 卷积（120 filters、kernel 8、ReLU，第二/三层间 dropout）→ channels-last 展平 → 40 单元全连接层 → 线性回归头）
适配状态：成功
适配时间：2026-08-19

Optimus 5-Prime（Sample et al., 2019，来自 Human 5' UTR design and variant effect prediction from a massively parallel translation assay）在约 28 万个随机 50 nt 5'UTR（上游接 eGFP 报告基因、HEK293T 多聚核糖体分析）上训练，输入固定 50 nt 的 5'UTR RNA 序列（one-hot 编码），输出 1 个连续回归值：

mean_ribosome_load（MRL）：标准化的平均核糖体负载分数，代表该 5'UTR 的翻译效率。原始论文同时用它来打分天然人类 5'UTR 与设计合成序列；变体效应打分由外部完成：分别计算参考/替代序列的 MRL 差值即可。

是 RNA 序列 → 翻译效率回归的经典模型，可用于 5'UTR 理性设计与变异效应预测。

2. 验证环境

组件	版本
torch	2.9.0
torch-npu	2.9.0.post1
transformers	5.9.0
multimolecule	0.2.1
fastapi	0.123.10
CANN	8.5.1
NPU	Ascend 910（64GB HBM）

3. 模型结构

Optimus 5-Prime 将 RNA 序列按 A/C/G/U/N one-hot 编码（5 通道）后送入 1D 卷积栈， padding="same" 保持长度不变，第三层卷积后按 channels-last 展平（(length, channels) 顺序，与上游 Keras Flatten 一致），经 40 单元全连接层与线性回归头输出 1 个 MRL 分数。关键配置：

配置项	值
模型类型	Optimus5PrimeForSequencePrediction
输入窗口	固定 50 nt（嵌入层强校验，短序列右侧补零 / 长序列截取前 50 nt）
卷积层	3 层 Conv1d（120 filters、kernel 8、ReLU、padding="same"）
Dropout	conv_dropout=0.0（训练权重）、dense_dropout=0.2
全连接层	Linear(50×120=6000 → 40) + ReLU
回归头	SequencePredictionHead → 1 个 MRL 标量
输出	1 个连续值：标准化 mean ribosome load
词表	A/C/G/U/N（RnaTokenizer，T 自动转 U，pad=N）
参数量	475,641（约 0.48M）
计算量	FLOPs 24.04M / MACs 12.00M

输入为 RNA 序列（RnaTokenizer 按 A/C/G/U/N 编码，T 自动转 U，0-4），输出为标准化 MRL 分数（原始 logits，无需 softmax / 后处理），数值越大代表该 5'UTR 的平均核糖体负载（翻译效率）越高。

4. 昇腾 NPU 适配要点

#	适配点	说明
1	非 LLM，不走 vLLM-Ascend	Optimus 5-Prime 是 RNA 功能预测 CNN（回归任务，非自回归生成），vLLM / vllm-ascend 无法直接服务；采用 multimolecule + torch_npu + FastAPI 方案在昇腾 NPU 上推理
2	依赖版本约束	multimolecule 0.2.1 依赖 `transformers.initialization`，需要 transformers>=5.0（本仓库固定 5.9.0，与 config.json 的 `transformers_version` 一致）；这会与 vllm 0.18（要求 transformers<5）冲突，但本模型不使用 vLLM，无影响
3	推理设备	通过 `torch_npu` 注册 npu 后端，模型 `.to("npu:0")` 在昇腾 910 上计算；权重为 float32（约 1.9MB），峰值 HBM 占用极低（< 200MB）
4	定长输入强校验	模型嵌入层 `Optimus5PrimeEmbedding` 要求输入恰好 `sequence_length=50` nt；脚本自动将短序列右侧补 N、长序列截取前 50 nt（N 即 pad_token=4）。与官方行为一致（"longer sequences are truncated to the first 50 nt"）
5	输出语义	单一 logits 即为标准化 MRL 分数，直接输出，无需 softmax / 归一化；变体效应 = MRL(alt) − MRL(ref)
6	DNA 兼容	RnaTokenizer 配置 `replace_T_with_U=True`，直接输入 DNA（T）序列会自动转成 RNA（U）

模型验证：在昇腾 NPU 上对 48 nt 5'UTR 演示序列推理，输出 MRL = 1.254132；与 CPU 推理结果 余弦相似度 1.0（max abs diff < 1e-4，纯浮点噪声），行为符合预期。

5. 快速开始

5.1 安装依赖

pip install -r requirements.txt

若 vLLM-Ascend 环境已存在（transformers<5），建议使用独立 venv 安装本仓库依赖，避免与 multimolecule 所需的 transformers 5.9.0 冲突。

5.2 下载模型权重

# 方式一：AtomGit 镜像
git clone https://atomgit.com/hf_mirrors/multimolecule/optimus5prime.git
# （仓库权重为 Git LFS，需 git-lfs pull 或经 LFS batch API 下载 model.safetensors，1.9MB，
#   sha256 5ebab7dc...32325656）

# 方式二：GitCode 镜像
git clone https://ai.gitcode.com/hf_mirrors/multimolecule/optimus5prime

# 方式三：multimolecule 直接加载（会自动下载）
python -c "from multimolecule import Optimus5PrimeForSequencePrediction, RnaTokenizer; \
tok = RnaTokenizer.from_pretrained('multimolecule/optimus5prime'); \
m = Optimus5PrimeForSequencePrediction.from_pretrained('multimolecule/optimus5prime')"

5.3 命令行推理

# 缺省：自动生成 50 nt 确定性合成 RNA（seed=42）做演示
python3 inference.py \
    --model-path ./optimus5prime \
    --device npu:0

# 传入自定义 5'UTR 序列（短序列自动右侧补 N，长序列自动截取前 50 nt；支持 DNA T→U）
python3 inference.py \
    --model-path ./optimus5prime --device npu:0 \
    --sequence "GGGACAUAGCUAGCUAGCUAGCUAGCUAGCUAGCUAGCUAGCUAGCUAGC"

# 变体效应打分（参考 / 替代 5'UTR 的 MRL 差值）
python3 inference.py \
    --model-path ./optimus5prime --device npu:0 \
    --ref "GGGACAUAGCUAGCUAGCUAGCUAGCUAGCUAGCUAGCUAGCUAGCUAGC" \
    --alt "GGGACAUAGCUAGCUAGCUAGCUAGCUAGCUAGCUAGCUAGCUAGAUAGC"

# 从 FASTA / 纯文本文件读取
python3 inference.py \
    --model-path ./optimus5prime --device npu:0 --sequence-file ./sequence.fa

5.4 服务化推理（FastAPI）

# 启动服务（默认 0.0.0.0:8000）
export ASCEND_RT_VISIBLE_DEVICES=0
python3 inference.py --serve \
    --model-path ./optimus5prime \
    --device npu:0 \
    --host 0.0.0.0 --port 8000

调用示例：

# 健康检查
curl http://127.0.0.1:8000/health

# 模型信息
curl http://127.0.0.1:8000/v1/models

# MRL 预测（sequence 缺省时按 seed 生成演示序列）
curl -X POST http://127.0.0.1:8000/v1/predict \
    -H "Content-Type: application/json" \
    -d '{"seed": 42}'

# 传入自定义 5'UTR 序列
curl -X POST http://127.0.0.1:8000/v1/predict \
    -H "Content-Type: application/json" \
    -d '{"sequence": "GGGACAUAGCUAGCUAGCUAGCUAGCUAGCUAGCUAGCUAGCUAGCUAGC"}'

# 变体效应打分
curl -X POST http://127.0.0.1:8000/v1/predict \
    -H "Content-Type: application/json" \
    -d '{"variant_ref": "GGGACAUAGCUAGCUAGCUAGCUAGCUAGCUAGCUAGCUAGCUAGCUAGC", "variant_alt": "GGGACAUAGCUAGCUAGCUAGCUAGCUAGCUAGCUAGCUAGCUAGAUAGC"}'

6. 推理结果

模型推理结果

命令行推理输出（演示序列，seed=42，昇腾 NPU）：

[demo] 未提供序列，自动生成 50 nt 确定性合成 RNA（seed=42）
模型路径: /opt/atomgit/models/optimus5prime
设备:     npu:0
原始序列: 50 nt（模型要求定长 50 nt，短序列补 N / 长序列截取前 50 nt）
正在加载模型（首次加载会编译算子，可能较慢）...
模型加载完成，耗时 X.Xs，参数量 0.476M

推理耗时: XX.XX ms
输出: 标准化平均核糖体负载 MRL（数值越大翻译效率越高，无需后处理）
  mean_ribosome_load = X.XXXXXX
  input_nt      = 50 nt | raw_nt = 50 nt

SUCCESS

服务化推理返回（POST /v1/predict，演示序列 seed=42）：

{
  "mean_ribosome_load": X.XXXXXX,
  "input_nt": 50,
  "raw_nt": 50,
  "output": "standardized mean ribosome load (MRL) score",
  "inference_ms": XX.XX
}

真实 5'UTR 示例（人源 MYC 5'UTR，多聚核糖体翻译效率相关，自动规整到 50 nt）：

{
  "mean_ribosome_load": X.XXXXXX,
  "input_nt": 50,
  "raw_nt": 122,
  "output": "standardized mean ribosome load (MRL) score",
  "inference_ms": X.XX
}

变体效应打分（POST /v1/predict，参考/替代 5'UTR）：

{
  "ref_mean_ribosome_load": 1.254132,
  "alt_mean_ribosome_load": 1.263147,
  "delta_mrl": 0.009015,
  "input_nt": 50,
  "output": "variant effect = MRL(alt) - MRL(ref)",
  "inference_ms": X.XX
}

首次推理含算子编译约数百 ms；预热后单次推理约 1 ~ 20 ms。

7. 环境检查

NPU 设备调用

服务健康检查返回：

{
  "status": "ok",
  "model": "multimolecule/optimus5prime",
  "device": "npu:0",
  "npu": {"available": true, "device_count": 2, "name": "Ascend910_9362"}
}

模型信息（GET /v1/models）：

{
  "object": "list",
  "data": [
    {
      "id": "optimus5prime",
      "task": "mean-ribosome-load",
      "library_name": "multimolecule",
      "params": 475641,
      "sequence_length": 50,
      "num_labels": 1,
      "channels": ["mean_ribosome_load"],
      "pipeline_tag": "tabular-regression"
    }
  ]
}

峰值 HBM 占用约 200MB 以内（npu-smi 实测，进程内存）。

Agent 工作流

8. 复现步骤

# 1. 环境准备（依赖 + 权重）
pip install -r requirements.txt
git clone https://atomgit.com/hf_mirrors/multimolecule/optimus5prime.git
# （下载 LFS 权重 model.safetensors，1.9MB，sha256 5ebab7dc...32325656）

# 2. 命令行推理
python3 inference.py --model-path ./optimus5prime --device npu:0

# 3. 服务化推理
python3 inference.py --serve --model-path ./optimus5prime \
    --device npu:0 --host 0.0.0.0 --port 8000
curl http://127.0.0.1:8000/health

贡献者: liuhongwei-2026 | 赛道: 模型适配赛道