TCRT5 pre-trained（昇腾 NPU 适配）

源仓库：dkarthikeyan1/tcrt5_pre_tcrdb（HuggingFace） 任务：T 细胞受体（TCR / CDR3β）序列条件生成（pMHC → TCR，seq2seq）

模型简介

TCRT5 是基于 T5 架构、针对 T 细胞受体（TCR）序列条件生成的预训练模型。 给定目标肽-MHC（pMHC）输入，可条件生成对应的 TCR（CDR3β）序列； 预训练于约 1400 万 CDR3β 序列（TCRdb）与约 78 万 pMHC 对（IEDB）的掩码片段重建任务。

架构

T5ForConditionalGeneration（config.json）：

• 参数量：~42M（41,990,144，safetensors 头实测）
• 编码器-解码器：10 层 encoder + 10 层 decoder
• 注意力头：16 头，d_model=256，d_kv=64
• 前馈维度：d_ff=1024，relu 激活
• 词表大小：128（含 [PMHC]/[SEP]/[EOS]/[PAD]/[SOS]/[UNK] 等特殊 token）

环境要求

pip install -r requirements.txt

快速开始

# 默认用例（pMHC → TCR 生成，自动选择 NPU）
python3 inference.py --device npu:0

# 指定输入
python3 inference.py --input-text "[PMHC]KLGGALQAK[SEP]YFAMYQENVAQTDVDTLYIIYRDYTWAELAYTWY[EOS]" --device npu:0

inference.py 加载本仓 weights/ 的 T5ForConditionalGeneration 权重，执行 pMHC→TCR 条件生成，验证输出形状与数值。

验证结果（昇腾 NPU 实测）

• 设备：npu:0（Ascend 910，CANN 8.5.1，PyTorch 2.9.0 / torch_npu 2.9.0.post1）
• 输入：[PMHC]KLGGALQAK[SEP]YFAMYQENVAQTDVDTLYIIYRDYTWAELAYTWY[EOS]（编码后 46 tokens）
• 输出：beam=4 生成的 CDR3β token 序列
• 数值检查：输出 token 全为合法词表索引（无 NaN/越界），NPU_INFERENCE_PASS=True

仓库结构

inference.py        推理脚本（必选交付）
readme.md           部署说明文档（必选交付）
requirements.txt    环境依赖清单
weights/            模型权重（model.safetensors 168MB + tokenizer + config）

许可证

cc-by-nc-sa-4.0（非商用，署名-相同方式共享）