zhangtaolab/FungiHelixSeek-Taxonomy

FungiHelixSeek-Taxonomy

FungiHelixSeek-Taxonomy is a taxonomic classification model fine-tuned from the FungiHelixSeek DNA foundation model. It assigns fungal internal transcribed spacer (ITS) barcode sequences to six hierarchical taxonomic ranks — phylum → class → order → family → genus → species — through six independent linear heads over a max-pooled HelixSeek backbone.

Training data: UNITE-style ITS references (~1.5 M sequences), 10 epochs, hierarchical label smoothing (per-rank ε = 0.02, uniform ε = 0) + class-weighted cross-entropy, AdamW with cosine schedule.

Taxonomic heads

Performance

Accuracy per rank on the held-out filamentous-fungi test set (cleanly-labelled reference subset; the number of evaluated references varies by rank):

Species-level accuracy is computed over 6,692 evaluated references.

Dataset

Training data: fungi-its-barcodes · ModelScope — ~1.5 M UNITE-style fungal ITS references spanning 14,742 species, with yeast / filamentous / large-scale test splits.

Usage

Load with trust_remote_code=True:

import json, torch
from transformers import AutoModel, AutoTokenizer

model_dir = "models/FungiHelixSeek-Taxonomy"
tokenizer = AutoTokenizer.from_pretrained(model_dir)
model = AutoModel.from_pretrained(model_dir, trust_remote_code=True).eval()

sequence = "GTAACAAGGTTTCCGTAGGTGAACCTGCGGAAGGATCATTATAGAAAAAAATGGAAGGGCCATGCGCTTAATTGCGCGG"
enc = tokenizer(sequence, truncation=True, max_length=1024, return_tensors="pt")
with torch.no_grad():
    logits_list = model(input_ids=enc["input_ids"], attention_mask=enc["attention_mask"]).logits_list

level_names = ["phylum", "class", "order", "family", "genus", "species"]
label_encoders = json.load(open(f"{model_dir}/label_encoders.json"))
id2label = {lvl: {i: name for name, i in enc_map.items()} for lvl, enc_map in label_encoders.items()}
prediction = {
    lvl: id2label[lvl][logits.argmax(dim=-1).item()]
    for lvl, logits in zip(level_names, logits_list)
}
print(prediction)
# {'phylum': 'Ascomycota', 'class': 'Saccharomycetes', ..., 'species': '...'}

label_encoders.json (shipped in this repo) maps each level's label string → class id and is required to decode predictions. Ready-made inference script: `scripts/species_classification/predict.py`.

Model details

• —Architecture: HelixSeek backbone (hybrid Transformer / Delta linear attention / MLA / sparse MoE, 14 blocks) + 6 linear classification heads on mask-aware max-pooled hidden states.
• —Input: single ITS nucleotide sequence (character-level DNA tokenizer, vocab 11); sequences are truncated to 1,024 tokens.
• —Model code: configuration.py / model.py in this repository (trust_remote_code); state dict layout backbone.* + heads.{0..5}.*.
• —Language: English (model), DNA (input).

License

CC-BY-NC-4.0 (Attribution-NonCommercial 4.0 International). For commercial licensing, contact the authors.