Skip to content

Custom Classification Head Fine-Tuning

This tutorial shows how to fine-tune with a custom classification head. You will first train with a standard Transformer-compatible model, then switch to a specialized architecture (megaDNA) that requires a custom head implementation.

Full Notebook

View Full Notebook

Prerequisites

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

The megaDNA half (Part 2) additionally needs the FEASIBILITY-locked prerequisites: the lingxusb/megaDNA clone pinned to commit cb2f5ab4cc88dc0effe05c5f23358862c837014a plus MEGABYTE_pytorch==0.2.1 (never a floating clone — see the pinned install cell in Part 2).

Environment provenance

The provenance stamp cell prints the exact library versions the notebook ran with as key=value lines (D-21); this notebook has no flash-linear-attention dependency, so its fla_version line records not-used and the megaDNA pins are stamped alongside:

# Provenance stamp (D-21): the exact versions this notebook ran with, printed
# as key=value lines. No flash-linear-attention dependency here, so the fla
# line records not-used and the FEASIBILITY-locked megaDNA pins are stamped
# alongside (the pinned install itself happens in the megaDNA section below).
import torch
import transformers

print(f"transformers_version={transformers.__version__}")
print(f"torch_version={torch.__version__}")
print("fla_version=not-used")
print("megadna_commit=cb2f5ab4cc88dc0effe05c5f23358862c837014a")
print("megabyte_version=0.2.1")

Part 1: Standard Model with Default Head

Load Configuration

from dnallm import load_config

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

Load Model and Tokenizer

from dnallm import load_model_and_tokenizer

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

Prepare Dataset

from dnallm import DNADataset

data_name = "zhangtaolab/plant-multi-species-core-promoters"
datasets = DNADataset.from_modelscope(
    data_name,
    seq_col="sequence",
    label_col="label",
    tokenizer=tokenizer,
    max_length=512
)

sampled_datasets = datasets.sampling(0.1, overwrite=True)
sampled_datasets.encode_sequences()

Train

from dnallm import DNATrainer

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

metrics = trainer.train()
print(metrics)

Part 2: megaDNA with Custom Head

megaDNA is not compatible with the standard Transformers classification head. DNALLM handles this by allowing you to specify a custom head in the config.

Update Configuration

configs['task'].head_config.head = "megadna"
configs['finetune'].output_dir = "./outputs_megadna"

Install pinned prerequisites

The megaDNA model checkpoint unpickles classes from the megaDNA package, so the pinned prerequisites must be installed BEFORE the model-load cell. uv pip install targets the running kernel's VIRTUAL_ENV (the project venv for this notebook) and is reversible (exact versions):

git clone https://github.com/lingxusb/megaDNA.git
git -C megaDNA checkout cb2f5ab4cc88dc0effe05c5f23358862c837014a
uv pip install MEGABYTE_pytorch==0.2.1 ./megaDNA

Load megaDNA Model

model_name = "lingxusb/megaDNA_updated"
model, tokenizer = load_model_and_tokenizer(
    model_name,
    task_config=configs['task'],
    source="huggingface"
)

Prepare Dataset

datasets = DNADataset.from_modelscope(
    data_name,
    seq_col="sequence",
    label_col="label",
    tokenizer=tokenizer,
    max_length=1024
)
sampled_datasets = datasets.sampling(0.1, overwrite=True)
sampled_datasets.encode_sequences()

Train

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

metrics = trainer.train()
print(metrics)