Generation Model Fine-Tuning

This tutorial demonstrates how to fine-tune causal language models (DNAGPT and megaDNA) for DNA sequence generation. After fine-tuning on coding sequences, the models can generate new sequences from a short prompt.

Full Notebook

View Full Notebook

Prerequisites

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

The megaDNA half 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 first code 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:

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")

Prepare Training Data

Download and preprocess coding sequences from Arabidopsis. The !wget genome download cell runs BEFORE the Fasta(...) load cell that reads the file (wget -c resumes a partial download):

from pyfastx import Fasta

genome = Fasta("Arabidopsis_thaliana.TAIR10.cds.all.fa.gz")
with open("ath_cds.csv", "w") as f:
    print("seq_id,sequence", file=f)
    for seq in genome:
        print(f"{seq.name},{seq.seq}", file=f)

Load and split the dataset:

# Load the datasets
data_path = "ath_cds.csv"
datasets = DNADataset.load_local_data(data_path, seq_col="sequence", sep=",")

# Sampling the datasets
datasets.sampling(0.1, seed=42, overwrite=True)
datasets.split_data(seed=42)

Preview a sample sequence:

seq = datasets.dataset["test"][10]["sequence"]
prompt = seq[:10]
print("Length:", len(seq))
print("Prompt sequence:", prompt)
print("Full sequence:  ", seq)

Part 1: DNAGPT

Load Configuration

import copy
from dnallm import load_config

configs = load_config("./finetune_config.yaml")
configs["finetune"].output_dir = "./outputs_dnagpt"

Load Model

from dnallm import load_model_and_tokenizer

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

Encode and Train

data = copy.deepcopy(datasets)
data.encode_sequences(tokenizer=tokenizer)

from dnallm import DNATrainer

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

metrics = trainer.train()
print(metrics)

Generate Sequences

model.eval()
tokenizer.pad_token = tokenizer.eos_token

inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
outputs = model.generate(
    **inputs,
    max_length=len(seq) + 5,
    num_return_sequences=5,
    do_sample=True,
    top_k=50,
    top_p=0.95,
    temperature=1.0
)

Decode and compare with the original:

print("Prompt:               ", prompt)
for i, out in enumerate(outputs):
    out_seq = tokenizer.decode(out, skip_special_tokens=True)
    print(f"Generated sequence {i}: ", out_seq.replace(" ", ""))
print("Raw sequence:         ", seq)

Part 2: megaDNA

megaDNA uses an embedding task type and requires custom trainer logic.

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. The clone is pinned to commit cb2f5ab4cc88dc0effe05c5f23358862c837014a and MEGABYTE_pytorch to 0.2.1 (the repo's own requirements pin) — never a floating clone. uv pip install targets the kernel's VIRTUAL_ENV: under the isolated dnallm-megadna kernelspec this mutates only the throwaway venv.

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

Update Config

configs = load_config("./finetune_config.yaml")
configs["task"].task_type = "embedding"
configs["finetune"].output_dir = "./outputs_megadna"

Load Model

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

Encode and Prepare

data = copy.deepcopy(datasets)
data.encode_sequences(tokenizer=tokenizer)

megaDNA requires specific column renaming (token_type_ids only materializes on transformers < 5, so drop whichever non-input columns are present on the running stack):

_present = [
    column
    for column in ["seq_id", "sequence", "token_type_ids", "attention_mask"]
    if column in data.dataset.column_names
]
data.dataset = data.dataset.remove_columns(_present)
data.dataset = data.dataset.rename_column("input_ids", "ids")

Custom Trainer

# Define a custom trainer for MEGA-DNA
class MegaDNATrainer(type(trainer.trainer)):
    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
        loss = model(**inputs, return_value="loss")
        if return_outputs:
            logits = model(**inputs, return_value="logits")
            return (loss, logits)
        return loss

trainer.customize_trainer(MegaDNATrainer)
trainer.trainer.can_return_loss = True

Train

metrics = trainer.train()
print(metrics)

Generate

model.eval()
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
outputs = [
    model.generate(inputs["input_ids"], seq_len=len(seq) + 5, temperature=0.95, filter_thres=0.0)
    for _ in range(5)
]

print("Prompt:               ", prompt)
for i, out in enumerate(outputs):
    out_seq = tokenizer.decode(out[0], skip_special_tokens=True)
    print(f"Generated sequence {i}: ", out_seq.replace(" ", ""))
print("Raw sequence:         ", seq)

Related Tutorials

• Sequence Generation Inference
• Binary Classification