kermt-embed

kermt-embed

エンコーダーを含むあらゆる KERMT チェックポイントから、分子ごとの埋め込みを抽出します。 このスキルはワークフローをオーケストレーションします。チェックポイントを検証し、CSV を検証し、 SMILES をクリーンアップして、ランナーをブロッキングで起動し、readout ごとの .npy ファイルを返します。

ハードウェア要件

• GPU：1 基（シングル GPU）。
• VRAM：デフォルトの batch_size 64 では 4 GB 以上。
• ディスク：出力サイズによって異なります。readout ごとに hidden 800 の場合、分子 1,000 件あたり数 MB 程度で、4 つの readout 全体では分子 1,000 件あたり約 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 — 対話形式のプロンプトを表示せずに、リリース済みモデルを使用することを明示的に指定します（非対話形式／エージェントでの実行向け）。--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]」明示的な yes（または --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 のパスが 1 つ含まれます。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 つの readout。 それぞれの形状は (N_rows, hidden_size) です）
     • metadata.pkl — canonical_smiles（行ごとに RDKit で正規化された SMILES）、valid（行ごとのブール値。RDKit による解析に成功したかどうか）、およびその他の実行メタデータを含む辞書の pickle です。
   • マニフェスト：$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' → --skip-clean を指定して prepare を実行したものの、ソース 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 してください。