Bacterial 16S rRNA Species Classifier

This model is a 4-layer DNA BERT sequence classification transformer fine-tuned on real 16S rRNA gene sequences from NCBI across 5 bacterial species (Escherichia coli, Bacillus subtilis, Staphylococcus aureus, Pseudomonas aeruginosa, and Salmonella enterica).

📊 Dataset & Benchmarks

Primary Dataset

• Published Dataset: mabrudan/bacterial-16s-sequences
• Total Samples: 500 real bacterial 16S rRNA sequences (100 per species)

Evaluation Metrics (Held-out Test Set)

Model Architecture

• Tokenizer: Overlapping 3-mer Genomic Tokenizer (64 k-mer vocabulary + special tokens)
• Base Architecture: 4-layer BERT Transformer (hidden_size=128, num_attention_heads=4)
• Task: 5-class bacterial species identification from 16S rRNA DNA sequences

Usage Example

import torch
from transformers import BertForSequenceClassification
import itertools

class DNATokenizer:
    def __init__(self, k: int = 3):
        self.k = k
        bases = ['A', 'C', 'G', 'T']
        kmers = [''.join(p) for p in itertools.product(bases, repeat=k)]
        self.vocab = {'[PAD]': 0, '[UNK]': 1, '[CLS]': 2, '[SEP]': 3}
        for idx, kmer in enumerate(kmers, start=4):
            self.vocab[kmer] = idx
        self.cls_token_id = 2
        self.sep_token_id = 3

    def __call__(self, sequence: str, max_length=512):
        clean_seq = ''.join([c for c in sequence.upper() if c in 'ACGT'])
        kmers = [clean_seq[i:i+self.k] for i in range(len(clean_seq)-self.k+1)][:max_length-2]
        ids = [self.cls_token_id] + [self.vocab.get(km, 1) for km in kmers] + [self.sep_token_id]
        return {
            'input_ids': torch.tensor([ids], dtype=torch.long),
            'attention_mask': torch.tensor([[1]*len(ids)], dtype=torch.long)
        }

model_id = 'mabrudan/bacterial-16s-classifier'
model = BertForSequenceClassification.from_pretrained(model_id)
model.eval()

tokenizer = DNATokenizer(k=3)
dna_sequence = 'AGAGTTTGATCATGGCTCAGATTGAACGCTGGCGGCAGGCCTAACACATGCAAGTCGAAC'
inputs = tokenizer(dna_sequence)

with torch.no_grad():
    logits = model(**inputs).logits
    pred_idx = torch.argmax(logits, dim=-1).item()

id2label = model.config.id2label
print('Predicted Species:', id2label[pred_idx])