SuperAPI AI 工具集

kermt-embed

通用与其它

查看项目

更新于 10-16 02:33

技能介绍

用途:从任意含编码器的 KERMT 检查点提取逐分子嵌入。该技能本身是工作流编排器,负责校验检查点、校验 CSV、清洗 SMILES、阻塞式启动 runner,并返回各读出的 .npy 文件。

主要能力:支持 grover_base、cmim、hybrid、finetuned 等含编码器的检查点;省略 --ckpt 时可使用已发布的预训练 hybrid 模型 nvidia/NV-KERMT-70M-v2;输出四种标准读出(atom_from_atom.npy、bond_from_atom.npy、atom_from_bond.npy、bond_from_bond.npy),每行形状为 (N_rows, hidden_size),并生成 metadata.pkl(含规范化 SMILES 与逐行有效标记)。

硬件要求:单 GPU;默认 batch_size 64 需 ≥4 GB 显存;磁盘占用约为每千分子每读出几 MB,四种读出合计约 10–20 MB/千分子,另加少量 canonical_smiles.npy 与 validity.npy;需支持 CUDA 12.6。

主要输入:--csv <path>(首列为 SMILES,其余列忽略);--ckpt <path>(可选,默认使用已发布模型);--pretrained-release(非交互式使用已发布模型,与 --ckpt 互斥);--model-dir <dir>(默认 $KERMT_REPO/models/NV-KERMT-70M-v2/);--batch-size N、--gpus 0、--from-prepare <dir>。

使用方式:依次执行容器检查、计算 RUN_DIR、解析并校验检查点、校验数据、prepare_data(仅清洗)、阻塞式启动 run_extract_embeddings,最后报告 $RUN_DIR/out/ 下的嵌入文件与 $RUN_DIR/run.json。

硬性规则:未获用户明确同意或未传 --pretrained-release 时不得下载已发布模型;--ckpt 与 --pretrained-release 互斥;不修改用户检查点(只读);架构来自检查点,runner 无 --hidden-size 等参数。

限制与错误:prepare 清单缺少 clean_csv、--gpus '0,1' 非单 GPU 等会报错;若 ok_to_replay 为 false(启动时 kermt 仓库工作区不干净),需先按 repo.commit 固定提交再 checkout。官方文档未提供的收费、价格等信息暂未提供。