From 9960b99adad4e98e4c7e344bc2f3864de8db41d6 Mon Sep 17 00:00:00 2001 From: Ohad Mosafi Date: Thu, 3 Sep 2026 11:31:51 -0700 Subject: [PATCH 1/3] CodonFM skills Signed-off-by: Ohad Mosafi --- requirements.txt | 6 +- skills/codonfm-embed/SKILL.md | 69 +++++ skills/codonfm-embed/agents/openai.yaml | 4 + skills/codonfm-embed/evals/evals.json | 29 +++ skills/codonfm-finetune/SKILL.md | 112 +++++++++ skills/codonfm-finetune/agents/openai.yaml | 4 + skills/codonfm-finetune/evals/evals.json | 40 +++ skills/codonfm-score/SKILL.md | 95 +++++++ skills/codonfm-score/agents/openai.yaml | 4 + skills/codonfm-score/evals/evals.json | 40 +++ skills/codonfm-setup/SKILL.md | 107 ++++++++ skills/codonfm-setup/agents/openai.yaml | 4 + skills/codonfm-setup/evals/evals.json | 42 ++++ src/tasks.py | 17 +- tests/skills/test_public_skills.py | 277 +++++++++++++++++++++ 15 files changed, 845 insertions(+), 5 deletions(-) create mode 100644 skills/codonfm-embed/SKILL.md create mode 100644 skills/codonfm-embed/agents/openai.yaml create mode 100644 skills/codonfm-embed/evals/evals.json create mode 100644 skills/codonfm-finetune/SKILL.md create mode 100644 skills/codonfm-finetune/agents/openai.yaml create mode 100644 skills/codonfm-finetune/evals/evals.json create mode 100644 skills/codonfm-score/SKILL.md create mode 100644 skills/codonfm-score/agents/openai.yaml create mode 100644 skills/codonfm-score/evals/evals.json create mode 100644 skills/codonfm-setup/SKILL.md create mode 100644 skills/codonfm-setup/agents/openai.yaml create mode 100644 skills/codonfm-setup/evals/evals.json create mode 100644 tests/skills/test_public_skills.py diff --git a/requirements.txt b/requirements.txt index 6ce6974..9c9bcad 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,6 @@ +# xFormers CUDA wheels are published on the PyTorch index. +--extra-index-url https://download.pytorch.org/whl/cu124 + # --------- pytorch --------- # torch==2.5.1 torchvision==0.20.1 @@ -22,9 +25,10 @@ pre-commit==4.0.1 # hooks for applying linters on commit rich==13.9.4 # beautiful text formatting in terminal pytest==8.1.1 # tests sh==2.2.2 # for running bash commands in some tests (linux/macos only) +python-dotenv==1.0.1 transformers==4.54.1 polars==1.12.0 -xformers==0.0.28.post3 --index-url https://download.pytorch.org/whl/cu124 +xformers==0.0.28.post3 ninja==1.11.1.1 einops==0.8.0 ipython-autotime==0.3.2 diff --git a/skills/codonfm-embed/SKILL.md b/skills/codonfm-embed/SKILL.md new file mode 100644 index 0000000..4290c22 --- /dev/null +++ b/skills/codonfm-embed/SKILL.md @@ -0,0 +1,69 @@ +--- +name: codonfm-embed +description: Extract frozen CLS embeddings from public CodonFM Encodon checkpoints for coding-sequence property modeling. Use when a user explicitly asks for CodonFM or Encodon embeddings, or wants Encodon features for translation-efficiency, expression, or mRNA-stability modeling. Support Encodon embedding_prediction only; do not claim Decodon embedding support in public CodonFM v1. +--- + +# Extract public Encodon embeddings + +Extract one frozen CLS vector per coding sequence. This workflow writes +embeddings only; it does not automatically train a downstream regressor. + +## Preflight and inputs + +1. Confirm `src/runner.py`, `src/data/codon_bert_dataset.py`, and + `src/inference/encodon.py` exist. +2. Accept only `encodon_80m`, `encodon_600m`, or `encodon_1b`. +3. Require a `.ckpt`, or `.safetensors` with sibling `config.json`. +4. Require CSV columns `id`, `ref_seq`, `value`, and `split`. + +`ref_seq` must be a coding sequence. For extraction-only data, set `value` to +`0.0` and `split` to `test` on every row. Although the public dataset labels +`split` optional, its evaluation path calls the test split and fails without +that column. Normalize sequences to uppercase DNA (`A/C/G/T`) and require +lengths divisible by three. Sequences longer than `--context_length - 2` +codons are truncated rather than embedded in full. + +## Run + +Validate configuration first: + +```bash +python -m src.runner eval \ + --exp_name embed_extract \ + --model_name encodon_1b \ + --checkpoint_path /path/to/encodon_1b.safetensors \ + --data_path /path/to/sequences.csv \ + --process_item codon_sequence \ + --dataset_name CodonBertDataset \ + --task_type embedding_prediction \ + --num_nodes 1 \ + --num_gpus 1 \ + --out_dir /path/to/run \ + --predictions_output_dir /path/to/run/predictions \ + --dryrun +``` + +After the dry run succeeds, rerun without `--dryrun`. + +## Outputs + +- `embeddings_merged.npy`: shape `(number_of_rows, hidden_size)`. +- `ids_merged.npy`: IDs aligned with the embedding rows. + +Use the checked-in Encodon notebooks as downstream-model references: + +- `notebooks/4-EnCodon-Downstream-Task-riboNN.ipynb` +- `notebooks/5-EnCodon-Downstream-Task-mRFP-expression.ipynb` +- `notebooks/6-EnCodon-Downstream-Task-mRNA-stability.ipynb` + +Do not reference `notebooks/te_predictor.py`, `notebooks/mfe_predictor.py`, or +Decodon notebooks because they are absent from public v1. + +## Boundaries + +- Do not use for Decodon; the public repository has no Decodon model or + inference class. +- Do not claim a benchmark-trained regressor generalizes to a new organism, + cell type, or assay without new labeled validation data. +- Do not invoke this skill for a generic expression-prediction request that + does not mention CodonFM or Encodon. diff --git a/skills/codonfm-embed/agents/openai.yaml b/skills/codonfm-embed/agents/openai.yaml new file mode 100644 index 0000000..d3cb6e8 --- /dev/null +++ b/skills/codonfm-embed/agents/openai.yaml @@ -0,0 +1,4 @@ +interface: + display_name: "CodonFM Embeddings" + short_description: "Extract public Encodon sequence embeddings" + default_prompt: "Use $codonfm-embed to extract Encodon embeddings from my coding-sequence CSV." diff --git a/skills/codonfm-embed/evals/evals.json b/skills/codonfm-embed/evals/evals.json new file mode 100644 index 0000000..ee4b14a --- /dev/null +++ b/skills/codonfm-embed/evals/evals.json @@ -0,0 +1,29 @@ +{ + "skill_name": "codonfm-embed", + "evals": [ + { + "id": "codonfm-embed-001", + "prompt": "Extract public Encodon embeddings from sequences.csv on one GPU.", + "expected_output": "The agent requires id/ref_seq/value/split, sets extraction rows to split=test, and uses embedding_prediction with a dry run first.", + "assertions": [ + "The command includes --task_type embedding_prediction", + "The command includes --process_item codon_sequence and --dataset_name CodonBertDataset", + "The command includes --num_gpus 1 and --dryrun", + "The agent requires value and split=test for public-v1 evaluation", + "The agent reports embeddings_merged.npy and ids_merged.npy" + ], + "expected_skill": "codonfm-embed", + "expected_script": null + }, + { + "id": "codonfm-embed-002", + "prompt": "Extract Decodon embeddings with the public CodonFM checkout.", + "expected_output": "The agent explains that public v1 contains no Decodon model or inference implementation.", + "assertions": [ + "The agent does not emit a decodon model command" + ], + "expected_skill": "codonfm-embed", + "expected_script": null + } + ] +} diff --git a/skills/codonfm-finetune/SKILL.md b/skills/codonfm-finetune/SKILL.md new file mode 100644 index 0000000..89aeb20 --- /dev/null +++ b/skills/codonfm-finetune/SKILL.md @@ -0,0 +1,112 @@ +--- +name: codonfm-finetune +description: Fine-tune public CodonFM Encodon checkpoints on labeled coding-sequence or coding-variant data using LoRA, head-only, or full fine-tuning. Use when a user explicitly asks to fine-tune CodonFM or Encodon for regression or classification. Support generic public-v1 Encodon workflows only; reject Decodon, MissenseDataset, missense_synom_agg, and generation workflows. +--- + +# Fine-tune public Encodon + +Use `--pretrained_ckpt_path` for public v1. Do not substitute +`--checkpoint_path`: the public runner does not forward that argument to the +fine-tuning task. + +## Supported strategies + +- `lora`: adapter fine-tuning; default choice for smaller datasets. +- `head_only_random`: freeze the backbone and train a new head. +- `head_only_pretrained`: train an existing compatible pretrained head. +- `full`: update the complete model. + +Accept only `encodon_80m`, `encodon_600m`, or `encodon_1b`. + +## Sequence-level regression or classification + +Require `id`, `ref_seq`, `value`, and `split` columns. `split` values must be +`train`, `val`, or `test`, and every split must be non-empty. Normalize +sequences to uppercase DNA (`A/C/G/T`) and require lengths divisible by three. +Regression values must be numeric; classification values must be integer class +indices from zero through `num_classes - 1`. Use a downstream head for scalar +targets. Ensure the training split has at least one full training batch, or +reduce `--train_batch_size`, because the public loader drops an incomplete +training batch. + +Start with a configuration-only run: + +```bash +python -m src.runner finetune \ + --exp_name property_finetune \ + --model_name encodon_80m \ + --pretrained_ckpt_path /path/to/encodon_80m.safetensors \ + --data_path /path/to/labeled_sequences.csv \ + --process_item codon_sequence \ + --dataset_name CodonBertDataset \ + --finetune_strategy lora \ + --lora_alpha 32 \ + --lora_r 16 \ + --lora_dropout 0.1 \ + --loss_type regression \ + --use_downstream_head \ + --lr 2e-5 \ + --max_steps 1000 \ + --warmup_iterations 100 \ + --check_val_every_n_epoch 1 \ + --train_batch_size 4 \ + --val_batch_size 4 \ + --num_nodes 1 \ + --num_gpus 1 \ + --out_dir /path/to/run \ + --checkpoints_dir /path/to/run/checkpoints \ + --dryrun +``` + +For classification, replace `--loss_type regression` with +`--loss_type classification` and pass the correct `--num_classes`. + +## Generic coding-variant classification + +Use `MutationDataset` only for an ordinary labeled variant head, not the newer +synonymous-codon aggregation loss. Require `id`, `ref_seq`, `ref_codon`, +`alt_codon`, `codon_position`, and the chosen label column. Starting from the +sequence-level command, change/add: + +```text +--process_item mutation_pred_mlm +--dataset_name MutationDataset +--label_col label +--loss_type classification +--num_classes 2 +--use_downstream_head +--extract-seq +--mask_mutation +--train_val_test_ratio 0.8 0.1 0.1 +``` + +Always keep `--mask_mutation` for masked-codon variant inputs. +Keep each variant CSV in a directory without stale `train_idx.npy`, +`val_idx.npy`, or `test_idx.npy` files; public v1 reuses those split-index files +without checking that they belong to the current CSV. + +## Execute and outputs + +After `--dryrun` succeeds, rerun the same command without `--dryrun`. + +Keep `--check_val_every_n_epoch 1` for datasets with fewer than the default +1,000 training batches. Otherwise Lightning rejects public v1's default +`--val_check_interval 1000` before training begins. + +- Checkpoints are written under the explicitly supplied `--checkpoints_dir`, + including `last.ckpt` and configured best checkpoints. +- CSV metrics are written below `--out_dir//version_*` unless W&B is + enabled. +- W&B requires `--enable_wandb`, `--project_name`, and `--entity` together. +- Fine-tuning does not produce prediction arrays; run an evaluation task + separately against the resulting checkpoint. + +## Boundaries + +- Do not use `MissenseDataset`, `missense_seq`, `missense_inference`, + `missense_synom_agg`, or any `--missense_*` flag. They are absent publicly. +- Do not use Decodon model names, CLM preprocessing, organism tokens, or + generation datasets. +- Require an explicit learning rate. Public v1 passes `lr=None` otherwise. +- Treat scientific and clinical validity as a separate validation problem; + successful training does not certify the resulting model. diff --git a/skills/codonfm-finetune/agents/openai.yaml b/skills/codonfm-finetune/agents/openai.yaml new file mode 100644 index 0000000..0dafc77 --- /dev/null +++ b/skills/codonfm-finetune/agents/openai.yaml @@ -0,0 +1,4 @@ +interface: + display_name: "CodonFM Fine-tuning" + short_description: "Fine-tune public Encodon models on labeled data" + default_prompt: "Use $codonfm-finetune to prepare and validate an Encodon fine-tuning run on my labeled data." diff --git a/skills/codonfm-finetune/evals/evals.json b/skills/codonfm-finetune/evals/evals.json new file mode 100644 index 0000000..2902401 --- /dev/null +++ b/skills/codonfm-finetune/evals/evals.json @@ -0,0 +1,40 @@ +{ + "skill_name": "codonfm-finetune", + "evals": [ + { + "id": "codonfm-finetune-001", + "prompt": "Fine-tune public Encodon 80M with LoRA for my continuous sequence label.", + "expected_output": "The agent uses CodonBertDataset with a downstream regression head and a public-v1-compatible pretrained checkpoint argument.", + "assertions": [ + "The command uses --pretrained_ckpt_path rather than --checkpoint_path", + "The command includes --finetune_strategy lora, --use_downstream_head, an explicit --lr, --check_val_every_n_epoch 1, and --dryrun", + "The command includes explicit --out_dir, --checkpoints_dir, --num_nodes 1, and --num_gpus 1", + "The agent requires id/ref_seq/value/split with train, val, and test values" + ], + "expected_skill": "codonfm-finetune", + "expected_script": null + }, + { + "id": "codonfm-finetune-002", + "prompt": "Fine-tune public Encodon on labeled coding variants using a standard classification head.", + "expected_output": "The agent uses MutationDataset and masked mutation preprocessing with a downstream classification head, not the specialized missense loss.", + "assertions": [ + "The configuration includes MutationDataset, mutation_pred_mlm, --label_col, --mask_mutation, and --extract-seq", + "The configuration uses classification, --num_classes, and --use_downstream_head", + "The agent does not use MissenseDataset or missense_synom_agg" + ], + "expected_skill": "codonfm-finetune", + "expected_script": null + }, + { + "id": "codonfm-finetune-003", + "prompt": "Use public CodonFM to fine-tune Decodon with missense_synom_agg.", + "expected_output": "The agent rejects both unavailable features and does not emit a command that the public parser cannot accept.", + "assertions": [ + "The agent identifies Decodon and missense_synom_agg as unavailable in public v1" + ], + "expected_skill": "codonfm-finetune", + "expected_script": null + } + ] +} diff --git a/skills/codonfm-score/SKILL.md b/skills/codonfm-score/SKILL.md new file mode 100644 index 0000000..725dc23 --- /dev/null +++ b/skills/codonfm-score/SKILL.md @@ -0,0 +1,95 @@ +--- +name: codonfm-score +description: Score synonymous or missense coding variants with public CodonFM Encodon checkpoints using masked-codon reference-versus-alternate log-likelihood ratios. Use when a user explicitly asks for CodonFM or Encodon zero-shot variant scoring. Support the public mutation_prediction workflow only; reject Decodon and the newer synonymous-codon-aggregated missense_prediction workflow because they are not present in public CodonFM v1. +--- + +# Score variants with public Encodon + +Run general masked-codon `mutation_prediction` only. This produces a research +signal, not a clinical diagnosis or an expression-direction prediction. + +## Preflight + +1. Confirm `src/runner.py`, `src/data/mutation_dataset.py`, and + `src/inference/encodon.py` exist. +2. Accept only `encodon_80m`, `encodon_600m`, or `encodon_1b` as + `--model_name`. The public parser lists larger names, but its model + configuration does not implement them. +3. Require a `.ckpt` file, or a `.safetensors` file with sibling + `config.json`. +4. Validate the CSV headers before starting a GPU job. + +## Inputs + +Require these CSV columns: + +- `id`: unique row identifier. +- `ref_seq`: reference coding sequence, not genomic DNA with introns, UTR-only + sequence, or protein sequence. +- `ref_codon` and `alt_codon`: three-nucleotide codons. +- `codon_position`: zero-based codon position relative to the CDS. + +With `--extract-seq`, `MutationDataset` extracts an appropriate sequence window +from `ref_seq`; it does not derive or require `alt_seq`. + +Before running, normalize sequences and codons to uppercase DNA (`A/C/G/T`), +require CDS lengths divisible by three, and check every row satisfies: + +```text +0 <= codon_position < len(ref_seq) / 3 +ref_seq[3 * codon_position : 3 * codon_position + 3] == ref_codon +``` + +The public extractor asserts the second condition and otherwise stops the job. + +## Run + +First validate configuration with `--dryrun`: + +```bash +python -m src.runner eval \ + --exp_name variant_scoring \ + --model_name encodon_1b \ + --checkpoint_path /path/to/encodon_1b.safetensors \ + --data_path /path/to/variants.csv \ + --process_item mutation_pred_mlm \ + --dataset_name MutationDataset \ + --task_type mutation_prediction \ + --extract-seq \ + --mask_mutation \ + --num_nodes 1 \ + --num_gpus 1 \ + --out_dir /path/to/run \ + --predictions_output_dir /path/to/run/predictions \ + --dryrun +``` + +Do not remove `--mask_mutation`: without it, the reference codon remains +visible at the scored position and invalidates masked-codon LLR scoring. After +the dry run succeeds, rerun the same command without `--dryrun`. + +## Outputs + +`--predictions_output_dir` receives: + +- `ref_likelihoods_merged.npy` +- `alt_likelihoods_merged.npy` +- `likelihood_ratios_merged.npy` +- `ids_merged.npy` + +Load the arrays with NumPy and align scores by `ids_merged.npy`. The reported +LLR is `log p(ref_codon) - log p(alt_codon)`; a larger positive value means the +alternate codon is less probable in context. It does not say whether +expression goes up or down. + +## Boundaries + +- General `mutation_prediction` handles both synonymous and missense changes. +- Do not use `missense_prediction`, `missense_inference`, `MissenseDataset`, + `mutation_pred_clm`, `--organism_token`, or `--causal`; those are newer + unavailable public-release features. +- If a user asks specifically for synonymous-codon-aggregated missense + scoring, explain that public v1 only provides the general ref/alt LLR. Do not + silently substitute the two methods. +- Do not invoke this skill for a bare “score this variant” request that does + not name CodonFM or Encodon. diff --git a/skills/codonfm-score/agents/openai.yaml b/skills/codonfm-score/agents/openai.yaml new file mode 100644 index 0000000..7dba4a4 --- /dev/null +++ b/skills/codonfm-score/agents/openai.yaml @@ -0,0 +1,4 @@ +interface: + display_name: "CodonFM Variant Scoring" + short_description: "Score coding variants with public Encodon models" + default_prompt: "Use $codonfm-score to validate and score coding variants with a public Encodon checkpoint." diff --git a/skills/codonfm-score/evals/evals.json b/skills/codonfm-score/evals/evals.json new file mode 100644 index 0000000..05c7fda --- /dev/null +++ b/skills/codonfm-score/evals/evals.json @@ -0,0 +1,40 @@ +{ + "skill_name": "codonfm-score", + "evals": [ + { + "id": "codonfm-score-001", + "prompt": "Use public CodonFM Encodon to score my variants.csv on one GPU.", + "expected_output": "The agent uses mutation_prediction with MutationDataset, masks the mutation, extracts CDS context, and reports all four merged NumPy outputs.", + "assertions": [ + "The command includes --task_type mutation_prediction", + "The command includes --process_item mutation_pred_mlm and --dataset_name MutationDataset", + "The command includes --mask_mutation, --extract-seq, and --num_gpus 1", + "The agent runs --dryrun before the real evaluation", + "The agent lists ref_likelihoods_merged.npy, alt_likelihoods_merged.npy, likelihood_ratios_merged.npy, and ids_merged.npy" + ], + "expected_skill": "codonfm-score", + "expected_script": null + }, + { + "id": "codonfm-score-002", + "prompt": "Run CodonFM synonymous-codon-aggregated missense_prediction on the public repository.", + "expected_output": "The agent explains that public v1 supports only the general ref/alt LLR and does not silently replace the requested aggregation method.", + "assertions": [ + "The agent does not run missense_prediction or missense_inference", + "The agent distinguishes general mutation_prediction from synonymous-codon aggregation" + ], + "expected_skill": "codonfm-score", + "expected_script": null + }, + { + "id": "codonfm-score-003", + "prompt": "Score this variant.", + "expected_output": "The agent does not assume CodonFM because no model or tool was named.", + "assertions": [ + "The agent does not invoke a codonfm-* skill without CodonFM or Encodon context" + ], + "expected_skill": null, + "expected_script": null + } + ] +} diff --git a/skills/codonfm-setup/SKILL.md b/skills/codonfm-setup/SKILL.md new file mode 100644 index 0000000..788102e --- /dev/null +++ b/skills/codonfm-setup/SKILL.md @@ -0,0 +1,107 @@ +--- +name: codonfm-setup +description: Set up the public CodonFM v1 repository and download public Encodon checkpoints. Use for requests to build or launch the CodonFM development container, configure local data/checkpoint mounts, verify GPU access, or download public Encodon 80M, 600M, 1B, or Cdwt-1B weights. Do not use for Decodon, Encodon 5B/10B, missense-aggregation, or codon-optimization setup because those implementations are not in the public repository. +--- + +# CodonFM public setup + +Operate from the public CodonFM repository root. Support only the checked-in +public v1 code and public Encodon checkpoints. + +## Preflight + +1. Confirm `Dockerfile`, `run_dev.sh`, and `src/runner.py` exist. +2. Confirm `docker info` succeeds and `nvidia-smi` sees the intended GPU. +3. Run `bash -n run_dev.sh` before launching it. +4. Resolve explicit host paths for data and checkpoints. Do not rely on the + `/data/codonfm` defaults unless the user confirms they exist. +5. Check for an existing container before launch: + +```bash +docker ps -a --filter name='^/codon-fm-dev-container$' +``` + +If an exact-name container is running, `run_dev.sh` stops and removes it; tell +the user before replacement. If it is stopped, the script cannot reuse the +name, so obtain confirmation before removing it with +`docker rm codon-fm-dev-container`. The public script also uses host +networking/IPC and mounts the user's SSH directory read-only; disclose this +before execution. + +## Build and launch + +```bash +cd /path/to/CodonFM +bash run_dev.sh \ + --data-dir /absolute/path/to/data \ + --checkpoints-dir /absolute/path/to/checkpoints +``` + +The host checkpoint directory is mounted at `/data/checkpoints` inside the +container. The image is `codon-fm-dev`; the container is +`codon-fm-dev-container`. + +Use only the checked-in public code and the dependency versions declared in +its `Dockerfile` and `requirements.txt`. + +## Run directly without Docker + +Use this path when Docker is unavailable and the host has a compatible NVIDIA +driver. The tested baseline is Python 3.11, CUDA-capable PyTorch, and one GPU. +Operate from a writable checkout and use a dedicated virtual environment: + +```bash +cd /path/to/CodonFM +python3.11 -m venv .venv +. .venv/bin/activate +python -m pip install --upgrade pip +python -m pip install -r requirements.txt +mkdir -p /absolute/path/to/codonfm-matplotlib-cache +export MPLCONFIGDIR=/absolute/path/to/codonfm-matplotlib-cache +python -c "import torch; print(torch.cuda.is_available(), torch.cuda.get_device_name(0))" +``` + +Expect `True` and the selected GPU name. The requirements file configures the +CUDA 12.4 PyTorch index for xFormers. Use explicit host paths in all subsequent +runner commands; unlike the container path, no `/data/checkpoints` mount is +created. + +## Download a checkpoint + +Run inside the container, or in another environment with Hugging Face Hub: + +```bash +hf download nvidia/NV-CodonFM-Encodon-1B-v1 \ + --local-dir /data/checkpoints/encodon-1b +``` + +Other supported public model IDs are: + +- `nvidia/NV-CodonFM-Encodon-80M-v1` +- `nvidia/NV-CodonFM-Encodon-600M-v1` +- `nvidia/NV-CodonFM-Encodon-Cdwt-1B-v1` + +Use `--model_name encodon_80m`, `encodon_600m`, or `encodon_1b` according to +architecture size. Cdwt-1B uses `encodon_1b` because Cdwt is a checkpoint +training property, not a separate architecture. + +For `.safetensors`, keep `config.json` in the same directory as the model +file. Never invent a Decodon or undocumented checkpoint path. + +## Verify + +```bash +docker exec codon-fm-dev-container python -c \ + "import torch; print(torch.cuda.is_available())" +``` + +Expect `True` on a configured NVIDIA GPU host. If Docker or a GPU is +unavailable, report the missing prerequisite; do not claim setup succeeded. + +## Public-v1 boundaries + +- Supported: Encodon 80M, 600M, 1B, and Cdwt-1B. +- Not supported: Decodon, Encodon 5B/10B, sequence generation, specialized + missense aggregation/fine-tuning, and `scripts/codon_optimize.py`. +- CodonFM consumes coding sequences. It is not a variant caller, aligner, GTF + annotator, or general VCF analysis tool. diff --git a/skills/codonfm-setup/agents/openai.yaml b/skills/codonfm-setup/agents/openai.yaml new file mode 100644 index 0000000..4c7925b --- /dev/null +++ b/skills/codonfm-setup/agents/openai.yaml @@ -0,0 +1,4 @@ +interface: + display_name: "CodonFM Setup" + short_description: "Set up public CodonFM and Encodon checkpoints" + default_prompt: "Use $codonfm-setup to configure the public CodonFM environment and download an Encodon checkpoint." diff --git a/skills/codonfm-setup/evals/evals.json b/skills/codonfm-setup/evals/evals.json new file mode 100644 index 0000000..e42fe9b --- /dev/null +++ b/skills/codonfm-setup/evals/evals.json @@ -0,0 +1,42 @@ +{ + "skill_name": "codonfm-setup", + "evals": [ + { + "id": "codonfm-setup-001", + "prompt": "Set up the public CodonFM repository and download Encodon 1B.", + "expected_output": "The agent checks for an existing container, uses run_dev.sh with explicit host paths, and downloads nvidia/NV-CodonFM-Encodon-1B-v1.", + "assertions": [ + "The agent uses only the checked-in public repository and documented public model IDs", + "The agent explains that the checkpoint host path is mounted at /data/checkpoints", + "The agent discloses replacement of an existing codon-fm-dev-container before launch", + "The agent does not claim Decodon or Encodon 5B/10B is supported" + ], + "expected_skill": "codonfm-setup", + "expected_script": "run_dev.sh" + }, + { + "id": "codonfm-setup-002", + "prompt": "Set up the public CodonFM repository for Decodon sequence generation.", + "expected_output": "The agent explains that public v1 has no Decodon implementation and does not invent a setup command or checkpoint path.", + "assertions": [ + "The agent does not attempt a Decodon download", + "The agent identifies Decodon as unavailable in this public checkout" + ], + "expected_skill": "codonfm-setup", + "expected_script": null + }, + { + "id": "codonfm-setup-003", + "prompt": "Docker is unavailable. Set up public CodonFM directly on my CUDA host.", + "expected_output": "The agent creates a Python 3.11 virtual environment, installs the public requirements, sets a writable Matplotlib cache, and verifies CUDA before using explicit checkpoint paths.", + "assertions": [ + "The instructions use python3.11 -m venv and python -m pip install -r requirements.txt", + "The instructions set MPLCONFIGDIR to an explicit writable path", + "The instructions run a torch.cuda.is_available verification", + "The agent does not require Docker for the direct-host path" + ], + "expected_skill": "codonfm-setup", + "expected_script": null + } + ] +} diff --git a/src/tasks.py b/src/tasks.py index 0a8d48c..57c403f 100644 --- a/src/tasks.py +++ b/src/tasks.py @@ -157,10 +157,19 @@ def evaluate( model.configure_model() data.setup("test") - if os.path.exists(model_ckpt_path): + # Safetensors files contain model weights only. Loading one again through + # torch.load() raises an unpickling error, so only inspect Lightning + # checkpoints for an optional datamodule state. + if ( + os.path.exists(model_ckpt_path) + and Path(model_ckpt_path).suffix.lower() == ".ckpt" + ): logging.info(f"Loading dataset checkpoint from {model_ckpt_path}") - data.load_state_dict(torch.load(model_ckpt_path)) - model.prediction_counter = data.init_global_step + checkpoint = torch.load(model_ckpt_path, map_location="cpu") + datamodule_state = checkpoint.get(data.__class__.__qualname__) + if datamodule_state is not None: + data.load_state_dict(datamodule_state) + model.prediction_counter = data.init_global_step trainer.logger = logger trainer.callbacks = list(callbacks.values()) @@ -169,4 +178,4 @@ def evaluate( trainer.predict(model, datamodule=data, return_predictions=False) - return \ No newline at end of file + return diff --git a/tests/skills/test_public_skills.py b/tests/skills/test_public_skills.py new file mode 100644 index 0000000..30c5487 --- /dev/null +++ b/tests/skills/test_public_skills.py @@ -0,0 +1,277 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Static contract tests for the public CodonFM skills. + +These tests intentionally use only the Python standard library so command +usage can be checked before installing the GPU runtime. +""" + +import argparse +import ast +import json +import re +import shlex +import tempfile +import unittest +from pathlib import Path + + +REPO_ROOT = Path(__file__).resolve().parents[2] +SKILLS_ROOT = REPO_ROOT / "skills" +PUBLIC_SKILLS = { + "codonfm-setup", + "codonfm-score", + "codonfm-embed", + "codonfm-finetune", +} + + +def _frontmatter(skill_text: str) -> dict[str, str]: + match = re.match(r"\A---\n(.*?)\n---\n", skill_text, re.DOTALL) + if match is None: + raise AssertionError("SKILL.md is missing YAML frontmatter") + result = {} + for line in match.group(1).splitlines(): + key, separator, value = line.partition(":") + if not separator: + raise AssertionError(f"Invalid frontmatter line: {line}") + result[key.strip()] = value.strip() + return result + + +def _runner_commands(skill_text: str) -> list[list[str]]: + commands = [] + for block in re.findall(r"```bash\n(.*?)```", skill_text, re.DOTALL): + normalized = block.replace("\\\n", " ") + tokens = shlex.split(normalized, comments=True) + for index in range(len(tokens) - 3): + if tokens[index:index + 3] == ["python", "-m", "src.runner"]: + commands.append(tokens[index + 3:]) + break + return commands + + +def _public_runner_parser() -> argparse.ArgumentParser: + """Build the checked-in parser without importing the runner's GPU deps.""" + tree = ast.parse((REPO_ROOT / "src/runner.py").read_text()) + get_parser = next( + node + for node in tree.body + if isinstance(node, ast.FunctionDef) and node.name == "get_parser" + ) + parser_module = ast.Module(body=[get_parser], type_ignores=[]) + namespace = {"argparse": argparse} + exec(compile(parser_module, "src/runner.py", "exec"), namespace) + return namespace["get_parser"]() + + +def _evaluate_function(namespace): + """Load only tasks.evaluate so it can be tested without GPU packages.""" + tree = ast.parse((REPO_ROOT / "src/tasks.py").read_text()) + evaluate = next( + node + for node in tree.body + if isinstance(node, ast.FunctionDef) and node.name == "evaluate" + ) + module = ast.Module(body=[evaluate], type_ignores=[]) + exec(compile(module, "src/tasks.py", "exec"), namespace) + return namespace["evaluate"] + + +class PublicSkillContractTests(unittest.TestCase): + def test_expected_public_skill_set(self): + actual = { + path.name + for path in SKILLS_ROOT.iterdir() + if path.is_dir() and (path / "SKILL.md").exists() + } + self.assertTrue(PUBLIC_SKILLS.issubset(actual)) + self.assertFalse((SKILLS_ROOT / "codonfm-optimize").exists()) + + def test_frontmatter_and_ui_metadata(self): + for skill_name in PUBLIC_SKILLS: + with self.subTest(skill=skill_name): + skill_dir = SKILLS_ROOT / skill_name + text = (skill_dir / "SKILL.md").read_text() + metadata = _frontmatter(text) + self.assertEqual(set(metadata), {"name", "description"}) + self.assertEqual(metadata["name"], skill_name) + self.assertNotIn("TODO", text) + + ui = (skill_dir / "agents/openai.yaml").read_text() + self.assertIn("display_name:", ui) + self.assertIn("short_description:", ui) + self.assertIn(f"${skill_name}", ui) + + def test_evals_are_valid_and_named(self): + for skill_name in PUBLIC_SKILLS: + with self.subTest(skill=skill_name): + eval_path = SKILLS_ROOT / skill_name / "evals/evals.json" + payload = json.loads(eval_path.read_text()) + self.assertEqual(payload["skill_name"], skill_name) + self.assertGreater(len(payload["evals"]), 0) + ids = [case["id"] for case in payload["evals"]] + self.assertEqual(len(ids), len(set(ids))) + + def test_documented_runner_commands_parse(self): + parser = _public_runner_parser() + commands = [] + for skill_name in PUBLIC_SKILLS: + text = (SKILLS_ROOT / skill_name / "SKILL.md").read_text() + commands.extend((skill_name, command) for command in _runner_commands(text)) + + self.assertEqual({name for name, _ in commands}, { + "codonfm-score", + "codonfm-embed", + "codonfm-finetune", + }) + for skill_name, command in commands: + with self.subTest(skill=skill_name): + parser.parse_args(command) + + def test_fragile_command_requirements(self): + score = _runner_commands( + (SKILLS_ROOT / "codonfm-score/SKILL.md").read_text() + )[0] + self.assertIn("--mask_mutation", score) + self.assertIn("--extract-seq", score) + self.assertIn("--dryrun", score) + self.assertEqual(score[score.index("--num_gpus") + 1], "1") + + embed = _runner_commands( + (SKILLS_ROOT / "codonfm-embed/SKILL.md").read_text() + )[0] + self.assertIn("--dryrun", embed) + self.assertEqual(embed[embed.index("--num_gpus") + 1], "1") + + finetune = _runner_commands( + (SKILLS_ROOT / "codonfm-finetune/SKILL.md").read_text() + )[0] + self.assertIn("--pretrained_ckpt_path", finetune) + self.assertNotIn("--checkpoint_path", finetune) + for flag in ( + "--lr", + "--check_val_every_n_epoch", + "--checkpoints_dir", + "--use_downstream_head", + "--dryrun", + ): + self.assertIn(flag, finetune) + self.assertEqual( + finetune[finetune.index("--check_val_every_n_epoch") + 1], "1" + ) + + def test_no_unavailable_feature_in_runner_commands(self): + forbidden = { + "MissenseDataset", + "missense_prediction", + "missense_synom_agg", + "missense_inference", + "missense_seq", + "decodon_200m", + "decodon_1b", + "mutation_pred_clm", + } + for skill_name in PUBLIC_SKILLS: + text = (SKILLS_ROOT / skill_name / "SKILL.md").read_text() + for command in _runner_commands(text): + with self.subTest(skill=skill_name, command=command): + self.assertTrue(forbidden.isdisjoint(command)) + + def test_setup_usage_matches_public_script(self): + text = (SKILLS_ROOT / "codonfm-setup/SKILL.md").read_text() + self.assertIn("bash run_dev.sh", text) + self.assertIn("--data-dir", text) + self.assertIn("--checkpoints-dir", text) + self.assertIn("/data/checkpoints", text) + self.assertIn("hf download nvidia/NV-CodonFM-Encodon-1B-v1", text) + self.assertIn("python3.11 -m venv .venv", text) + self.assertIn("python -m pip install -r requirements.txt", text) + self.assertIn("export MPLCONFIGDIR=", text) + self.assertIn("Use only the checked-in public code", text) + + def test_safetensors_eval_is_not_loaded_as_a_lightning_checkpoint(self): + calls = {"torch_load": 0, "predict": 0} + + class FakeTorch: + @staticmethod + def load(*args, **kwargs): + calls["torch_load"] += 1 + raise AssertionError("torch.load must not read safetensors") + + class FakeLogger: + def log_hyperparams(self, config): + self.config = config + + class FakeData: + init_global_step = 0 + + def setup(self, stage): + self.stage = stage + + def load_state_dict(self, state): + self.state = state + + class FakeModel: + prediction_counter = 0 + + def configure_model(self): + self.configured = True + + class FakeTrainer: + def __init__(self, **kwargs): + self.kwargs = kwargs + + def predict(self, *args, **kwargs): + calls["predict"] += 1 + + namespace = { + "Any": object, + "Dict": dict, + "Path": Path, + "Trainer": FakeTrainer, + "logging": type("Logging", (), {"info": staticmethod(lambda message: None)}), + "os": __import__("os"), + "seed_everything": lambda *args, **kwargs: None, + "torch": FakeTorch, + } + evaluate = _evaluate_function(namespace) + + with tempfile.TemporaryDirectory() as temp_dir: + model_path = Path(temp_dir) / "model.safetensors" + model_path.touch() + data = FakeData() + model = FakeModel() + evaluate( + config={ + "log": FakeLogger(), + "data": data, + "trainer": {}, + "model": model, + "callbacks": {}, + }, + config_dict={}, + model_ckpt_path=str(model_path), + out_dir=temp_dir, + ) + + self.assertEqual(calls["torch_load"], 0) + self.assertEqual(calls["predict"], 1) + self.assertEqual(data.stage, "test") + self.assertTrue(model.configured) + + def test_checked_in_references_exist(self): + for relative_path in ( + "notebooks/4-EnCodon-Downstream-Task-riboNN.ipynb", + "notebooks/5-EnCodon-Downstream-Task-mRFP-expression.ipynb", + "notebooks/6-EnCodon-Downstream-Task-mRNA-stability.ipynb", + "src/data/codon_bert_dataset.py", + "src/data/mutation_dataset.py", + "src/inference/encodon.py", + ): + self.assertTrue((REPO_ROOT / relative_path).is_file(), relative_path) + + +if __name__ == "__main__": + unittest.main() From 87006b1e98828c0fed0925a29375d4e73e7b9245 Mon Sep 17 00:00:00 2001 From: Ohad Mosafi Date: Thu, 3 Sep 2026 16:01:33 -0700 Subject: [PATCH 2/3] Add author metadata Signed-off-by: Ohad Mosafi --- skills/codonfm-embed/SKILL.md | 2 ++ skills/codonfm-finetune/SKILL.md | 2 ++ skills/codonfm-score/SKILL.md | 2 ++ skills/codonfm-setup/SKILL.md | 2 ++ 4 files changed, 8 insertions(+) diff --git a/skills/codonfm-embed/SKILL.md b/skills/codonfm-embed/SKILL.md index 4290c22..32a5be1 100644 --- a/skills/codonfm-embed/SKILL.md +++ b/skills/codonfm-embed/SKILL.md @@ -1,6 +1,8 @@ --- name: codonfm-embed description: Extract frozen CLS embeddings from public CodonFM Encodon checkpoints for coding-sequence property modeling. Use when a user explicitly asks for CodonFM or Encodon embeddings, or wants Encodon features for translation-efficiency, expression, or mRNA-stability modeling. Support Encodon embedding_prediction only; do not claim Decodon embedding support in public CodonFM v1. +metadata: + author: "NVIDIA BioNeMo " --- # Extract public Encodon embeddings diff --git a/skills/codonfm-finetune/SKILL.md b/skills/codonfm-finetune/SKILL.md index 89aeb20..d32e6ef 100644 --- a/skills/codonfm-finetune/SKILL.md +++ b/skills/codonfm-finetune/SKILL.md @@ -1,6 +1,8 @@ --- name: codonfm-finetune description: Fine-tune public CodonFM Encodon checkpoints on labeled coding-sequence or coding-variant data using LoRA, head-only, or full fine-tuning. Use when a user explicitly asks to fine-tune CodonFM or Encodon for regression or classification. Support generic public-v1 Encodon workflows only; reject Decodon, MissenseDataset, missense_synom_agg, and generation workflows. +metadata: + author: "NVIDIA BioNeMo " --- # Fine-tune public Encodon diff --git a/skills/codonfm-score/SKILL.md b/skills/codonfm-score/SKILL.md index 725dc23..5f55a8a 100644 --- a/skills/codonfm-score/SKILL.md +++ b/skills/codonfm-score/SKILL.md @@ -1,6 +1,8 @@ --- name: codonfm-score description: Score synonymous or missense coding variants with public CodonFM Encodon checkpoints using masked-codon reference-versus-alternate log-likelihood ratios. Use when a user explicitly asks for CodonFM or Encodon zero-shot variant scoring. Support the public mutation_prediction workflow only; reject Decodon and the newer synonymous-codon-aggregated missense_prediction workflow because they are not present in public CodonFM v1. +metadata: + author: "NVIDIA BioNeMo " --- # Score variants with public Encodon diff --git a/skills/codonfm-setup/SKILL.md b/skills/codonfm-setup/SKILL.md index 788102e..b0c1749 100644 --- a/skills/codonfm-setup/SKILL.md +++ b/skills/codonfm-setup/SKILL.md @@ -1,6 +1,8 @@ --- name: codonfm-setup description: Set up the public CodonFM v1 repository and download public Encodon checkpoints. Use for requests to build or launch the CodonFM development container, configure local data/checkpoint mounts, verify GPU access, or download public Encodon 80M, 600M, 1B, or Cdwt-1B weights. Do not use for Decodon, Encodon 5B/10B, missense-aggregation, or codon-optimization setup because those implementations are not in the public repository. +metadata: + author: "NVIDIA BioNeMo " --- # CodonFM public setup From db15d24f2e23ae732a247321132f179c7001d10d Mon Sep 17 00:00:00 2001 From: Ohad Mosafi Date: Thu, 3 Sep 2026 16:07:33 -0700 Subject: [PATCH 3/3] nvskills workflow file Signed-off-by: Ohad Mosafi --- .github/workflows/request-nvskills-ci.yml | 26 +++++++++++++++++++++++ 1 file changed, 26 insertions(+) create mode 100644 .github/workflows/request-nvskills-ci.yml diff --git a/.github/workflows/request-nvskills-ci.yml b/.github/workflows/request-nvskills-ci.yml new file mode 100644 index 0000000..a88a057 --- /dev/null +++ b/.github/workflows/request-nvskills-ci.yml @@ -0,0 +1,26 @@ +name: Request NVSkills CI + +on: + issue_comment: + types: [created] + pull_request: + types: [opened, reopened, synchronize, ready_for_review] + push: + +jobs: + request: + if: > + github.event_name == 'pull_request' || + (github.event_name == 'issue_comment' && + github.event.issue.pull_request && + startsWith(github.event.comment.body, '/nvskills-ci')) || + (github.event_name == 'push' && + github.actor == (vars.NVSKILLS_SIGNATURE_PUSH_ACTOR || 'nv-skills-ci[bot]') && + startsWith(github.event.head_commit.message, vars.NVSKILLS_SIGNATURE_COMMIT_TITLE || 'Attach NVSkills validation signatures')) + permissions: + contents: read + pull-requests: read + statuses: read + uses: NVIDIA/skills/.github/workflows/team-request.yml@main + secrets: + NVSKILLS_CI_DISPATCH_TOKEN: ${{ secrets.NVSKILLS_CI_DISPATCH_TOKEN }} \ No newline at end of file