NER Fine-Tuning¶

This tutorial demonstrates how to fine-tune a DNA language model for Named Entity Recognition (NER) on genomic sequences, identifying features such as exons and introns at the token level.

Full Notebook¶

View Full Notebook

Prerequisites¶

Install DNALLM with the fine-tuning extras:

uv pip install -e '.[base,finetune,cuda124]'

For generating NER training data from scratch, see the NER Data Generation tutorial.

Load Configuration¶

from dnallm import load_config

configs = load_config("./ner_task_config.yaml")

Load Model and Tokenizer¶

from dnallm import load_model_and_tokenizer

model_name = "zhangtaolab/plant-nucleotide-transformer-BPE"
model, tokenizer = load_model_and_tokenizer(
    model_name,
    task_config=configs['task'],
    source="modelscope"
)

Prepare Dataset¶

Load a pre-generated NER dataset (in pickle format) and encode it for token classification:

from dnallm import DNADataset

datasets = DNADataset.load_local_data(
    "./rice_gene_ner_BPE.pkl",
    seq_col="sequence",
    label_col="labels",
    tokenizer=tokenizer,
    max_length=1024
)

Encode with the task-specific collator and split into train/test/validation:

datasets.encode_sequences(
    task=configs['task'].task_type,
    remove_unused_columns=True
)
datasets.split_data()

Initialize Trainer¶

from dnallm import DNATrainer

trainer = DNATrainer(
    model=model,
    config=configs,
    datasets=datasets
)

Start Training¶

metrics = trainer.train()
print(metrics)

Run Inference¶

trainer.infer()