Speed Performance

[Electrocyte]

I've been using this embedding tool for nearly a year now and really appreciate it. Do you have a more lightweight version that is faster as for large datasets, this tool is very slow.

[goldsteinc]

Hi! if you have used it, any chance you can help us understand how to use this model? we are not sure about the data format and the fine-tuning of it..

[Electrocyte]

It is important to first get your cuda environment working if you haven't already (as it runs very slowly on CPU)

Here is a basic script to get started:

#!/usr/bin/env python3

from transformers import AutoTokenizer, AutoModel
import torch
from torch.utils.data import DataLoader, Dataset

# Helper functions for reading fastq files in group of e.g. 4 lines (2 lines for fasta)
def grouper(iterable, n, fillvalue=None):
    """Group iterable into chunks of n"""
    from itertools import zip_longest
    args = [iter(iterable)] * n
    return zip_longest(*args, fillvalue=fillvalue)

def reverse_complement(dna_seq):
    """Get reverse complement of DNA sequence"""
    complement = {'A': 'T', 'T': 'A', 'G': 'C', 'C': 'G'}
    return ''.join(complement.get(base, base) for base in reversed(dna_seq))

class DNASequencesDataset(Dataset):
    def __init__(self, file_path):
        self.sequences = []
        self.read_ids = []
        
        with open(file_path, 'r') as infile:
            for readID, sequence, plus, quality in grouper(infile, 4, None):
                if sequence is None:
                    continue
                readID = readID.split(" ")[0].split("@")[-1]
                self.sequences.append(sequence.strip())
                self.read_ids.append(readID)

    def __len__(self):
        return len(self.sequences)

    def __getitem__(self, idx):
        dna_seq = self.sequences[idx]
        read_id = self.read_ids[idx]
        reverse_dna_seq = reverse_complement(dna_seq)
        return {
            "dna": dna_seq,
            "reverse_dna": reverse_dna_seq,
            "read_id": read_id
        }

# Load model and tokenizer
# you can also download it locally if you have connectivity issues
tokenizer = AutoTokenizer.from_pretrained("zhihan1996/DNABERT-S", trust_remote_code=True)
model = AutoModel.from_pretrained("zhihan1996/DNABERT-S", trust_remote_code=True)

# Setup device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# Load your FASTQ data
fq = "your_file.fastq"  # Replace with your file path
# should be in the format
# READ ID
# SEQUENCE BASES
# +
# SEQUENCE BASE QUALITIES
batch_size = 32

dataset = DNASequencesDataset(fq)
loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

# Process batches
all_embeddings = []
all_read_ids = []

for batch in loader:
    dna = batch['dna']  # List of DNA sequences
    reverse_dna = batch['reverse_dna']  # List of reverse complement sequences
    read_ids = batch['read_id']  # List of read IDs
    
    # Tokenize DNA sequences
    inputs = tokenizer(
        dna, 
        return_tensors='pt', 
        padding=True, # optional
        truncation=True, # optional
        max_length=512 # optional
    )
    inputs = {k: v.to(device) for k, v in inputs.items()}
    
    # Get embeddings
    with torch.no_grad():
        outputs = model(**inputs)
        embeddings = outputs.last_hidden_state  # Shape: (batch_size, seq_len, 768)
        
        # Get sequence-level embeddings (average pooling)
        sequence_embeddings = embeddings.mean(dim=1)  # Shape: (batch_size, 768)
    
    all_embeddings.append(sequence_embeddings.cpu())
    all_read_ids.extend(read_ids)
    
    print(f"Processed batch with {len(dna)} sequences")

# Combine all embeddings
final_embeddings = torch.cat(all_embeddings, dim=0)
print(f"Final embeddings shape: {final_embeddings.shape}")
print(f"Total sequences processed: {len(all_read_ids)}")

You can download the model using this script:
https://gist.github.com/Electrocyte/e90598b827d4db545c03385151de4c2f