mabrudan/bacterial-16s-classifier

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)

Species	Precision	Recall	F1-Score	Support
Staphylococcus aureus	0.59	1.00	0.74	17
Escherichia coli	0.44	1.00	0.61	18
Pseudomonas aeruginosa	0.50	0.65	0.56	17
Bacillus subtilis	0.67	0.15	0.24	27
Salmonella enterica	0.00	0.00	0.00	21
Overall Accuracy	—	—	50.0%	100

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