NER Fine-Tuning¶

This tutorial demonstrates how to fine-tune a DNA large 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 base extras:

uv pip install -e '.[base,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()