kermt-embed

kermt-embed

从任何包含编码器的 KERMT 检查点中提取逐分子嵌入。 该技能负责协调工作流：验证检查点、验证 CSV、清理 SMILES、阻塞式启动运行程序，并返回逐读取结果的 .npy 文件。

硬件要求

• GPU：1 个（单 GPU）。
• 显存：使用默认 batch_size 64 时需 ≥ 4 GB。
• 磁盘：取决于输出大小——每个读取结果在 hidden 800 时，每 1k 个分子约需几 MB，因此 4 个读取结果合计每 1k 个分子约需 10–20 MB。此外，每次运行还会生成较小的 canonical_smiles.npy + validity.npy。
• 驱动 / CUDA：任何支持 CUDA 12.6 的主机。

输入

必需：

• --csv <path> — SMILES CSV。第一列为 smiles；其他列 会被忽略（不需要目标值）。

检查点（可选 — 若省略，则默认为已发布的模型）：

• --ckpt <path> — 任何包含编码器的检查点。Grover_base、cmim、 hybrid 和微调后的检查点均可接受。验证器只会拒绝 不含编码器的检查点。如果省略，该技能会询问是否下载 已发布的预训练 hybrid 模型 nvidia/NV-KERMT-70M-v2，并使用它生成嵌入—— 参见“解析并验证检查点”（工作流步骤 3）。
• --pretrained-release — 显式选择使用已发布模型，不显示交互式提示 （适用于非交互式 / Agent 运行）。不能与 --ckpt 同时使用。
• --model-dir <dir> — 下载的模型包保存位置（默认为 $KERMT_REPO/models/NV-KERMT-70M-v2/）。如果该目录中已有完整模型包， 则会复用，不会重新下载。

可选：

• --batch-size N — 覆盖配置的默认值（64）。
• --gpus 0 — 单 GPU ID（默认为 0）。
• --from-prepare <dir> — 跳过 prepare 步骤，复用 <dir> 中现有的 prepare_data.json。

工作流

令 $KERMT_REPO 为 kermt 仓库检出的路径。

1. 运行前检查：容器 + 系统探测。

   $KERMT_REPO/agent/scripts/kermt_container.sh check_system
   

2. 计算运行目录。

   RUN_DIR=$KERMT_REPO/runs/embed_$(date -u +%Y-%m-%dT%H-%M-%SZ)
   

3. 解析并验证检查点。

   解析 — 仅当省略了 --ckpt 时执行。 默认为已发布的 预训练 hybrid 模型 nvidia/NV-KERMT-70M-v2：

   • 同意确认。 除非已传入 --pretrained-release，否则询问用户： “未提供检查点 — 是否下载已发布模型 nvidia/NV-KERMT-70M-v2 （NVIDIA Open Model License，https://huggingface.co/nvidia/NV-KERMT-70M-v2） 并使用它生成嵌入？[y/N]”。未经明确同意，绝不下载（除非传入 --pretrained-release）。如果同时传入 --ckpt 和 --pretrained-release，则 中止 — 两者冲突。
   • **保存位置。**默认为 $KERMT_REPO/models/NV-KERMT-70M-v2/；若提供 --model-dir <dir>，则遵照指定位置。若已有完整模型包，则会复用。
   • 下载（前台运行；首次获取约 282 MB）：

     $KERMT_REPO/agent/scripts/kermt_container.sh run --model-dir <save-dir> -- \
         "python agent/scripts/fetch_released_model.py --out /model"
     

     解析 JSON；若 ok: false，则中止（显示 errors）。成功后设置 <user-ckpt> = <save-dir>/kermt_contrastive_v2.0.pt。

   验证已解析（或由用户提供的）检查点：

   $KERMT_REPO/agent/scripts/kermt_container.sh run --ckpt <user-ckpt> -- \
       "python agent/scripts/check_checkpoint.py --mode embed --ckpt /ckpt"
   

   解析 JSON。若 ok: false，则中止。验证器只会拒绝不含编码器的 检查点（罕见）。

4. 验证数据。

   $KERMT_REPO/agent/scripts/kermt_container.sh run --data <user-csv> -- \
       "python agent/scripts/check_data.py --mode embed --csv /data/<basename>"
   

5. 准备数据（仅清理 — 不执行特征步骤）。

   $KERMT_REPO/agent/scripts/kermt_container.sh run --data <user-csv> --run-dir $RUN_DIR -- \
       "python agent/scripts/prepare_data.py --mode embed \\
            --csv /data/<basename> --out /runs/data"
   

   输出位于 $RUN_DIR/data/prepare_data.json，其中只有一个 clean_csv 路径。task/extract_embeddings.py 会直接从 SMILES 在线生成特征。

6. 启动运行程序（阻塞式）。

   $KERMT_REPO/agent/scripts/kermt_container.sh run \\
       --ckpt <user-ckpt> --run-dir $RUN_DIR -- \\
       "python agent/scripts/run_extract_embeddings.py \\
            --ckpt /ckpt \\
            --prepare-manifest /runs/data/prepare_data.json \\
            --out /runs \\
            [--gpus 0 --batch-size N]"
   

7. 向用户报告。

   • 嵌入目录：$RUN_DIR/out/
     • atom_from_atom.npy、bond_from_atom.npy、 atom_from_bond.npy、bond_from_bond.npy（4 个标准读取结果； 每个形状为 (N_rows, hidden_size)）
     • metadata.pkl — 包含以下内容的字典 pickle：canonical_smiles （逐行经 RDKit 规范化的 SMILES）、valid（逐行布尔值：RDKit 是否成功 解析），以及其他运行元数据。
   • 清单：$RUN_DIR/run.json
   • 日志：$RUN_DIR/logs/embed.log

硬性规则

• 未经同意，绝不下载已发布的模型。 省略 --ckpt 时，只有在用户明确回复“是”或明确指定 --pretrained-release 标志后，才下载 nvidia/NV-KERMT-70M-v2。--ckpt 和 --pretrained-release 不能同时使用。
• 绝不修改用户的 ckpt。 运行器只通过 task/extract_embeddings.py 的 --checkpoint <path> 标志读取文件。
• 架构由 ckpt 决定。 该运行器不提供 --hidden-size 等标志；task/extract_embeddings.py 会从 ckpt 的 saved_args 中读取架构。

常见错误

• prepare_data manifest is missing required output 'clean_csv' → 运行 prepare 时使用了 --skip-clean，但未提供源 CSV。请重新运行 prepare，并移除该选项。
• --gpus '0,1' is single-GPU only → 请传入单个 ID。

可重放性

$(jq -r .cmd_replay $RUN_DIR/run.json)

如果 ok_to_replay: false（启动时 kermt 仓库工作树不干净），请通过 repo.commit 固定提交，并先运行 git checkout 切换到该提交。