kermt-embed

Extraire les embeddings par molécule depuis n'importe quel checkpoint KERMT porteur d'encoder. La skill est l'orchestrateur de workflow : valider le ckpt, valider le CSV, nettoyer les SMILES, lancer le runner en mode bloquant, retourner les fichiers .npy par readout.

Chemins de la skill et du runtime

Définissez SKILL_DIR au chemin absolu du répertoire de cette skill installée. Exportez KERMT_REPO comme le chemin absolu du checkout KERMT utilisé pour l'exécution du modèle. L'assistant conteneur bundlé monte ce checkout à /workspace et cette skill à /skill (lecture seule). Les commandes à l'intérieur du conteneur utilisent /skill/scripts/; les défauts sont bundlés dans config/. Voir Released models pour les exigences du bundle de checkpoint.

Téléchargements et sorties locales

La branche released-model optionnelle lit config/released_model.json pour le dépôt Hugging Face, la révision épinglée, et les noms de fichiers. Le script bundlé scripts/fetch_released_model.py télécharge le bundle de modèle en HTTPS dans le répertoire hôte que l'utilisateur sélectionne. Les modèles publics fonctionnent sans identifiants ; si HF_TOKEN est défini, l'assistant conteneur le transmet pour l'authentification Hugging Face. Les données préparées, logs, et résultats de workflow vont dans le répertoire de run choisi.

Exigences matérielles

GPUs : 1 (mono-GPU).
VRAM : ≥ 4 GB pour le batch_size 64 par défaut.
Disque : dépend de la taille de la sortie — environ quelques MB pour 1k molécules à hidden 800 par readout, soit ~10–20 MB pour 1k molécules sur les 4 readouts. Plus un petit canonical_smiles.npy + validity.npy par run.
Driver / CUDA : tout hôte supportant CUDA 12.6.

Entrées

Obligatoires :

--csv <chemin> — CSV de SMILES. La première colonne est smiles; les autres colonnes sont ignorées (aucune cible nécessaire).

Checkpoint (optionnel — par défaut le modèle released si omis) :

--ckpt <chemin> — tout checkpoint porteur d'encoder. Les checkpoints Grover_base, cmim, hybrid, et fine-tuned sont tous acceptés. Le validateur refuse uniquement les checkpoints sans encoder. S'il est omis, la skill propose de télécharger le modèle hybrid pretrained released nvidia/NV-KERMT-70M-v2 et d'embedding avec lui — voir « Resolve & validate the checkpoint » (étape 3 du workflow).
--pretrained-release — opt-in explicite pour utiliser le modèle released sans le prompt interactif (pour les runs non-interactifs / agent). Mutuellement exclusif avec --ckpt.
--model-dir <dir> — où sauvegarder le bundle téléchargé (défaut $KERMT_REPO/models/NV-KERMT-70M-v2/). Un bundle complet déjà présent là est réutilisé, non re-téléchargé.

Optionnels :

--batch-size N — surcharger le défaut configuré (64).
--gpus 0 — ID mono-GPU (défaut 0).
--from-prepare <dir> — sauter l'étape prepare et réutiliser un prepare_data.json existant dans <dir>.

Workflow

Soit $KERMT_REPO le chemin vers votre checkout du repo kermt.

Pré-vol : conteneur + sonde système.

"$SKILL_DIR/scripts/kermt_container.sh" check_system

Calculer le répertoire de run.

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

Résoudre et valider le checkpoint.

Résoudre — uniquement si --ckpt a été omis. Par défaut le modèle hybrid pretrained released nvidia/NV-KERMT-70M-v2 :
- Portail de consentement. Sauf si --pretrained-release a été passé, demander à l'utilisateur : « No checkpoint given — download the released model nvidia/NV-KERMT-70M-v2 (NVIDIA Open Model License, https://huggingface.co/nvidia/NV-KERMT-70M-v2) and embed with it? [y/N] ». Ne jamais télécharger sans un oui explicite (ou --pretrained-release). Si --ckpt et --pretrained-release sont donnés, abandonner — ils entrent en conflit.
- Lieu de sauvegarde. Par défaut $KERMT_REPO/models/NV-KERMT-70M-v2/; respecter --model-dir <dir> s'il est donné. Un bundle complet existant est réutilisé.
- Télécharger (foreground; ~282 MB au premier appel) :
```
"$SKILL_DIR/scripts/kermt_container.sh" run --model-dir <save-dir> -- \
    "python /skill/scripts/fetch_released_model.py --out /model"
```
  Parser le JSON; abandonner si ok: false (afficher errors). En cas de succès définir <user-ckpt> = <save-dir>/kermt_contrastive_v2.0.pt.
Valider le ckpt résolu (ou fourni par l'utilisateur) :
```
"$SKILL_DIR/scripts/kermt_container.sh" run --ckpt <user-ckpt> -- \
    "python /skill/scripts/check_checkpoint.py --mode embed --ckpt /ckpt"
```
Parser le JSON. Abandonner si ok: false. Le validateur refuse uniquement les checkpoints sans encoder (rares).

Valider les données.

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

Préparer les données (clean-only — pas d'étape features).
```
"$SKILL_DIR/scripts/kermt_container.sh" run --data <user-csv> --run-dir $RUN_DIR -- \
    "python /skill/scripts/prepare_data.py --mode embed \\
         --csv /data/<basename> --out /runs/data"
```
Les sorties atterrissent à $RUN_DIR/data/prepare_data.json avec un chemin clean_csv unique. task/extract_embeddings.py featurise depuis SMILES à la volée.

Lancer le runner (bloquant).

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

Rapporter à l'utilisateur.
- Répertoire des embeddings : $RUN_DIR/out/
  - atom_from_atom.npy, bond_from_atom.npy, atom_from_bond.npy, bond_from_bond.npy (les 4 readouts standards; chacun de forme (N_rows, hidden_size))
  - metadata.pkl — pickle d'un dict contenant canonical_smiles (SMILES canonicalisés RDKit par ligne), valid (booléen par ligne : RDKit l'a-t-il parsé), plus d'autres métadonnées de run.
- Manifest : $RUN_DIR/run.json
- Log : $RUN_DIR/logs/embed.log

Règles strictes

Ne jamais télécharger le modèle released sans consentement. Quand --ckpt est omis, télécharger nvidia/NV-KERMT-70M-v2 uniquement après un « oui » utilisateur explicite ou un flag --pretrained-release explicite. --ckpt et --pretrained-release sont mutuellement exclusifs.
Ne jamais modifier le ckpt de l'utilisateur. Le runner lit en lecture seule via le flag --checkpoint <chemin> de task/extract_embeddings.py.
L'architecture vient du ckpt. Aucun flag --hidden-size etc. sur ce runner; task/extract_embeddings.py lit l'architecture depuis les saved_args du ckpt.

Erreurs courantes

prepare_data manifest is missing required output 'clean_csv' → prepare a couru avec --skip-clean mais aucun CSV source donné. Relancer prepare sans lui.
--gpus '0,1' is single-GPU only → passer un ID unique.

Rejouabilité

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

Si ok_to_replay: false (dirty kermt repo worktree au moment du lancement), épingler le commit via repo.commit et faire git checkout d'abord.