From 24a94010963b511889395942d524fc3580574081 Mon Sep 17 00:00:00 2001 From: Leo-T-Zang <85547488+Leo-T-Zang@users.noreply.github.com> Date: Mon, 24 Aug 2026 16:15:38 -0400 Subject: [PATCH 1/3] Score several assays per Evo 2 checkpoint load The 40B checkpoint is 76.6 GiB and takes over a minute to load, so one process per assay spends most of its time loading. Add --row_ids, which takes the same range syntax as a slurm array (0-8,11-32) and scores each assay in turn after a single load. --row_id is unchanged, so score_evo2.sh and any other caller keep working exactly as before. Two guards come with it. --require_fp8 aborts unless the constructed model reports use_fp8_input_projections=True: Evo2 falls back to bf16 for 7B checkpoints when Transformer Engine is unavailable, and it does so silently, so without the check a run can be bf16 while its provenance records FP8. The bf16 fallback is only valid for the 7B in any case, since the package refuses to build the other checkpoints without FP8. Scored CSVs are now written to a temporary file and renamed into place. A run that dies mid-write would otherwise leave a truncated CSV that the resume logic reads as a finished assay. Token budgeting also had to learn what the model actually processes. --max_tokens_per_batch divided by the raw sequence length, but a BOS token adds one position and FP8 pads the sequence dimension to a multiple of 16 inside every input projection, so an assay of length 87 runs at 96. effective_length accounts for both. --- .../baselines/Evo/score_evo2_single_dms.py | 314 ++++++++++++------ fitness/model_registry.py | 4 + fitness/performance_fitness.py | 1 - 3 files changed, 225 insertions(+), 94 deletions(-) diff --git a/fitness/baselines/Evo/score_evo2_single_dms.py b/fitness/baselines/Evo/score_evo2_single_dms.py index a550541..2e31a2c 100644 --- a/fitness/baselines/Evo/score_evo2_single_dms.py +++ b/fitness/baselines/Evo/score_evo2_single_dms.py @@ -1,6 +1,6 @@ #!/usr/bin/env python3 """ -Score a single RNAGym DMS assay with an Evo 2 model. +Score RNAGym DMS assays with an Evo 2 model. Evo 2 is an autoregressive genomic language model (StripedHyena 2). For each variant sequence we compute the mean per-token log-likelihood under the model @@ -9,17 +9,16 @@ This is the Evo 2 counterpart of ``score_evo_single_dms.py`` (Evo 1 / 1.5). It uses the official ``evo2`` package (https://github.com/ArcInstitute/evo2) rather -than the ``evo`` package, and supports the large ``evo2_40b`` checkpoint, which -requires FP8 via Transformer Engine on Hopper GPUs and is automatically sharded -across every visible GPU by the Vortex inference engine. +than the ``evo`` package, and supports every checkpoint in ``evo2.utils.MODEL_NAMES`` +(``evo2_1b_base``, ``evo2_7b``, ``evo2_20b``, ``evo2_40b``, ...). Notes on multi-GPU ------------------ Vortex places and (for large models) shards the model across all CUDA devices that are visible. Select the GPUs with ``CUDA_VISIBLE_DEVICES`` and do NOT move the model manually with ``.to(device)``. ``evo2_40b`` does not fit on a single -80 GB GPU and needs at least two (e.g. 2xH100-80GB — but note the -40B/20B/1B checkpoints require FP8 + Transformer Engine, i.e. a Hopper GPU). +80 GB GPU and needs at least two (e.g. 2xH100-80GB). The 40B/20B/7B/1B +checkpoints all request FP8 via Transformer Engine, i.e. a Hopper GPU. Offline weights --------------- @@ -28,6 +27,15 @@ pre-merge the checkpoint once (see ``download_weights.sh``) and pass the merged file via ``--local_path`` so no network access is needed at run time. +Batching +-------- +All variants of one assay share a single sequence length, so batches never need +padding between sequences and the batch size cannot change a score beyond +floating point roundoff. ``--max_tokens_per_batch`` sizes each batch by a token +budget rather than a sequence count. When FP8 input projections are enabled, +Vortex pads the sequence dimension up to a multiple of 16 inside every +projection, so the budget is applied to that padded length. + Usage ----- python score_evo2_single_dms.py \ @@ -38,11 +46,16 @@ --model_name evo2_40b \ --local_path /path/to/evo2_40b.pt \ --batch_size 1 + + # several assays in one process, so the checkpoint is loaded once + python score_evo2_single_dms.py --row_ids 0-8,11-32 ... """ import argparse +import math import os import sys +import tempfile from pathlib import Path import numpy as np @@ -63,17 +76,61 @@ def preprocess_sequence(sequence: str) -> str: return sequence.strip().upper().replace("U", "T") +def parse_row_ids(spec: str) -> list: + """Parse a row selection such as ``0-8,11-32,40`` into a sorted list. + + Empty components are rejected rather than skipped: ``0-8,,11`` is far more + likely to be a typo than an intention, and silently dropping it would score + a different set of assays than the caller asked for. + """ + rows = set() + for part in spec.split(","): + part = part.strip() + if not part: + raise ValueError(f"Empty component in --row_ids: {spec!r}") + if "-" in part.lstrip("-"): + start, end = part.split("-", 1) + start, end = int(start), int(end) + if end < start: + raise ValueError(f"Empty range in --row_ids: {part}") + rows.update(range(start, end + 1)) + else: + rows.add(int(part)) + if not rows: + raise ValueError(f"No rows selected by --row_ids {spec!r}") + return sorted(rows) + + +def effective_length(seq_len: int, prepend_bos: bool, fp8: bool) -> int: + """The sequence length the model actually processes. + + ``prepare_batch`` prepends one token when ``prepend_bos`` is set, and Vortex's + ``pad_to_multiple`` pads the sequence dimension to a multiple of 16 inside + every input projection when FP8 is enabled. + """ + length = seq_len + int(prepend_bos) + if fp8: + length = 16 * math.ceil(length / 16) + return length + + def parse_args(): """Parse command line arguments.""" parser = argparse.ArgumentParser( - description="Run Evo 2 inference on the sequences of a single DMS assay." + description="Run Evo 2 inference on the sequences of one or more DMS assays." ) - parser.add_argument( + rows = parser.add_mutually_exclusive_group(required=True) + rows.add_argument( "--row_id", type=int, - required=True, help="Row ID in the reference sheet to process", ) + rows.add_argument( + "--row_ids", + type=str, + help="Several reference sheet rows, e.g. '0-8,11-32'. They are scored in " + "one process so the checkpoint is loaded once.", + ) parser.add_argument( "--ref_sheet", type=str, @@ -120,9 +177,10 @@ def parse_args(): type=int, default=None, help="If set, the batch size is derived per assay as " - "max(1, max_tokens_per_batch // seq_len), overriding --batch_size. Keeps " - "GPU memory roughly constant across assays of very different lengths " - "while maximising throughput (e.g. 8192).", + "max(1, max_tokens_per_batch // effective_length), overriding " + "--batch_size. The effective length accounts for the BOS token and for " + "Vortex's multiple-of-16 padding under FP8. Keeps GPU memory roughly " + "constant across assays of very different lengths (e.g. 8192).", ) parser.add_argument( "--reduce_method", @@ -149,6 +207,14 @@ def parse_args(): "comparison. Pass --no-average_reverse_complement for forward strand " "only (~2x faster).", ) + parser.add_argument( + "--require_fp8", + action="store_true", + help="Abort unless the model was actually built with FP8 input " + "projections. Evo2.load_evo2_model silently falls back to bf16 for 7B " + "checkpoints when Transformer Engine is unavailable, so without this a " + "run can be bf16 while everything around it records FP8.", + ) parser.add_argument( "--overwrite", action="store_true", @@ -157,8 +223,8 @@ def parse_args(): return parser.parse_args() -def load_reference_data(ref_sheet_path: str, row_id: int) -> str: - """Return the DMS_ID for ``row_id`` in the reference sheet.""" +def load_reference_data(ref_sheet_path: str, row_ids) -> list: + """Return the DMS_IDs for ``row_ids`` in the reference sheet.""" try: ref_df = pd.read_csv(ref_sheet_path) except FileNotFoundError: @@ -168,15 +234,18 @@ def load_reference_data(ref_sheet_path: str, row_id: int) -> str: ref_df.columns = [c.lstrip("") for c in ref_df.columns] if "DMS_ID" not in ref_df.columns: raise KeyError("Reference sheet must contain a 'DMS_ID' column") - if row_id < 0 or row_id >= len(ref_df): - raise ValueError( - f"Row ID {row_id} out of range (reference sheet has {len(ref_df)} rows)" - ) - dms_id = ref_df.loc[row_id, "DMS_ID"] - if pd.isna(dms_id): - raise ValueError(f"DMS_ID is missing for row {row_id}") - return str(dms_id) + dms_ids = [] + for row_id in row_ids: + if row_id < 0 or row_id >= len(ref_df): + raise ValueError( + f"Row ID {row_id} out of range (reference sheet has {len(ref_df)} rows)" + ) + dms_id = ref_df.loc[row_id, "DMS_ID"] + if pd.isna(dms_id): + raise ValueError(f"DMS_ID is missing for row {row_id}") + dms_ids.append(str(dms_id)) + return dms_ids def load_dms_data(dms_dir_path: str, dms_id: str) -> pd.DataFrame: @@ -193,87 +262,146 @@ def load_dms_data(dms_dir_path: str, dms_id: str) -> pd.DataFrame: return df +def write_csv_atomically(df, output_file): + """Write the scored assay, then rename it into place. + + These CSVs are the published scores, so a partial file must never appear + under the final name: an interrupted or out-of-quota write would otherwise + leave a truncated CSV that later looks like a completed assay to the + resume logic, and --overwrite would destroy a good file to produce it. + """ + output_file = Path(output_file) + handle, tmp_path = tempfile.mkstemp(dir=str(output_file.parent), + prefix=f".{output_file.name}.", suffix=".tmp") + os.close(handle) + try: + df.to_csv(tmp_path, index=False) + os.replace(tmp_path, output_file) + except BaseException: + if os.path.exists(tmp_path): + os.remove(tmp_path) + raise + + +def score_one_assay(evo2_model, args, dms_id, fp8_enabled): + """Score one assay and write its CSV. Returns the Spearman correlation.""" + output_file = Path(args.output_dir_path) / f"{dms_id}.csv" + dms_df = load_dms_data(args.dms_dir_path, dms_id) + + # Preprocess sequences (RNA -> DNA), tracking any rows we cannot score. + print("Preprocessing sequences...") + raw = dms_df["sequence"] + valid_mask = raw.notna() & (raw.astype(str).str.strip() != "") + n_skipped = int((~valid_mask).sum()) + if n_skipped: + print(f"Skipping {n_skipped} rows with empty/NaN sequence") + sequences = [preprocess_sequence(s) for s in raw[valid_mask].astype(str)] + max_len = max((len(s) for s in sequences), default=0) + print(f"Scoring {len(sequences)} sequences (max length {max_len} nt)") + + # Choose the batch size (token-budget adaptive if requested). + batch_size = args.batch_size + if args.max_tokens_per_batch is not None and sequences: + eff_len = effective_length(max_len, args.prepend_bos, fp8_enabled) + batch_size = max(1, args.max_tokens_per_batch // eff_len) + print(f"Token budget {args.max_tokens_per_batch}: seq_len={max_len} " + f"prepend_bos={args.prepend_bos} fp8={fp8_enabled} " + f"effective_length={eff_len} -> batch_size={batch_size}") + + print(f"Running inference (batch_size={batch_size}, " + f"reduce_method={args.reduce_method}, prepend_bos={args.prepend_bos}, " + f"rc={args.average_reverse_complement})...") + scores = evo2_model.score_sequences( + sequences, + batch_size=batch_size, + prepend_bos=args.prepend_bos, + reduce_method=args.reduce_method, + average_reverse_complement=args.average_reverse_complement, + ) + scores = np.asarray(scores, dtype=float) + + # Write scores back onto the scored rows (NaN for skipped ones). + score_column = f"{args.model_name}_score" + dms_df[score_column] = np.nan + dms_df.loc[valid_mask, score_column] = scores + + # Spearman on the rows we actually scored. + scored = dms_df.loc[valid_mask, ["DMS_score", score_column]].dropna() + if len(scored) >= 2: + correlation, pvalue = spearmanr(scored["DMS_score"], scored[score_column]) + else: + correlation, pvalue = float("nan"), float("nan") + + write_csv_atomically(dms_df, output_file) + + print("\nSummary:") + print(f" DMS ID: {dms_id}") + print(f" Sequences scored: {len(sequences)}") + print(f" Score column: {score_column}") + print(f" Spearman vs DMS: {correlation:.3f} (p={pvalue:.2e})") + print(f" Saved to: {output_file}") + return correlation + + def main(): args = parse_args() output_dir = Path(args.output_dir_path) output_dir.mkdir(parents=True, exist_ok=True) - try: - dms_id = load_reference_data(args.ref_sheet, args.row_id) - print(f"Processing DMS ID: {dms_id}") + row_ids = [args.row_id] if args.row_id is not None else parse_row_ids(args.row_ids) + dms_ids = load_reference_data(args.ref_sheet, row_ids) + print(f"Rows {row_ids} -> DMS IDs: {dms_ids}") + + todo = [] + for row_id, dms_id in zip(row_ids, dms_ids): output_file = output_dir / f"{dms_id}.csv" if output_file.exists() and not args.overwrite: print(f"Output already exists (use --overwrite to redo): {output_file}") - return - - dms_df = load_dms_data(args.dms_dir_path, dms_id) - - # Preprocess sequences (RNA -> DNA), tracking any rows we cannot score. - print("Preprocessing sequences...") - raw = dms_df["sequence"] - valid_mask = raw.notna() & (raw.astype(str).str.strip() != "") - n_skipped = int((~valid_mask).sum()) - if n_skipped: - print(f"Skipping {n_skipped} rows with empty/NaN sequence") - sequences = [preprocess_sequence(s) for s in raw[valid_mask].astype(str)] - print(f"Scoring {len(sequences)} sequences (max length " - f"{max((len(s) for s in sequences), default=0)} nt)") - - if not torch.cuda.is_available(): - print("WARNING: CUDA not available — Evo 2 requires a GPU.", - file=sys.stderr) - print(f"Visible GPUs: {torch.cuda.device_count()}") - - # Initialize model. Vortex handles device placement / multi-GPU sharding; - # do NOT call .to(device). - print(f"Loading Evo 2 model: {args.model_name} " - f"(local_path={args.local_path})...") - evo2_model = Evo2(args.model_name, local_path=args.local_path) - - # Choose the batch size (token-budget adaptive if requested). - batch_size = args.batch_size - if args.max_tokens_per_batch is not None and sequences: - max_len = max(len(s) for s in sequences) - batch_size = max(1, args.max_tokens_per_batch // max_len) - - print(f"Running inference (batch_size={batch_size}, " - f"reduce_method={args.reduce_method}, prepend_bos={args.prepend_bos}, " - f"rc={args.average_reverse_complement})...") - scores = evo2_model.score_sequences( - sequences, - batch_size=batch_size, - prepend_bos=args.prepend_bos, - reduce_method=args.reduce_method, - average_reverse_complement=args.average_reverse_complement, - ) - scores = np.asarray(scores, dtype=float) - - # Write scores back onto the scored rows (NaN for skipped ones). - score_column = f"{args.model_name}_score" - dms_df[score_column] = np.nan - dms_df.loc[valid_mask, score_column] = scores - - # Spearman on the rows we actually scored. - scored = dms_df.loc[valid_mask, ["DMS_score", score_column]].dropna() - if len(scored) >= 2: - correlation, pvalue = spearmanr(scored["DMS_score"], scored[score_column]) - else: - correlation, pvalue = float("nan"), float("nan") - - dms_df.to_csv(output_file, index=False) - - print("\nSummary:") - print(f" DMS ID: {dms_id}") - print(f" Sequences scored: {len(sequences)}") - print(f" Score column: {score_column}") - print(f" Spearman vs DMS: {correlation:.3f} (p={pvalue:.2e})") - print(f" Saved to: {output_file}") - - except Exception as e: - print(f"Error: {str(e)}", file=sys.stderr) - raise + continue + todo.append((row_id, dms_id)) + if not todo: + print("Nothing to score.") + return + + if not torch.cuda.is_available(): + print("WARNING: CUDA not available - Evo 2 requires a GPU.", file=sys.stderr) + print(f"Visible GPUs: {torch.cuda.device_count()}") + + # Initialize model. Vortex handles device placement / multi-GPU sharding; + # do NOT call .to(device). The checkpoint is loaded once for every assay. + print(f"Loading Evo 2 model: {args.model_name} " + f"(local_path={args.local_path})...") + evo2_model = Evo2(args.model_name, local_path=args.local_path) + # Always ask the built model, never the packaged YAML: load_evo2_model can + # turn FP8 off for 7B when Transformer Engine is missing, and the batch-size + # arithmetic below has to follow the config the model was actually built with. + config = evo2_model.model.config + fp8_enabled = bool(config.get("use_fp8_input_projections", False)) + print(f"use_fp8_input_projections={fp8_enabled}") + if args.require_fp8 and not fp8_enabled: + raise SystemExit( + "--require_fp8 was given but the model resolved to " + "use_fp8_input_projections=False. For a 7B checkpoint this happens " + "silently when Transformer Engine is unavailable; for the others it " + "means Transformer Engine is not providing FP8. Refusing to score, because " + "the surrounding provenance would claim FP8.") + + failures = [] + for row_id, dms_id in todo: + print(f"\n=== row {row_id}: {dms_id} ===") + try: + score_one_assay(evo2_model, args, dms_id, fp8_enabled) + except Exception as e: + print(f"Error scoring {dms_id}: {str(e)}", file=sys.stderr) + failures.append(dms_id) + if len(todo) == 1: + raise + + if failures: + print(f"\nFAILED assays ({len(failures)}): {failures}", file=sys.stderr) + sys.exit(1) if __name__ == "__main__": diff --git a/fitness/model_registry.py b/fitness/model_registry.py index a1f57f2..46d0af4 100644 --- a/fitness/model_registry.py +++ b/fitness/model_registry.py @@ -4,6 +4,8 @@ "evo1", "evo1.5", "evo2", + "evo2_1b_base", + "evo2_20b", "evo2_40b", "GenSLM", "NT", @@ -23,6 +25,8 @@ "evo1": "evo_1_131k_base_score", "evo1.5": "evo_1.5_8k_base_score", "evo2": "evo2_7b_score", + "evo2_1b_base": "evo2_1b_base_score", + "evo2_20b": "evo2_20b_score", "evo2_40b": "evo2_40b_score", "GenSLM": "logit_scores", "NT": "kmer_pseudo_LL", diff --git a/fitness/performance_fitness.py b/fitness/performance_fitness.py index 916841b..aa215b8 100755 --- a/fitness/performance_fitness.py +++ b/fitness/performance_fitness.py @@ -332,7 +332,6 @@ def save_results( ) transposed.to_csv(output_dir / "assay_level_results_transposed.csv", index=False) - def select_assays(reference: pd.DataFrame, assay_type: str) -> pd.DataFrame: """Select an assay group.""" if assay_type == "all": From 6c77558bf5620016ffa02e6400fe2607de25d605 Mon Sep 17 00:00:00 2001 From: Leo-T-Zang <85547488+Leo-T-Zang@users.noreply.github.com> Date: Mon, 24 Aug 2026 16:15:38 -0400 Subject: [PATCH 2/3] Add the Evo 2 1B and 20B checkpoints to the leaderboard The leaderboard carried the 7B and the 40B. Adding the 1B and the 20B completes the family, so the four sizes can be read together on the same 31 ncRNA assays. evo2_20b scores 0.2129, second behind aido_rna_650m and ahead of the 40B. evo2_1b_base scores 0.1962, which needs a caveat: its Ribozyme mean is 0.0087 and the macro weights the three categories equally, so a strong tRNA mean over only three assays carries it. It is also the only base-pretrained 8k-context checkpoint here, where the 7B, 20B and 40B are mid-trained at 1M, so it is not a size point comparable to the rest. Both were scored with the packaged FP8 configuration, which is what the Evo 2 authors require for every checkpoint except the 7B. The 7B and the 40B were rescored the same way as a check and are not published: they reproduce the existing columns, the 40B bitwise on all 31 assays and the 7B to the resolution of its own float32 serialisation, so publishing them would add rows that duplicate what is already there. That reproduction is what allows the existing 7B and 40B columns to be read alongside these two as one comparable family. The row count in test_every_leaderboard_row_is_a_registered_model goes from 16 to 18 with them. --- leaderboard/fitness/leaderboard_signed_3ncRNA.csv | 2 ++ 1 file changed, 2 insertions(+) diff --git a/leaderboard/fitness/leaderboard_signed_3ncRNA.csv b/leaderboard/fitness/leaderboard_signed_3ncRNA.csv index 0326f70..8913c4e 100644 --- a/leaderboard/fitness/leaderboard_signed_3ncRNA.csv +++ b/leaderboard/fitness/leaderboard_signed_3ncRNA.csv @@ -1,7 +1,9 @@ model,Ribozyme,tRNA,Aptamer,macro_3ncRNA aido_rna_650m,0.0660332868987325,0.4893565654410639,0.0934325517551516,0.216274134698316 +evo2_20b,0.1104601596335257,0.4345270738070895,0.0937093504359844,0.2128988612921999 evo2_40b,0.1081020224381099,0.4309848333293184,0.09695482733614616,0.21201389436785814 aido_rna,0.0608547145460644,0.4883787106096389,0.0718036455462622,0.2070123569006552 +evo2_1b_base,0.0086596951000284,0.4466199808506075,0.1334226383526211,0.1962341047677523 RNAErnie,0.13427241327809655,0.416115138737617,0.030641733377440288,0.1936764284643846 evo2,0.06511905079221095,0.3867345488695757,0.11920502912796216,0.1903528762632496 rnagenesis,0.0750044159987944,0.4380555746849195,0.0343114577905016,0.1824571494914051 From de3d738f77582476c408c21da44b84fba699e7ca Mon Sep 17 00:00:00 2001 From: murfalo Date: Fri, 28 Aug 2026 18:45:35 -0400 Subject: [PATCH 3/3] Fitness PR 2 cleanup --- fitness/README.md | 8 +- .../baselines/Evo/score_evo2_single_dms.py | 490 +++++++----------- fitness/performance_fitness.py | 1 + fitness/sh/coverage.sh | 1 + fitness/sh/format.sh | 2 + fitness/sh/lint.sh | 2 + leaderboard/fitness/README.md | 41 +- tests/test_fitness.py | 83 +++ 8 files changed, 303 insertions(+), 325 deletions(-) diff --git a/fitness/README.md b/fitness/README.md index 22827ac..ff68c8a 100644 --- a/fitness/README.md +++ b/fitness/README.md @@ -114,9 +114,9 @@ pixi run fmt ## Models -The default merge registry includes Evo 1, Evo 1.5, Evo 2 (7B and 40B), GenSLM, -Nucleotide Transformer, RNA-ERNIE, RNA-FM, RiNALMo, RNAGenesis, Orthrus and the -five released AIDO.RNA checkpoints. EVmutation is merged separately with -`--assays_with_MSAs_only` because it only covers assays with an MSA. +The default merge registry includes Evo 1, Evo 1.5, Evo 2 (1B base, 7B, 20B and +40B), GenSLM, Nucleotide Transformer, RNA-ERNIE, RNA-FM, RiNALMo, RNAGenesis, +Orthrus and the five released AIDO.RNA checkpoints. EVmutation is merged +separately with `--assays_with_MSAs_only` because it only covers assays with an MSA. [0]: https://marks.hms.harvard.edu/rnagym/fitness_prediction diff --git a/fitness/baselines/Evo/score_evo2_single_dms.py b/fitness/baselines/Evo/score_evo2_single_dms.py index 2e31a2c..cd57f92 100644 --- a/fitness/baselines/Evo/score_evo2_single_dms.py +++ b/fitness/baselines/Evo/score_evo2_single_dms.py @@ -1,407 +1,293 @@ #!/usr/bin/env python3 -""" -Score RNAGym DMS assays with an Evo 2 model. - -Evo 2 is an autoregressive genomic language model (StripedHyena 2). For each -variant sequence we compute the mean per-token log-likelihood under the model -and use it as the fitness score, then report the Spearman correlation against -the experimental ``DMS_score``. - -This is the Evo 2 counterpart of ``score_evo_single_dms.py`` (Evo 1 / 1.5). It -uses the official ``evo2`` package (https://github.com/ArcInstitute/evo2) rather -than the ``evo`` package, and supports every checkpoint in ``evo2.utils.MODEL_NAMES`` -(``evo2_1b_base``, ``evo2_7b``, ``evo2_20b``, ``evo2_40b``, ...). - -Notes on multi-GPU ------------------- -Vortex places and (for large models) shards the model across all CUDA devices -that are visible. Select the GPUs with ``CUDA_VISIBLE_DEVICES`` and do NOT move -the model manually with ``.to(device)``. ``evo2_40b`` does not fit on a single -80 GB GPU and needs at least two (e.g. 2xH100-80GB). The 40B/20B/7B/1B -checkpoints all request FP8 via Transformer Engine, i.e. a Hopper GPU. - -Offline weights ---------------- -The 40B checkpoint ships as two ~41 GB shards that ``evo2`` merges into a single -``evo2_40b.pt`` on first load (a network call). On air-gapped compute nodes, -pre-merge the checkpoint once (see ``download_weights.sh``) and pass the merged -file via ``--local_path`` so no network access is needed at run time. - -Batching --------- -All variants of one assay share a single sequence length, so batches never need -padding between sequences and the batch size cannot change a score beyond -floating point roundoff. ``--max_tokens_per_batch`` sizes each batch by a token -budget rather than a sequence count. When FP8 input projections are enabled, -Vortex pads the sequence dimension up to a multiple of 16 inside every -projection, so the budget is applied to that padded length. - -Usage ------ - python score_evo2_single_dms.py \ - --row_id 0 \ - --ref_sheet reference_sheet_final.csv \ - --dms_dir_path fitness_processed_assays \ - --output_dir_path evo2_40b_output \ - --model_name evo2_40b \ - --local_path /path/to/evo2_40b.pt \ - --batch_size 1 - - # several assays in one process, so the checkpoint is loaded once - python score_evo2_single_dms.py --row_ids 0-8,11-32 ... +"""Score one or more RNAGym assays with the official Evo 2 predictor. + +Multiple ``--row_ids`` share one model load. Vortex controls model placement +and sharding across the visible GPUs, so this script never moves the model. """ import argparse -import math -import os +import re import sys -import tempfile from pathlib import Path +from tempfile import NamedTemporaryFile import numpy as np import pandas as pd import torch from scipy.stats import spearmanr -from evo2 import Evo2 +def _nonnegative_int(value: str) -> int: + """Parse a nonnegative command-line integer.""" + parsed = int(value) + if parsed < 0: + raise argparse.ArgumentTypeError("must be nonnegative") + return parsed -def preprocess_sequence(sequence: str) -> str: - """Preprocess an RNA/DNA sequence for the Evo 2 (DNA) model. - - Convert RNA (U) to DNA (T) - - Uppercase - - Strip surrounding whitespace - """ - return sequence.strip().upper().replace("U", "T") +def _positive_int(value: str) -> int: + """Parse a positive command-line integer.""" + parsed = int(value) + if parsed < 1: + raise argparse.ArgumentTypeError("must be positive") + return parsed -def parse_row_ids(spec: str) -> list: - """Parse a row selection such as ``0-8,11-32,40`` into a sorted list. +def effective_length(sequence_length: int, prepend_bos: bool, fp8: bool) -> int: + """Return the sequence dimension used by the largest input projection.""" + length = sequence_length + int(prepend_bos) + return ((length + 15) // 16) * 16 if fp8 else length - Empty components are rejected rather than skipped: ``0-8,,11`` is far more - likely to be a typo than an intention, and silently dropping it would score - a different set of assays than the caller asked for. - """ - rows = set() - for part in spec.split(","): - part = part.strip() - if not part: - raise ValueError(f"Empty component in --row_ids: {spec!r}") - if "-" in part.lstrip("-"): - start, end = part.split("-", 1) - start, end = int(start), int(end) - if end < start: - raise ValueError(f"Empty range in --row_ids: {part}") - rows.update(range(start, end + 1)) - else: - rows.add(int(part)) - if not rows: - raise ValueError(f"No rows selected by --row_ids {spec!r}") - return sorted(rows) +def load_dms_data(dms_dir: Path, dms_id: str) -> pd.DataFrame: + """Load and validate an assay table.""" + path = dms_dir / f"{dms_id}.csv" + data = pd.read_csv(path) + required = {"mutant", "DMS_score", "sequence"} + missing = sorted(required - set(data.columns)) + if missing: + raise ValueError(f"{path} is missing columns: {missing}") + return data -def effective_length(seq_len: int, prepend_bos: bool, fp8: bool) -> int: - """The sequence length the model actually processes. - ``prepare_batch`` prepends one token when ``prepend_bos`` is set, and Vortex's - ``pad_to_multiple`` pads the sequence dimension to a multiple of 16 inside - every input projection when FP8 is enabled. - """ - length = seq_len + int(prepend_bos) - if fp8: - length = 16 * math.ceil(length / 16) - return length +def load_reference_data(reference_file: Path, row_ids: list[int]) -> list[str]: + """Return the DMS IDs at the selected reference-sheet rows.""" + reference = pd.read_csv(reference_file, encoding="utf-8-sig") + if "DMS_ID" not in reference: + raise ValueError(f"{reference_file} is missing column 'DMS_ID'") + invalid = [row_id for row_id in row_ids if row_id < 0 or row_id >= len(reference)] + if invalid: + raise ValueError( + f"Reference rows {invalid} fall outside 0-{len(reference) - 1}" + ) + dms_ids = reference.iloc[row_ids]["DMS_ID"] + missing = [row_id for row_id, dms_id in zip(row_ids, dms_ids) if pd.isna(dms_id)] + if missing: + raise ValueError(f"DMS_ID is missing for reference rows {missing}") + return dms_ids.astype(str).tolist() -def parse_args(): - """Parse command line arguments.""" + +def parse_args(argv=None) -> argparse.Namespace: + """Parse command-line arguments.""" parser = argparse.ArgumentParser( - description="Run Evo 2 inference on the sequences of one or more DMS assays." + description="Score one or more RNAGym assays with Evo 2", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, ) rows = parser.add_mutually_exclusive_group(required=True) rows.add_argument( "--row_id", - type=int, - help="Row ID in the reference sheet to process", + type=_nonnegative_int, + help="Reference-sheet row to score", ) rows.add_argument( "--row_ids", - type=str, - help="Several reference sheet rows, e.g. '0-8,11-32'. They are scored in " - "one process so the checkpoint is loaded once.", + help="Comma-separated rows and ranges to score in one model load, e.g. 0-8,11-32", ) parser.add_argument( "--ref_sheet", - type=str, + type=Path, required=True, - help="Path to reference sheet containing a DMS_ID column", + help="Reference sheet containing DMS_ID", ) parser.add_argument( "--dms_dir_path", - type=str, + type=Path, required=True, - help="Directory containing DMS assay CSV files ({DMS_ID}.csv)", + help="Directory containing assay CSVs", ) parser.add_argument( "--output_dir_path", - type=str, + type=Path, required=True, - help="Directory to save the scored output CSV", + help="Directory for scored CSVs", ) parser.add_argument( "--model_name", - type=str, default="evo2_40b", - help="Evo 2 checkpoint name (default: evo2_40b). The score column is " - "named '{model_name}_score', e.g. evo2_40b_score.", + help="Evo 2 checkpoint name", ) parser.add_argument( "--local_path", - type=str, - default=None, - help="Path to a pre-merged Evo 2 .pt checkpoint. When given, the model " - "is loaded fully offline (no HuggingFace network access). Recommended " - "for air-gapped compute nodes.", + help="Optional local checkpoint path", ) parser.add_argument( "--batch_size", - type=int, + type=_positive_int, default=1, - help="Number of sequences scored per forward pass (default: 1). All " - "variants of an assay share one length, so batching is padding-free; " - "raise it for short assays, keep it small for very long ones / 40B.", + help="Sequences per forward pass", ) parser.add_argument( "--max_tokens_per_batch", - type=int, - default=None, - help="If set, the batch size is derived per assay as " - "max(1, max_tokens_per_batch // effective_length), overriding " - "--batch_size. The effective length accounts for the BOS token and for " - "Vortex's multiple-of-16 padding under FP8. Keeps GPU memory roughly " - "constant across assays of very different lengths (e.g. 8192).", + type=_positive_int, + help="Derive each assay's batch size from this token budget", ) parser.add_argument( "--reduce_method", - type=str, + choices=("mean", "sum"), default="mean", - choices=["mean", "sum"], - help="Reduce per-token log-likelihoods by mean (mean PLL, default) or " - "sum (PLL).", + help="Per-sequence log-likelihood reduction", ) parser.add_argument( "--prepend_bos", action="store_true", - help="Prepend the BOS/EOD token before scoring (default: off, matching " - "the evo2 package default).", + help="Prepend the BOS/EOD token", ) parser.add_argument( "--average_reverse_complement", action=argparse.BooleanOptionalAction, default=True, - help="Score each sequence as the mean of its forward and " - "reverse-complement log-likelihood (default: ON). Evo 2 is a " - "strand-symmetric DNA model and the RNAGym evo2 baselines use " - "reverse-complement averaging, so this is the default for a fair " - "comparison. Pass --no-average_reverse_complement for forward strand " - "only (~2x faster).", + help="Average forward and reverse-complement scores", ) parser.add_argument( "--require_fp8", action="store_true", - help="Abort unless the model was actually built with FP8 input " - "projections. Evo2.load_evo2_model silently falls back to bf16 for 7B " - "checkpoints when Transformer Engine is unavailable, so without this a " - "run can be bf16 while everything around it records FP8.", + help="Require FP8 input projections in the constructed model", ) parser.add_argument( "--overwrite", action="store_true", - help="Re-score even if the output CSV already exists.", + help="Replace existing scored CSVs", ) - return parser.parse_args() + return parser.parse_args(argv) -def load_reference_data(ref_sheet_path: str, row_ids) -> list: - """Return the DMS_IDs for ``row_ids`` in the reference sheet.""" - try: - ref_df = pd.read_csv(ref_sheet_path) - except FileNotFoundError: - raise FileNotFoundError(f"Reference sheet not found: {ref_sheet_path}") - - # Tolerate a UTF-8 BOM on the DMS_ID column header. - ref_df.columns = [c.lstrip("") for c in ref_df.columns] - if "DMS_ID" not in ref_df.columns: - raise KeyError("Reference sheet must contain a 'DMS_ID' column") - - dms_ids = [] - for row_id in row_ids: - if row_id < 0 or row_id >= len(ref_df): - raise ValueError( - f"Row ID {row_id} out of range (reference sheet has {len(ref_df)} rows)" - ) - dms_id = ref_df.loc[row_id, "DMS_ID"] - if pd.isna(dms_id): - raise ValueError(f"DMS_ID is missing for row {row_id}") - dms_ids.append(str(dms_id)) - return dms_ids - - -def load_dms_data(dms_dir_path: str, dms_id: str) -> pd.DataFrame: - """Load the DMS assay CSV for ``dms_id``.""" - dms_file = Path(dms_dir_path) / f"{dms_id}.csv" - if not dms_file.exists(): - raise FileNotFoundError(f"DMS file not found: {dms_file}") - - df = pd.read_csv(dms_file) - required_cols = ["mutant", "DMS_score", "sequence"] - missing_cols = [col for col in required_cols if col not in df.columns] - if missing_cols: - raise ValueError(f"Missing required columns in DMS file: {missing_cols}") - return df - - -def write_csv_atomically(df, output_file): - """Write the scored assay, then rename it into place. - - These CSVs are the published scores, so a partial file must never appear - under the final name: an interrupted or out-of-quota write would otherwise - leave a truncated CSV that later looks like a completed assay to the - resume logic, and --overwrite would destroy a good file to produce it. - """ - output_file = Path(output_file) - handle, tmp_path = tempfile.mkstemp(dir=str(output_file.parent), - prefix=f".{output_file.name}.", suffix=".tmp") - os.close(handle) - try: - df.to_csv(tmp_path, index=False) - os.replace(tmp_path, output_file) - except BaseException: - if os.path.exists(tmp_path): - os.remove(tmp_path) - raise - - -def score_one_assay(evo2_model, args, dms_id, fp8_enabled): - """Score one assay and write its CSV. Returns the Spearman correlation.""" - output_file = Path(args.output_dir_path) / f"{dms_id}.csv" - dms_df = load_dms_data(args.dms_dir_path, dms_id) - - # Preprocess sequences (RNA -> DNA), tracking any rows we cannot score. - print("Preprocessing sequences...") - raw = dms_df["sequence"] - valid_mask = raw.notna() & (raw.astype(str).str.strip() != "") - n_skipped = int((~valid_mask).sum()) - if n_skipped: - print(f"Skipping {n_skipped} rows with empty/NaN sequence") - sequences = [preprocess_sequence(s) for s in raw[valid_mask].astype(str)] - max_len = max((len(s) for s in sequences), default=0) - print(f"Scoring {len(sequences)} sequences (max length {max_len} nt)") - - # Choose the batch size (token-budget adaptive if requested). +def parse_row_ids(specification: str) -> list[int]: + """Expand a row selection such as ``0-8,11-32``.""" + rows = set() + for selection in specification.split(","): + match = re.fullmatch(r"\s*(\d+)(?:-(\d+))?\s*", selection) + if match is None: + raise ValueError(f"Invalid --row_ids selection: {selection!r}") + start = int(match.group(1)) + end = int(match.group(2) or start) + if end < start: + raise ValueError(f"Descending --row_ids range: {selection!r}") + rows.update(range(start, end + 1)) + return sorted(rows) + + +def preprocess_sequence(sequence: str) -> str: + """Convert an RNA or DNA sequence to uppercase DNA.""" + return sequence.strip().upper().replace("U", "T") + + +def score_assay(model, args: argparse.Namespace, dms_id: str, fp8: bool) -> float: + """Score one assay, write its prediction table, and return Spearman rho.""" + data = load_dms_data(args.dms_dir_path, dms_id) + raw_sequences = data["sequence"] + valid = raw_sequences.notna() & raw_sequences.astype(str).str.strip().ne("") + sequences = [preprocess_sequence(value) for value in raw_sequences[valid]] + if not sequences: + raise ValueError(f"{dms_id} has no nonempty sequences") + + lengths = {len(sequence) for sequence in sequences} + if len(lengths) != 1: + raise ValueError(f"{dms_id} contains mixed sequence lengths: {sorted(lengths)}") + sequence_length = lengths.pop() batch_size = args.batch_size - if args.max_tokens_per_batch is not None and sequences: - eff_len = effective_length(max_len, args.prepend_bos, fp8_enabled) - batch_size = max(1, args.max_tokens_per_batch // eff_len) - print(f"Token budget {args.max_tokens_per_batch}: seq_len={max_len} " - f"prepend_bos={args.prepend_bos} fp8={fp8_enabled} " - f"effective_length={eff_len} -> batch_size={batch_size}") - - print(f"Running inference (batch_size={batch_size}, " - f"reduce_method={args.reduce_method}, prepend_bos={args.prepend_bos}, " - f"rc={args.average_reverse_complement})...") - scores = evo2_model.score_sequences( - sequences, - batch_size=batch_size, - prepend_bos=args.prepend_bos, - reduce_method=args.reduce_method, - average_reverse_complement=args.average_reverse_complement, + if args.max_tokens_per_batch is not None: + length = effective_length(sequence_length, args.prepend_bos, fp8) + batch_size = max(1, args.max_tokens_per_batch // length) + + print( + f"{dms_id}: scoring {len(sequences)} sequences of length {sequence_length} " + f"in batches of {batch_size}" ) - scores = np.asarray(scores, dtype=float) + scores = np.asarray( + model.score_sequences( + sequences, + batch_size=batch_size, + prepend_bos=args.prepend_bos, + reduce_method=args.reduce_method, + average_reverse_complement=args.average_reverse_complement, + ), + dtype=float, + ) + if scores.shape != (len(sequences),): + raise ValueError( + f"{dms_id} returned score shape {scores.shape}, expected {(len(sequences),)}" + ) + if not np.isfinite(scores).all(): + raise FloatingPointError(f"{dms_id} returned nonfinite model scores") - # Write scores back onto the scored rows (NaN for skipped ones). score_column = f"{args.model_name}_score" - dms_df[score_column] = np.nan - dms_df.loc[valid_mask, score_column] = scores - - # Spearman on the rows we actually scored. - scored = dms_df.loc[valid_mask, ["DMS_score", score_column]].dropna() - if len(scored) >= 2: - correlation, pvalue = spearmanr(scored["DMS_score"], scored[score_column]) + data[score_column] = np.nan + data.loc[valid, score_column] = scores + pairs = data[["DMS_score", score_column]].replace([np.inf, -np.inf], np.nan) + pairs = pairs.dropna() + if len(pairs) < 2: + correlation = pvalue = float("nan") else: - correlation, pvalue = float("nan"), float("nan") + result = spearmanr(pairs["DMS_score"], pairs[score_column]) + correlation, pvalue = result.statistic, result.pvalue - write_csv_atomically(dms_df, output_file) - - print("\nSummary:") - print(f" DMS ID: {dms_id}") - print(f" Sequences scored: {len(sequences)}") - print(f" Score column: {score_column}") - print(f" Spearman vs DMS: {correlation:.3f} (p={pvalue:.2e})") - print(f" Saved to: {output_file}") + output_file = args.output_dir_path / f"{dms_id}.csv" + write_csv_atomically(data, output_file) + print(f"{dms_id}: Spearman={correlation:.3f} p={pvalue:.2e} -> {output_file}") return correlation -def main(): - args = parse_args() +def write_csv_atomically(data: pd.DataFrame, output_file: Path) -> None: + """Replace an output only after its complete CSV has been written.""" + permissions = output_file.stat().st_mode & 0o777 if output_file.exists() else 0o644 + with NamedTemporaryFile( + dir=output_file.parent, + prefix=f".{output_file.name}.", + suffix=".tmp", + delete=False, + ) as handle: + temporary_file = Path(handle.name) + try: + data.to_csv(temporary_file, index=False) + temporary_file.chmod(permissions) + temporary_file.replace(output_file) + finally: + temporary_file.unlink(missing_ok=True) - output_dir = Path(args.output_dir_path) - output_dir.mkdir(parents=True, exist_ok=True) +def run(args: argparse.Namespace, model_factory=None) -> None: + """Load Evo 2 once and score all requested assays.""" row_ids = [args.row_id] if args.row_id is not None else parse_row_ids(args.row_ids) - dms_ids = load_reference_data(args.ref_sheet, row_ids) - print(f"Rows {row_ids} -> DMS IDs: {dms_ids}") - - todo = [] - for row_id, dms_id in zip(row_ids, dms_ids): - output_file = output_dir / f"{dms_id}.csv" - if output_file.exists() and not args.overwrite: - print(f"Output already exists (use --overwrite to redo): {output_file}") - continue - todo.append((row_id, dms_id)) - if not todo: - print("Nothing to score.") + args.output_dir_path.mkdir(parents=True, exist_ok=True) + pending = [ + (row_id, dms_id) + for row_id, dms_id in zip(row_ids, dms_ids) + if args.overwrite or not (args.output_dir_path / f"{dms_id}.csv").exists() + ] + if not pending: + print("Nothing to score") return if not torch.cuda.is_available(): - print("WARNING: CUDA not available - Evo 2 requires a GPU.", file=sys.stderr) - print(f"Visible GPUs: {torch.cuda.device_count()}") - - # Initialize model. Vortex handles device placement / multi-GPU sharding; - # do NOT call .to(device). The checkpoint is loaded once for every assay. - print(f"Loading Evo 2 model: {args.model_name} " - f"(local_path={args.local_path})...") - evo2_model = Evo2(args.model_name, local_path=args.local_path) - # Always ask the built model, never the packaged YAML: load_evo2_model can - # turn FP8 off for 7B when Transformer Engine is missing, and the batch-size - # arithmetic below has to follow the config the model was actually built with. - config = evo2_model.model.config - fp8_enabled = bool(config.get("use_fp8_input_projections", False)) - print(f"use_fp8_input_projections={fp8_enabled}") - if args.require_fp8 and not fp8_enabled: - raise SystemExit( - "--require_fp8 was given but the model resolved to " - "use_fp8_input_projections=False. For a 7B checkpoint this happens " - "silently when Transformer Engine is unavailable; for the others it " - "means Transformer Engine is not providing FP8. Refusing to score, because " - "the surrounding provenance would claim FP8.") + print("WARNING: Evo 2 requires CUDA", file=sys.stderr) + print(f"Loading {args.model_name} on {torch.cuda.device_count()} visible GPUs") + if model_factory is None: + from evo2 import Evo2 + + model_factory = Evo2 + model = model_factory(args.model_name, local_path=args.local_path) + fp8 = bool(model.model.config.get("use_fp8_input_projections", False)) + if args.require_fp8 and not fp8: + raise RuntimeError("The constructed model does not use FP8 input projections") failures = [] - for row_id, dms_id in todo: - print(f"\n=== row {row_id}: {dms_id} ===") + for row_id, dms_id in pending: try: - score_one_assay(evo2_model, args, dms_id, fp8_enabled) - except Exception as e: - print(f"Error scoring {dms_id}: {str(e)}", file=sys.stderr) - failures.append(dms_id) - if len(todo) == 1: + score_assay(model, args, dms_id, fp8) + except Exception as error: + if len(pending) == 1: raise - + print(f"Row {row_id} ({dms_id}) failed: {error}", file=sys.stderr) + failures.append(dms_id) if failures: - print(f"\nFAILED assays ({len(failures)}): {failures}", file=sys.stderr) - sys.exit(1) + raise RuntimeError(f"Failed assays: {', '.join(failures)}") + + +def main() -> None: + """Run the Evo 2 scoring command.""" + run(parse_args()) if __name__ == "__main__": diff --git a/fitness/performance_fitness.py b/fitness/performance_fitness.py index aa215b8..916841b 100755 --- a/fitness/performance_fitness.py +++ b/fitness/performance_fitness.py @@ -332,6 +332,7 @@ def save_results( ) transposed.to_csv(output_dir / "assay_level_results_transposed.csv", index=False) + def select_assays(reference: pd.DataFrame, assay_type: str) -> pd.DataFrame: """Select an assay group.""" if assay_type == "all": diff --git a/fitness/sh/coverage.sh b/fitness/sh/coverage.sh index 0bef7b4..4f11388 100644 --- a/fitness/sh/coverage.sh +++ b/fitness/sh/coverage.sh @@ -5,6 +5,7 @@ export COVERAGE_FILE="$PWD/.pixi/.coverage" cd .. python -m pytest -q tests/test_fitness.py \ --cov=fitness.analyze_fill_strategies \ + --cov=fitness.baselines.Evo.score_evo2_single_dms \ --cov=fitness.baselines.masked_lm \ --cov=fitness.merge_scoring_files \ --cov=fitness.model_registry \ diff --git a/fitness/sh/format.sh b/fitness/sh/format.sh index 1e86c4c..cbf115a 100644 --- a/fitness/sh/format.sh +++ b/fitness/sh/format.sh @@ -4,6 +4,7 @@ set -euo pipefail taplo format pixi.toml ruff check --fix --extend-select I \ analyze_fill_strategies.py \ + baselines/Evo/score_evo2_single_dms.py \ baselines/{AIDO_RNA,Orthrus,RNAGenesis,RNA_FM,RiNALMo}/*.py \ baselines/masked_lm \ merge_scoring_files.py \ @@ -12,6 +13,7 @@ ruff check --fix --extend-select I \ ../tests/test_fitness.py ruff format \ analyze_fill_strategies.py \ + baselines/Evo/score_evo2_single_dms.py \ baselines/{AIDO_RNA,Orthrus,RNAGenesis,RNA_FM,RiNALMo}/*.py \ baselines/masked_lm \ merge_scoring_files.py \ diff --git a/fitness/sh/lint.sh b/fitness/sh/lint.sh index bc3c8e1..2bc3663 100644 --- a/fitness/sh/lint.sh +++ b/fitness/sh/lint.sh @@ -4,6 +4,7 @@ set -euo pipefail taplo lint pixi.toml ruff check --extend-select I \ analyze_fill_strategies.py \ + baselines/Evo/score_evo2_single_dms.py \ baselines/{AIDO_RNA,Orthrus,RNAGenesis,RNA_FM,RiNALMo}/*.py \ baselines/masked_lm \ merge_scoring_files.py \ @@ -12,6 +13,7 @@ ruff check --extend-select I \ ../tests/test_fitness.py ruff format --check \ analyze_fill_strategies.py \ + baselines/Evo/score_evo2_single_dms.py \ baselines/{AIDO_RNA,Orthrus,RNAGenesis,RNA_FM,RiNALMo}/*.py \ baselines/masked_lm \ merge_scoring_files.py \ diff --git a/leaderboard/fitness/README.md b/leaderboard/fitness/README.md index 3daff9a..17dcaf1 100644 --- a/leaderboard/fitness/README.md +++ b/leaderboard/fitness/README.md @@ -14,24 +14,27 @@ Two changes from before: | Rank | Model | Ribozyme (n=26) | tRNA (n=3) | Aptamer (n=2) | Macro (3 ncRNA) | |---:|:--|--:|--:|--:|--:| | 1 | AIDO.RNA (650M) | 0.0660 | 0.4894 | 0.0934 | 0.2163 | -| 2 | Evo 2 (40B) | 0.1081 | 0.4310 | 0.0970 | 0.2120 | -| 3 | AIDO.RNA (1.6B) | 0.0609 | 0.4884 | 0.0718 | 0.2070 | -| 4 | RNA-ERNIE | 0.1343 | 0.4161 | 0.0306 | 0.1937 | -| 5 | Evo 2 (7B) | 0.0651 | 0.3867 | 0.1192 | 0.1904 | -| 6 | RNAGenesis | 0.0750 | 0.4381 | 0.0343 | 0.1825 | -| 7 | RiNALMo | -0.0243 | 0.4856 | 0.0459 | 0.1690 | -| 8 | AIDO.RNA (300M) | -0.0256 | 0.4551 | 0.0372 | 0.1556 | -| 9 | AIDO.RNA (25M) | -0.0299 | 0.4569 | 0.0348 | 0.1540 | -| 10 | Evo 1.5 | 0.0278 | 0.3850 | 0.0007 | 0.1378 | -| 11 | RNA-FM | -0.0225 | 0.4147 | -0.0043 | 0.1293 | -| 12 | Nucleotide Transformer | 0.1329 | 0.3166 | -0.0886 | 0.1203 | -| 13 | AIDO.RNA (1M) | 0.0014 | 0.1777 | 0.0626 | 0.0806 | -| 14 | Orthrus | -0.0349 | 0.0922 | 0.1578 | 0.0717 | -| 15 | Evo 1 | -0.0216 | 0.0948 | 0.0058 | 0.0263 | -| 16 | GenSLM | -0.0045 | -0.0934 | -0.0036 | -0.0338 | - -AIDO.RNA (650M), Evo 2 (40B) and AIDO.RNA (1.6B) span 0.0093. With only 31 assays, this difference -should not be treated as a resolved ordering. +| 2 | Evo 2 (20B) | 0.1105 | 0.4345 | 0.0937 | 0.2129 | +| 3 | Evo 2 (40B) | 0.1081 | 0.4310 | 0.0970 | 0.2120 | +| 4 | AIDO.RNA (1.6B) | 0.0609 | 0.4884 | 0.0718 | 0.2070 | +| 5 | Evo 2 (1B base) | 0.0087 | 0.4466 | 0.1334 | 0.1962 | +| 6 | RNA-ERNIE | 0.1343 | 0.4161 | 0.0306 | 0.1937 | +| 7 | Evo 2 (7B) | 0.0651 | 0.3867 | 0.1192 | 0.1904 | +| 8 | RNAGenesis | 0.0750 | 0.4381 | 0.0343 | 0.1825 | +| 9 | RiNALMo | -0.0243 | 0.4856 | 0.0459 | 0.1690 | +| 10 | AIDO.RNA (300M) | -0.0256 | 0.4551 | 0.0372 | 0.1556 | +| 11 | AIDO.RNA (25M) | -0.0299 | 0.4569 | 0.0348 | 0.1540 | +| 12 | Evo 1.5 | 0.0278 | 0.3850 | 0.0007 | 0.1378 | +| 13 | RNA-FM | -0.0225 | 0.4147 | -0.0043 | 0.1293 | +| 14 | Nucleotide Transformer | 0.1329 | 0.3166 | -0.0886 | 0.1203 | +| 15 | AIDO.RNA (1M) | 0.0014 | 0.1777 | 0.0626 | 0.0806 | +| 16 | Orthrus | -0.0349 | 0.0922 | 0.1578 | 0.0717 | +| 17 | Evo 1 | -0.0216 | 0.0948 | 0.0058 | 0.0263 | +| 18 | GenSLM | -0.0045 | -0.0934 | -0.0036 | -0.0338 | + +The top four models span 0.0093. With only 31 assays, this difference should not be treated as a +resolved ordering. Evo 2 1B is a base-pretrained 8k-context checkpoint, while the 7B, 20B and 40B +checkpoints are mid-trained at 1M context, so it is not a like-for-like size point. All five released AIDO.RNA checkpoints are listed as separate entries rather than in a side table, since they are separate models scored the same way. Their scores rise steeply to 650M and then stop: @@ -145,7 +148,7 @@ a few hundredths as unresolved rather than as an ordering. ## Scoring scripts Scoring scripts, paths relative to the repository root: -- Evo 2 40B: `fitness/baselines/Evo/score_evo2_single_dms.py` and `score_evo2.sh` +- Evo 2: `fitness/baselines/Evo/score_evo2_single_dms.py` and `score_evo2.sh` - Orthrus: `fitness/baselines/Orthrus/score_orthrus_single_dms.py` and `score_orthrus.sh` - AIDO.RNA: `fitness/baselines/AIDO_RNA/score_aido_rna_single_dms.py` and `score_aido_rna.sh` - RNAGenesis: `fitness/baselines/RNAGenesis/score_rnagenesis_single_dms.py` and `score_rnagenesis.sh` diff --git a/tests/test_fitness.py b/tests/test_fitness.py index c5fd542..c5d7bc0 100644 --- a/tests/test_fitness.py +++ b/tests/test_fitness.py @@ -6,12 +6,14 @@ import shutil import sys from pathlib import Path +from types import SimpleNamespace import numpy as np import pandas as pd import pytest import torch from fitness import analyze_fill_strategies +from fitness.baselines.Evo import score_evo2_single_dms as evo2 from fitness.baselines.masked_lm import ( MASK_CHAR, MaskedLMAdapter, @@ -390,6 +392,87 @@ def test_metrics_omit_nonfinite_pairs(): assert observed == pytest.approx(expected) +def test_evo2_workflow_and_guards(tmp_path): + """Score complete assays in one model load and reject failed inference.""" + + class FixtureEvo2: + calls = [] + fp8 = True + loads = 0 + nonfinite = False + + def __init__(self, model_name, local_path=None): + type(self).loads += 1 + self.model = SimpleNamespace( + config={"use_fp8_input_projections": type(self).fp8} + ) + + def score_sequences(self, sequences, **kwargs): + type(self).calls.append((sequences[0], kwargs)) + scores = np.array( + [sequence.count("T") / len(sequence) for sequence in sequences] + ) + if type(self).nonfinite: + scores[0] = np.nan + return scores + + assay_dir = tmp_path / "assays" + shutil.copytree(ASSAY_DIR, assay_dir) + domingo_file = assay_dir / "Domingo_2018_tRNA.csv" + domingo = pd.read_csv(domingo_file) + domingo.loc[0, "sequence"] = np.nan + domingo.to_csv(domingo_file, index=False) + output_dir = tmp_path / "evo2" + command = ( + f"--row_ids 0-2 --ref_sheet {REFERENCE_FILE} " + f"--dms_dir_path {assay_dir} --output_dir_path {output_dir} " + "--model_name evo2_1b_base --max_tokens_per_batch 4096 --require_fp8" + ) + args = evo2.parse_args(shlex.split(command)) + + evo2.run(args, FixtureEvo2) + assert FixtureEvo2.loads == 1 + assert len(FixtureEvo2.calls) == 3 + for sequence, options in FixtureEvo2.calls: + assert options == { + "batch_size": max( + 1, 4096 // evo2.effective_length(len(sequence), False, True) + ), + "prepend_bos": False, + "reduce_method": "mean", + "average_reverse_complement": True, + } + for assay_file in sorted(assay_dir.glob("*.csv")): + output_file = output_dir / assay_file.name + observed = pd.read_csv(output_file) + source = pd.read_csv(assay_file) + score_column = "evo2_1b_base_score" + valid = source["sequence"].notna() + expected = source.loc[valid, "sequence"].map( + lambda sequence: ( + sequence.upper().replace("U", "T").count("T") / len(sequence) + ) + ) + np.testing.assert_allclose(observed.loc[valid, score_column], expected) + assert observed.loc[~valid, score_column].isna().all() + assert output_file.stat().st_mode & 0o777 == 0o644 + assert not list(output_dir.glob("*.tmp")) + + evo2.run(args, FixtureEvo2) + assert FixtureEvo2.loads == 1 + args.row_id, args.row_ids = 0, None + args.output_dir_path = tmp_path / "invalid_evo2" + FixtureEvo2.fp8 = False + with pytest.raises(RuntimeError, match="does not use FP8"): + evo2.run(args, FixtureEvo2) + + FixtureEvo2.fp8 = True + FixtureEvo2.nonfinite = True + with pytest.raises(FloatingPointError, match="nonfinite model scores"): + evo2.run(args, FixtureEvo2) + assert not list(args.output_dir_path.glob("*.csv")) + + def test_merge_rejects_incomplete_predictions(tmp_path): """Reject missing assays and partial rows before replacing merged output.""" assay_name = "Domingo_2018_tRNA.csv"