Merge pull request #26 from plantcad/feat/ncbi-utr-derivation

BUSCO evaluation

Author: zongyanliu

README.md

@@ -1,8 +1,21 @@
+---
+title: GeneCAD
+emoji: 🧬
+colorFrom: green
+colorTo: blue
+sdk: gradio
+sdk_version: "5.33.0"
+app_file: app.py
+pinned: false
+license: apache-2.0
+python_version: "3.12"
+---
+
 <div align="center">
 
 # GeneCAD: Plant Genome Annotation with a DNA Foundation Model
 
-![](https://img.shields.io/badge/version-0.2.0-blue)
+![](https://img.shields.io/badge/version-0.2.2-blue)
 [![License](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
 [![CI](https://github.com/plantcad/genecad/actions/workflows/ci.yaml/badge.svg)](https://github.com/plantcad/genecad/actions/workflows/ci.yaml)
 [![bioRxiv](https://img.shields.io/badge/bioRxiv-10.1101/2025.10.31.685877-b31b1b.svg)](https://doi.org/10.1101/2025.10.31.685877)

predict.sh

@@ -81,6 +81,7 @@ MIN_TRANSCRIPT_LENGTH="3"
 CPU_WORKERS="1"
 LAUNCHER_ARG="${LAUNCHER:-}"
 MODEL_CHECKPOINT_ARG=""
+BASE_MODEL_ARG=""
 
 while [[ $# -gt 0 ]]; do
   case $1 in
@@ -95,6 +96,7 @@ while [[ $# -gt 0 ]]; do
     -g|--gpus)       GPUS_ARG="$2";       shift 2 ;;
     --launcher)      LAUNCHER_ARG="$2";   shift 2 ;;
     --model-checkpoint) MODEL_CHECKPOINT_ARG="$2"; shift 2 ;;
+    --base-model)    BASE_MODEL_ARG="$2"; shift 2 ;;
     -h|--help) usage ;;
     *) echo "Unknown option: $1"; usage ;;
   esac
@@ -152,6 +154,10 @@ esac
 if [[ -n "$MODEL_CHECKPOINT_ARG" ]]; then
     HEAD_MODEL="$MODEL_CHECKPOINT_ARG"
 fi
+if [[ -n "$BASE_MODEL_ARG" ]]; then
+    BASE_MODEL="$BASE_MODEL_ARG"
+    TOKENIZER_PATH="$BASE_MODEL"
+fi
 
 # =================================================================
 # GPU Resolution

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "genecad"
-version = "0.2.1"
+version = "0.2.2"
 description = "GeneCAD: A long-context genome annotation model built on PlantCAD2"
 readme = "README.md"
 requires-python = ">=3.12,<3.13"

scripts/evaluate.py

@@ -495,9 +495,12 @@ def eval_splice_sites(pred_exon_genes, genome):
 def extract_transcripts(pred_exon_genes, genome, output_path):
     """
     Extract spliced transcript sequences directly from the predicted exon
-    structure — replaces the need for gffread.
+    structure — replaces the need for gffread. Transcripts containing
+    ambiguous bases are excluded because older MetaEuk/MMseqs versions can
+    crash while translating them.
     """
-    count = skipped = 0
+    count = skipped = skipped_ambiguous = 0
+    valid_bases = set("ACGTacgt")
     with open(output_path, "w") as fh:
         for gdata in pred_exon_genes.values():
             ch, st = gdata["chrom"], gdata["strand"]
@@ -512,11 +515,20 @@ def extract_transcripts(pred_exon_genes, genome, output_path):
                 tx_seq = "".join(seq[s - 1 : e] for (s, e) in exons)
                 if st == "-":
                     tx_seq = _revcomp(tx_seq)
+                if not set(tx_seq) <= valid_bases:
+                    skipped_ambiguous += 1
+                    continue
                 fh.write(f">{tx_id}\n{tx_seq}\n")
                 count += 1
     print(f"[INFO] Extracted {count} transcript sequences.", file=sys.stderr)
     if skipped:
         print(f"[WARN] Skipped {skipped} genes (chrom not in FASTA).", file=sys.stderr)
+    if skipped_ambiguous:
+        print(
+            f"[WARN] Skipped {skipped_ambiguous} transcripts containing "
+            "ambiguous bases for BUSCO.",
+            file=sys.stderr,
+        )
     return count
 
 

scripts/extract.py

@@ -216,26 +216,87 @@ def _extract_gff_features(
                 # For now, we use the longest transcript as the canonical transcript
                 transcript_is_canonical = is_longest_transcript(transcript)
 
-                # Process each feature
+                # Collect sub-features by type for this transcript
+                exon_sub_features: list[SeqFeature] = []
+                cds_sub_features: list[SeqFeature] = []
+                utr_sub_features: list[SeqFeature] = []
                 for feature in transcript.sub_features:
-                    # Skip exon features if requested (they are structural, not modeling features)
-                    if skip_exon_features and feature.type == "exon":
-                        continue
-
-                    if feature.type not in [
+                    if feature.type == "exon":
+                        exon_sub_features.append(feature)
+                    elif feature.type == GffFeatureType.CDS:
+                        cds_sub_features.append(feature)
+                    elif feature.type in (
                         GffFeatureType.FIVE_PRIME_UTR,
-                        GffFeatureType.CDS,
                         GffFeatureType.THREE_PRIME_UTR,
-                    ]:
+                    ):
+                        utr_sub_features.append(feature)
+                    else:
                         raise ValueError(
                             f"Found unexpected transcript feature type: {feature.type}"
                         )
 
+                # Derive UTRs from exon-CDS when explicit UTR features are absent
+                derived_utr_features: list[tuple[int, int, GffFeatureType]] = []
+                if (
+                    not skip_exon_features
+                    and exon_sub_features
+                    and cds_sub_features
+                    and not utr_sub_features
+                ):
+                    cds_coords = [
+                        (int(f.location.start), int(f.location.end))
+                        for f in cds_sub_features
+                    ]
+                    cds_min = min(s for s, _ in cds_coords)
+                    cds_max = max(e for _, e in cds_coords)
+                    strand = transcript_info.strand
+                    for ef in exon_sub_features:
+                        es = int(ef.location.start)
+                        ee = int(ef.location.end)
+                        if strand == 1:
+                            if es < cds_min:
+                                derived_utr_features.append(
+                                    (
+                                        es,
+                                        min(ee, cds_min),
+                                        GffFeatureType.FIVE_PRIME_UTR,
+                                    )
+                                )
+                            if ee > cds_max:
+                                derived_utr_features.append(
+                                    (
+                                        max(es, cds_max),
+                                        ee,
+                                        GffFeatureType.THREE_PRIME_UTR,
+                                    )
+                                )
+                        else:
+                            if ee > cds_max:
+                                derived_utr_features.append(
+                                    (
+                                        max(es, cds_max),
+                                        ee,
+                                        GffFeatureType.FIVE_PRIME_UTR,
+                                    )
+                                )
+                            if es < cds_min:
+                                derived_utr_features.append(
+                                    (
+                                        es,
+                                        min(ee, cds_min),
+                                        GffFeatureType.THREE_PRIME_UTR,
+                                    )
+                                )
+
+                # Process CDS and UTR features (explicit or derived)
+                modeling_features: list[SeqFeature] = (
+                    cds_sub_features + utr_sub_features
+                )
+                for feature in modeling_features:
                     feature_id = get_feature_id(feature)
                     feature_info = get_position_info(feature)
                     feature_name = get_feature_name(feature)
 
-                    # Create a validated GeneFeatureData object
                     feature_data = SequenceFeature(
                         species_id=species_id,
                         species_name=species_name,
@@ -263,6 +324,38 @@ def _extract_gff_features(
                     )
                     features_data.append(feature_data)
 
+                for feat_start, feat_stop, feat_type in derived_utr_features:
+                    if feat_start >= feat_stop:
+                        continue
+                    strand_char = "+" if transcript_info.strand == 1 else "-"
+                    derived_id = f"{feat_type}:{transcript_id}:{feat_start}:{feat_stop}:{strand_char}"
+                    feature_data = SequenceFeature(
+                        species_id=species_id,
+                        species_name=species_name,
+                        chromosome_id=chrom_id,
+                        chromosome_name=chrom_name,
+                        chromosome_length=chrom_length,
+                        gene_id=gene_id,
+                        gene_name=gene_name,
+                        gene_strand=gene_info.strand,
+                        gene_start=gene_info.start,
+                        gene_stop=gene_info.stop,
+                        transcript_id=transcript_id,
+                        transcript_name=transcript_name,
+                        transcript_strand=transcript_info.strand,
+                        transcript_is_canonical=transcript_is_canonical,
+                        transcript_start=transcript_info.start,
+                        transcript_stop=transcript_info.stop,
+                        feature_id=derived_id,
+                        feature_name=None,
+                        feature_strand=transcript_info.strand,
+                        feature_type=feat_type,
+                        feature_start=feat_start,
+                        feature_stop=feat_stop,
+                        filename=filename,
+                    )
+                    features_data.append(feature_data)
+
     logger.info(
         f"[species={species_config.id}] Total features extracted: {len(features_data)}"
     )

scripts/install_release.sh

@@ -8,7 +8,7 @@ set -euo pipefail
 #   GENECAD_VERSION=0.2.0
 #   PYTORCH_CUDA_INDEX=https://download.pytorch.org/whl/cu128
 
-GENECAD_VERSION="${GENECAD_VERSION:-0.2.0}"
+GENECAD_VERSION="${GENECAD_VERSION:-0.2.2}"
 PYTORCH_CUDA_INDEX="${PYTORCH_CUDA_INDEX:-https://download.pytorch.org/whl/cu128}"
 WHEEL_URL="https://github.com/plantcad/genecad/releases/download/v${GENECAD_VERSION}/genecad-${GENECAD_VERSION}-py3-none-any.whl"
 

src/reelprotein.py

@@ -110,11 +110,9 @@ def build_cds_string_for_gene(gene, chrom_seq):
         raw_seq = chrom_seq[a - 1 : b]
         seq_str = str(raw_seq).upper()
         if strand == "+":
-            trimmed = seq_str[cds["phase"] :]
+            parts.append(seq_str)
         else:
-            rc = str(Seq(seq_str).reverse_complement())
-            trimmed = rc[cds["phase"] :]
-        parts.append(trimmed)
+            parts.append(str(Seq(seq_str).reverse_complement()))
 
     return "".join(parts)
 

train.sh

@@ -325,6 +325,11 @@ DATA_DIR="$PIPELINE_DIR/data/$DOMAIN"
 GFF_DIR="$DATA_DIR/gff/training"
 FASTA_DIR="$DATA_DIR/fasta/training"
 
+# Default: skip exon features (Ensembl/plant GFFs have explicit UTR annotations).
+# Set to "no" for NCBI-style GFFs that lack UTR features; extract.py will then
+# derive five_prime_UTR / three_prime_UTR from exon-CDS difference.
+GFF_SKIP_EXON="yes"
+
 if [[ -z "$ANIMAL_INPUT_DIR" ]]; then
   ANIMAL_INPUT_DIR="$FASTA_DIR"
 fi
@@ -686,6 +691,11 @@ else
     fi
   done
 
+  # Custom NCBI GFFs lack explicit UTR features; derive them from exon-CDS in extract.py.
+  if [[ $contains_ensembl -eq 0 ]]; then
+    GFF_SKIP_EXON="no"
+  fi
+
   if [[ $contains_ensembl -eq 1 ]]; then
     mkdir -p "$ANIMAL_GFF_DIR" "$ANIMAL_INPUT_DIR"
 
@@ -867,6 +877,7 @@ if [ ! -f "$EXTRACT_DIR/raw_features.parquet" ]; then
           --input-dir "$GFF_DIR" \
           --species-id "$sid" \
           --output "$out_file" \
+          --skip-exon-features "$GFF_SKIP_EXON" \
           "${SPECIES_CONFIG_ARG[@]}"
       ) &
 
@@ -903,6 +914,7 @@ PY
       --input-dir "$GFF_DIR" \
       --species-id $SPECIES_IDS \
       --output "$EXTRACT_DIR/raw_features.parquet" \
+      --skip-exon-features "$GFF_SKIP_EXON" \
       "${SPECIES_CONFIG_ARG[@]}"
   fi
   echo "  Created: raw_features.parquet"