hyejong/genomeocean-sub-classifier

GenomeOcean Sub Classifier

A GenomeOcean 100M v1.2 fine-tuned sequence classifier that splits genomic FASTA contigs already identified as NCLDV/Mirus candidates into:

Label	Meaning
NCLDV	Nucleocytoviricota (giant virus) sequences
Mirus	Mirus lineage sequences

This model is the second stage of a two-stage hierarchical Giant Virus classifier and was trained independently on the full set of prepared NCLDV and Mirus sequences (not only on candidates predicted by stage one). It expects sequences that have already been narrowed down to NCLDV/Mirus candidates — typically the output of the companion hyejong/genomeocean-main-classifier model. Running it directly on arbitrary Cellular or other viral sequences will still produce an NCLDV/Mirus label, but that label is not meaningful outside the NCLDV/Mirus scope. Full pipeline code, preprocessing details, and CLI tooling are available at Genomeocean_Giant_Virus_Classifier on GitHub.

Model architecture

• Base: GenomeOcean 100M v1.2, a Mistral-architecture genomic language model, fine-tuned for sequence classification (MistralForSequenceClassification).
• Hidden size 768, 12 layers, 8 attention heads, vocab size 4,096, max position embeddings 32,768.
• Custom modeling code (modeling_mistral.py, configuration_mistral.py) is included in each fold folder, so loading requires trust_remote_code=True.
• config.json does not define id2label; the class-to-label mapping is 0 -> NCLDV, 1 -> Mirus (see usage snippet below).

5-fold ensemble

This repository hosts five independently fine-tuned folds (fold1 – fold5), each in its own subfolder with a full set of model, tokenizer, and config files. The recommended way to use this model is to run all five folds on the same input and average (soft-vote) the softmax probabilities — this is what the reference pipeline below does automatically. A single fold can also be used on its own, at the cost of losing the variance/agreement signal the ensemble provides.

Input format

Inputs are 5,000 bp genomic chunks derived from FASTA contigs, preprocessed as follows before tokenization:

1. Convert the sequence to uppercase.
2. Remove characters other than A/C/G/T/N.
3. Remove N.
4. Split the cleaned sequence into 5,000 bp windows with a 5,000 bp stride (non-overlapping).
5. Contigs shorter than 5,000 bp, and the incomplete tail of longer contigs, are not used for prediction.

Tokenization uses the fold's own tokenizer with max_length=1250.

Usage

Recommended: reference CLI package

The GitHub repository ships an installable CLI (genomeocean-sub) that handles FASTA parsing, chunking, batching, 5-fold ensembling, and result aggregation for you:

git clone https://github.com/hyej0ng/Genomeocean_Giant_Virus_Classifier.git
cd Genomeocean_Giant_Virus_Classifier
python -m pip install -r requirements.txt
python -m pip install ./genomeocean-sub-classifier

genomeocean-sub predict \
  --input /path/to/ncldv_mirus_candidates.fna \
  --output-dir /path/to/sub_results \
  --model-id hyejong/genomeocean-sub-classifier \
  --subfolder fold1 --subfolder fold2 --subfolder fold3 \
  --subfolder fold4 --subfolder fold5 \
  --device cuda

See the Sub Classifier README for the full CLI reference, or the Integrated Pipeline README to run Main and Sub together and get final Cellular / NCLDV / Mirus / Other Viruses labels in one command, starting from any FASTA.

Direct transformers usage (single fold)

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

model_id = "hyejong/genomeocean-sub-classifier"
fold = "fold1"

tokenizer = AutoTokenizer.from_pretrained(model_id, subfolder=fold, trust_remote_code=True)
model = AutoModelForSequenceClassification.from_pretrained(
    model_id, subfolder=fold, trust_remote_code=True
)
model.eval()

id2label = {0: "NCLDV", 1: "Mirus"}

sequence = "ACGT..."  # a single, already-preprocessed 5,000 bp chunk from an
                       # NCLDV/Mirus candidate contig
inputs = tokenizer(sequence, return_tensors="pt", truncation=True, max_length=1250)
with torch.inference_mode():
    probs = torch.softmax(model(**inputs).logits, dim=-1)[0]

predicted = id2label[int(probs.argmax())]
print(predicted, probs.tolist())

To reproduce the full 5-fold ensemble manually, repeat the above for fold1 – fold5 and average the resulting probability vectors before taking the argmax.

Training data

Trained independently on the full set of prepared genomic contigs/fragments for each class:

Label	Scope
NCLDV	Full NCLDV sequences
Mirus	Full Mirus sequences

Limitations

• This model performs a binary split and assumes the input is already an NCLDV/Mirus candidate; it does not detect Cellular or other viral sequences.
• Confidence scores are hierarchical/comparative, not calibrated probabilities of biological truth.
• Sequence length, assembly quality, and distance from the training distribution can affect predictions; ensemble agreement across folds (ensemble_agreement, confidence_std in the reference pipeline output) should be checked for low-confidence or borderline calls.
• For candidate selection from arbitrary FASTA input, pair this model with hyejong/genomeocean-main-classifier.

License

MIT