Protein Embedding Visualizer

基于 Hugging Face 蛋白质语言模型生成 protein-level embeddings，并输出可追溯的 t-SNE 与 UMAP 可视化结果。输入、模型、推理、输出和绘图参数统一由 config.yaml 管理。

功能

• 读取自定义 CSV 序列列和家族列，不修改原始输入；
• 支持 ESMC、ESM-2、ProtT5 以及自定义 Hugging Face checkpoint；
• 支持 mean、max、CLS pooling 和可选序列截断；
• 生成 embedding 时用 tqdm 按已完成序列数显示进度；
• 以 embeddings.npy、records.csv 和 JSON 元数据保存可追溯结果；
• 支持 t-SNE、UMAP、家族配色和 PNG/SVG 输出。

项目结构

.
├── config.yaml                         # Example configuration
├── data/example_proteins.csv           # Example input
├── protein_embedding_visualizer/
│   ├── cli.py                          # Command-line workflow
│   ├── config.py                       # Configuration models and validation
│   ├── data_input.py                   # CSV loading and validation
│   ├── embedding.py                    # Model inference and artifact storage
│   ├── io_utils.py                     # Shared JSON and artifact helpers
│   └── visualization.py                # t-SNE, UMAP, and plotting
├── environment.yml                     # Conda environment and pinned dependencies
└── pyproject.toml                      # Package metadata and dependencies

环境与安装

请先使用 environment.yml 创建并激活 Visualize_prtebd Conda 环境：

conda env create -f environment.yml
conda activate Visualize_prtebd

项目依赖包含固定 commit 的 Biohub Transformers fork，因此首次安装需要访问 GitHub。大模型权重由 Hugging Face 在首次推理时下载，所需磁盘和显存取决于 所选 checkpoint。

快速开始

运行完整的 embedding 与可视化工作流：

python -m protein_embedding_visualizer --config config.yaml

输入格式

CSV 必须包含序列列和家族列：

Sequence,Family
MKTAYIAKQRQISFVKSHFSRQ,Example_family_A
GAVLILKKKGHHEAELKPLAQ,Example_family_B

列名、分隔符和编码均可配置：

input:
  path: example_data/my_proteins.csv
  sequence_column: protein_seq
  family_column: group
  delimiter: ","
  encoding: utf-8-sig

默认会移除序列空白、转换为大写并验证蛋白质字符。解析错误会包含原始 CSV 行号。空字符串、缺失字段和非法字符不会被静默替换。

配置

所有可变参数位于 config.yaml，主要分为五部分：

配置段	用途
input	CSV 路径、列名、编码和序列验证
model	模型家族、规模、checkpoint、revision 和 offload
embedding	batch size、token 长度、pooling、设备和精度
output	输出目录和覆盖策略
visualization	方法、配色、图形格式和降维参数

model.model_id 非空时优先于内置 size 别名。默认 truncate: false，超长序列 会明确报错；只有确认允许丢失尾部残基时才应设为 true。

支持的内置规模：

家族	size 可选值
ESMC	300m, 600m, 6b
ESM-2	8m, 35m, 150m, 650m, 3b, 15b
ProtT5	base_uniref50, xl_bfd, xl_uniref50, xl_half_uniref50_encoder, xxl_bfd, xxl_uniref50

输出

output.directory 中会生成：

文件	内容
embeddings.npy	N × D embedding 数组
records.csv	embedding 行号、序列、家族和原始 CSV 行号
metadata.json	模型、revision、参数量和配置

Embedding 产物完成后才会替换目标文件；中断后不会将临时文件误认为完整结果。 目标文件已存在时默认拒绝覆盖，确认需要重新生成后设置 output.overwrite: true。

可视化结果默认位于 output.directory/visualizations：

• tsne.png、tsne.svg、umap.png 和 umap.svg；
• tsne_coordinates.csv 和 umap_coordinates.csv；
• visualization_metadata.json。

固定随机种子可以复现同一软件环境中的布局，但二维簇间距离不应被解释为蛋白质 家族关系的定量证据。

大模型与 Flash Attention 2

大模型通常需要多 GPU 或 CPU/disk offload：

model:
  device_map: auto
  max_memory:
    0: 20GiB
    cpu: 64GiB
  offload_folder: outputs/offload
embedding:
  dtype: bfloat16
  batch_size: 1

Flash Attention 2 通常用于支持 CUDA 的 Linux/WSL 环境：

python -m pip install -e ".[flash-attention]"

同时设置：

model:
  attn_implementation: flash_attention_2
embedding:
  device: cuda
  dtype: bfloat16

程序会在加载模型前校验 CUDA、半精度和 flash_attn。Windows 本机未安装 flash_attn 时会明确报错，默认 attention 路径不受影响。

参考

• Biohub ESM/ESMC
• Meta ESM-2
• Rostlab ProtT5 encoder